Initial commit

2026-11-02 17:10:39 +08:00
commit 102fdbaf47
79 changed files with 56509 additions and 0 deletions
@@ -0,0 +1,149 @@
+"""S10 actor history and actor/critic networks, independent of the quadruped learner.
+
+Two Actor observation schemas coexist:
+
+* v4 / v4-upward / new v7: 53 values/frame, including measured wheel speed;
+* legacy v7: 49 values/frame, no measured wheel speed (the last action still
+  contains the four wheel commands).
+
+The learner, replay and deployment exporter select the schema explicitly. Old
+checkpoints are recognized from their first-layer width.
+"""
+
+import torch
+from torch import nn
+from .world_model import RunningNorm, ANG_VEL, GRAVITY, LEG_POS, LEG_VEL, WHEEL_VEL
+
+LEGACY_TERMS = (3, 3, 12, 12, 4, 16, 3)
+TERMS = (3, 3, 12, 12, 16, 3)
+LEGACY_FRAME_DIM = sum(LEGACY_TERMS)
+FRAME_DIM = sum(TERMS)
+CRITIC_DIM = 38 + 3 + 7 + 16
+POLICY_FORMAT = "s10-flat-policy-v1"
+
+
+def frame(state, last_env_action, command, bias, noise=False, wheel_vel=True):
+  """Exactly the term order/scaling in env.flat (noise before scaling)."""
+  values = [state[:, ANG_VEL], state[:, GRAVITY], state[:, LEG_POS] + bias,
+            state[:, LEG_VEL]]
+  widths = [0.2, 0.05, 0.01, 1.5]
+  if wheel_vel:
+    values.append(state[:, WHEEL_VEL])
+    widths.append(1.5)
+  values.extend((last_env_action, command))
+  widths.extend((0., 0.))
+  if noise:
+    values = [v + (torch.rand_like(v) * 2 - 1) * w if w else v
+              for v, w in zip(values, widths)]
+  # Leg velocity is always index 3; wheel velocity, when present, is index 4.
+  velocity_indices = (3, 4) if wheel_vel else (3,)
+  return torch.cat([v * (0.05 if i in velocity_indices else 1.)
+                    for i, v in enumerate(values)], -1)
+
+
+def push_history(history, latest, length, terms):
+  """Term-wise histories, oldest to newest, matching mjlab flattening."""
+  result, hidx, fidx = [], 0, 0
+  for dim in terms:
+    result.append(torch.cat((history[:, hidx + dim:hidx + dim * length],
+                             latest[:, fidx:fidx + dim]), -1))
+    hidx += dim * length
+    fidx += dim
+  return torch.cat(result, -1)
+
+
+def critic_input(state, command, context, previous):
+  return torch.cat((state, command, context, previous), -1)
+
+
+class Network(nn.Module):
+  def __init__(self, inputs, outputs, hidden):
+    super().__init__()
+    self.norm = RunningNorm(inputs)
+    layers = []
+    for n in hidden:
+      layers.extend((nn.Linear(inputs, n), nn.ELU()))
+      inputs = n
+    layers.append(nn.Linear(inputs, outputs))
+    self.net = nn.Sequential(*layers)
+
+  def forward(self, x):
+    return self.net(self.norm(x))
+
+
+class Actor(Network):
+  def __init__(self, history=5, hidden=(256, 128, 64), actor_wheel_vel=True):
+    self.actor_wheel_vel = bool(actor_wheel_vel)
+    self.terms = LEGACY_TERMS if self.actor_wheel_vel else TERMS
+    self.frame_dim = sum(self.terms)
+    super().__init__(self.frame_dim * history, 16, hidden)
+    self.history, self.hidden = history, tuple(hidden)
+    nn.init.uniform_(self.net[-1].weight, -0.01, 0.01)
+    nn.init.zeros_(self.net[-1].bias)
+
+  def forward(self, obs):
+    # Same bounded action distribution during collection and imagination.
+    return super().forward(obs).tanh()
+
+
+def checkpoint_wheel_vel(checkpoint):
+  """Infer legacy observation schema from weights; validate explicit metadata."""
+  flag = checkpoint.get("actor_wheel_vel")
+  if "actor" in checkpoint:
+    dim = checkpoint["actor"]["net.0.weight"].shape[1]
+    history = checkpoint["history"]
+    if dim not in (FRAME_DIM * history, LEGACY_FRAME_DIM * history):
+      raise ValueError("Unsupported actor observation dimension")
+    inferred = dim == LEGACY_FRAME_DIM * history
+    if flag is not None and bool(flag) != inferred:
+      raise ValueError("Actor observation metadata disagrees with weights")
+    return inferred
+  return bool(flag) if flag is not None else False
+
+
+def load_actor(checkpoint, device="cpu"):
+  if checkpoint.get("format") not in (POLICY_FORMAT, "s10-terrain-policy-v1"):
+    raise ValueError("Expected a S10 policy checkpoint, not other-robot or dynamics-only weights")
+  actor = Actor(checkpoint["history"], checkpoint["actor_hidden"],
+                checkpoint_wheel_vel(checkpoint)).to(device)
+  actor.load_state_dict(checkpoint["actor"])
+  return actor.eval()
+
+
+def evaluation_env_cfg(checkpoint, level=0, kind=None):
+  """Select the matching task; terrain remains privileged at deployment."""
+  from s10_supermimic.tasks import get_task
+  actor_wheel_vel = checkpoint_wheel_vel(checkpoint)
+  if checkpoint.get("format") == "s10-terrain-policy-v1":
+    checkpoint_task = checkpoint.get("config", {}).get("task")
+    task = get_task(checkpoint_task) if checkpoint_task else None
+    if task is not None:
+      if task.terrain_revision != checkpoint["config"].get("terrain_revision", 1):
+        raise ValueError("Checkpoint task revision differs from its config")
+      if task.actor.wheel_vel != actor_wheel_vel:
+        raise ValueError("Checkpoint task Actor schema differs from its weights")
+      cfg = task.build_env_cfg(play=True, history=checkpoint["history"],
+                               seed=checkpoint["config"]["seed"])
+    else:
+      from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
+      cfg = s10_velocity_terrain_env_cfg(
+        play=True, history=checkpoint["history"],
+        seed=checkpoint["config"]["seed"],
+        revision=checkpoint["config"].get("terrain_revision", 1),
+        actor_wheel_vel=actor_wheel_vel)
+    # A single staircase can overlap the robot's pre-reset origin at compile
+    # time. Reserve contacts for that transient pose as well as normal rolling.
+    cfg.sim.nconmax = max(cfg.sim.nconmax or 0, 512)
+    cfg.sim.njmax = max(cfg.sim.njmax or 0, 2048)
+    if level not in range(6):
+      raise ValueError("Terrain level must be 0..5")
+    gen = cfg.scene.terrain.terrain_generator
+    # Evaluation fixes difficulty; one row, same seeded generator.
+    gen.num_rows = 1
+    gen.difficulty_range = (level/5,level/5)
+    if kind is not None:
+      gen.sub_terrains = {kind: gen.sub_terrains[kind]}
+    return cfg
+  from s10_supermimic.env.flat import s10_velocity_flat_env_cfg
+  return s10_velocity_flat_env_cfg(play=True, history=checkpoint["history"],
+                                   actor_wheel_vel=actor_wheel_vel)