89 lines
3.8 KiB
Python
89 lines
3.8 KiB
Python
"""Explicit name-based mapping between mjlab and canonical S10 model tensors."""
|
|
|
|
import torch
|
|
|
|
from s10_supermimic.robot.config import (
|
|
LEG_JOINT_NAMES, WHEEL_BODY_NAMES, WHEEL_GEOM_NAMES, WHEEL_JOINT_NAMES,
|
|
)
|
|
from .world_model import STATE_DIM, ACTION_DIM
|
|
|
|
|
|
class S10Bridge:
|
|
def __init__(self, env):
|
|
self.env = env
|
|
self.robot = env.scene["robot"]
|
|
self.legs, _ = self.robot.find_joints(LEG_JOINT_NAMES, preserve_order=True)
|
|
self.wheels, _ = self.robot.find_joints(WHEEL_JOINT_NAMES, preserve_order=True)
|
|
if len(self.legs) != 12 or len(self.wheels) != 4:
|
|
raise ValueError("S10 requires 12 leg joints and 4 wheel joints")
|
|
pos = env.action_manager.get_term("joint_pos")
|
|
vel = env.action_manager.get_term("wheel_vel")
|
|
if list(env.cfg.actions) != ["joint_pos", "wheel_vel"]:
|
|
raise ValueError("Expected joint_pos then wheel_vel action terms")
|
|
names = list(pos.target_names) + list(vel.target_names)
|
|
canonical = LEG_JOINT_NAMES + WHEEL_JOINT_NAMES
|
|
if len(names) != ACTION_DIM or set(names) != set(canonical):
|
|
raise ValueError("Unexpected S10 action targets")
|
|
self.to_canonical = [names.index(n) for n in canonical]
|
|
self.to_env = [canonical.index(n) for n in names]
|
|
geoms, _ = self.robot.find_geoms(WHEEL_GEOM_NAMES, preserve_order=True)
|
|
self.geoms = self.robot.indexing.geom_ids[geoms]
|
|
body, _ = self.robot.find_bodies("base_link")
|
|
self.base = self.robot.indexing.body_ids[body[0]]
|
|
# XML reference, not an already-randomized simulator value after reset.
|
|
from s10_supermimic.robot.config import get_spec
|
|
self.nominal_com = torch.as_tensor(get_spec().body("base_link").ipos.copy(),
|
|
device=env.device, dtype=torch.float32)
|
|
|
|
def canonical_action(self, action):
|
|
return action[:, self.to_canonical]
|
|
|
|
def env_action(self, action):
|
|
return action[:, self.to_env]
|
|
|
|
def state(self):
|
|
d = self.robot.data
|
|
s = torch.cat((d.root_link_pos_w[:, 2:3], d.projected_gravity_b,
|
|
d.root_link_lin_vel_b, d.root_link_ang_vel_b,
|
|
(d.joint_pos - d.default_joint_pos)[:, self.legs],
|
|
d.joint_vel[:, self.legs], d.joint_vel[:, self.wheels]), -1)
|
|
assert s.shape[-1] == STATE_DIM
|
|
return s
|
|
|
|
def pose(self):
|
|
d = self.robot.data
|
|
return torch.cat((d.root_link_pos_w[:, :2], d.heading_w[:, None]), -1)
|
|
|
|
def context(self):
|
|
model = self.env.sim.model
|
|
friction = model.geom_friction
|
|
com = model.body_ipos
|
|
mu = friction[:, self.geoms, 0] if friction.ndim == 3 else friction[self.geoms, 0].expand(self.env.num_envs, -1)
|
|
offset = com[:, self.base] if com.ndim == 3 else com[self.base].expand(self.env.num_envs, -1)
|
|
return torch.cat((mu, offset - self.nominal_com), -1)
|
|
|
|
def contacts(self):
|
|
"""Endpoint targets: contact flags4, net wheel forces12 in base frame (N).
|
|
|
|
Explicit names prevent silent FL/FR swaps in sensor expansion order.
|
|
Terminal/reset endpoints are excluded by replay validity in the learner.
|
|
"""
|
|
from mjlab.utils.lab_api.math import quat_apply_inverse
|
|
sensor = self.env.scene["wheel_ground_contact"]
|
|
names = [name.split('/')[-1] for name in sensor.primary_names]
|
|
order = [names.index(name) for name in WHEEL_BODY_NAMES]
|
|
found = (sensor.data.found[:,order] > 0).float().reshape(-1,4)
|
|
force = sensor.data.force[:,order].reshape(-1,4,3)
|
|
quat = self.robot.data.root_link_quat_w[:,None,:].expand(-1,4,-1)
|
|
force = quat_apply_inverse(quat.reshape(-1,4),force.reshape(-1,3)).reshape(-1,12)
|
|
return torch.cat((found,force),-1)
|
|
|
|
|
|
def pose_delta(before, after):
|
|
dxy = after[:, :2] - before[:, :2]
|
|
c, s = before[:, 2].cos(), before[:, 2].sin()
|
|
yaw = after[:, 2] - before[:, 2]
|
|
return torch.stack((c * dxy[:, 0] + s * dxy[:, 1],
|
|
-s * dxy[:, 0] + c * dxy[:, 1],
|
|
torch.atan2(yaw.sin(), yaw.cos())), -1)
|