Initial commit

s10_supermimic/core/bridge.py

@@ -0,0 +1,88 @@
+"""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)