Initial commit

s10_supermimic/tasks/__init__.py

@@ -0,0 +1,12 @@
+"""Selectable S10 training tasks.
+
+Import a task module here to register it. The training entry point accepts
+``--task <name>`` and applies that task's defaults before explicit CLI flags.
+Shared algorithm code stays revision-agnostic; only the observation contract
+and environment wiring differ between tasks.
+"""
+
+from .base import ActorSpec, TaskSpec, get_task, list_tasks, register, task_names
+from . import task_v4, task_v4_upward, task_v7
+
+__all__ = ("ActorSpec", "TaskSpec", "get_task", "list_tasks", "register", "task_names")

s10_supermimic/tasks/base.py

@@ -0,0 +1,95 @@
+"""Parallel S10 training tasks: architecture, environment and recipe.
+
+A task owns the observation contract (the part that cannot be changed by a
+resume) and the training defaults that go with its terrain revision. Shared
+learner, reward and world-model code lives outside this package, so adding a
+new task means adding one module and registering it in ``tasks/__init__.py``.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+
+@dataclass(frozen=True)
+class ActorSpec:
+  """Actor observation contract used by training, export and playback."""
+  name: str
+  terms: tuple[int, ...]
+  frame_terms: tuple[dict[str, Any], ...]
+  wheel_vel: bool
+
+  @property
+  def frame_dim(self) -> int:
+    return sum(self.terms)
+
+  def history_dim(self, history: int) -> int:
+    if history < 1:
+      raise ValueError("history must be positive")
+    return self.frame_dim * history
+
+  def metadata(self, history: int) -> dict[str, Any]:
+    return {
+      "actor_schema": self.name,
+      "actor_wheel_vel": self.wheel_vel,
+      "frame_dim": self.frame_dim,
+      "observation_dim": self.history_dim(history),
+      "frame_terms": [dict(term) for term in self.frame_terms],
+    }
+
+
+@dataclass(frozen=True)
+class TaskSpec:
+  """One selectable training task.
+
+  ``defaults`` are argparse destinations applied before parsing, so explicit
+  command-line flags still override them. That keeps a task reproducible while
+  allowing a one-off run to change length or batch size.
+  """
+  name: str
+  description: str
+  terrain_revision: int
+  actor: ActorSpec
+  defaults: dict[str, Any] = field(default_factory=dict)
+  env_cfg: Callable[..., Any] | None = None
+
+  @property
+  def contacts(self) -> bool:
+    return self.terrain_revision >= 4
+
+  @property
+  def courses(self) -> bool:
+    return self.terrain_revision >= 5
+
+  @property
+  def collisions(self) -> bool:
+    return self.terrain_revision >= 7
+
+  def build_env_cfg(self, play: bool, history: int, seed: int):
+    if self.env_cfg is None:
+      raise ValueError(f"Task {self.name} has no environment builder")
+    return self.env_cfg(play=play, history=history, seed=seed)
+
+
+_REGISTRY: dict[str, TaskSpec] = {}
+
+
+def register(task: TaskSpec) -> TaskSpec:
+  if task.name in _REGISTRY:
+    raise ValueError(f"Duplicate task name: {task.name}")
+  _REGISTRY[task.name] = task
+  return task
+
+
+def get_task(name: str) -> TaskSpec:
+  try:
+    return _REGISTRY[name]
+  except KeyError as exc:
+    raise KeyError(f"Unknown task {name!r}; available: {', '.join(task_names())}") from exc
+
+
+def list_tasks() -> tuple[TaskSpec, ...]:
+  return tuple(_REGISTRY.values())
+
+
+def task_names() -> tuple[str, ...]:
+  return tuple(_REGISTRY)

s10_supermimic/tasks/task_v4.py

@@ -0,0 +1,65 @@
+"""S10 terrain revision 4: legacy 53-value Actor frame with measured wheel speed."""
+
+from .base import ActorSpec, TaskSpec, register
+
+ACTOR = ActorSpec(
+  name="s10-v4-53",
+  terms=(3, 3, 12, 12, 4, 16, 3),
+  frame_terms=(
+    {"name": "base_ang_vel", "dim": 3, "scale": 1.0},
+    {"name": "projected_gravity", "dim": 3, "scale": 1.0},
+    {"name": "joint_pos_rel", "dim": 12, "scale": 1.0},
+    {"name": "joint_vel_rel", "dim": 12, "scale": 0.05},
+    {"name": "wheel_vel_rel", "dim": 4, "scale": 0.05},
+    {"name": "last_action", "dim": 16, "scale": 1.0},
+    {"name": "command", "dim": 3, "scale": 1.0},
+  ),
+  wheel_vel=True,
+)
+
+
+def env_cfg(play: bool, history: int, seed: int):
+  from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
+  return s10_velocity_terrain_env_cfg(
+    play=play, history=history, seed=seed, revision=4,
+    actor_wheel_vel=ACTOR.wheel_vel)
+
+
+TASK = register(TaskSpec(
+  name="s10-v4",
+  description="S10 terrain v4 legacy task: 53-value Actor frame, wheel speed, 20k validated recipe",
+  terrain_revision=4,
+  actor=ACTOR,
+  env_cfg=env_cfg,
+  defaults={
+    "terrain": True,
+    "terrain_revision": 4,
+    "num_envs": 4096,
+    "iterations": 20000,
+    "num_steps_per_env": 1,
+    "buffer_len": 512,
+    "wm_window": 8,
+    "wm_batch": 2048,
+    "wm_updates": 4,
+    "imag_batch": 2048,
+    "horizon": 32,
+    "warmup": 50,
+    "critic_epochs": 5,
+    "critic_minibatches": 4,
+    "save_interval": 500,
+    "log_interval": 10,
+    "command_curriculum": True,
+    "command_curriculum_start": 0,
+    "prioritized_replay": True,
+    "archive_size": 100000,
+    "actor_ramp_iterations": 1000,
+    "actor_hold_iterations": 0,
+    "lr_decay_iterations": 30000,
+    "lr_final_frac": 0.03,
+    "history": 5,
+    "wm_hidden": [512, 512, 512],
+    "actor_hidden": [256, 128, 64],
+    "critic_hidden": [256, 128, 64],
+    "compile_imagination": False,
+  },
+))

s10_supermimic/tasks/task_v4_upward.py

@@ -0,0 +1,172 @@
+"""S10 terrain v4 + upward recovery task.
+
+This is a separate task, not an inheritance of v4 or v7. It keeps the v4
+53-value Actor frame, terrain revision, command schedule, terrain proportions,
+rolling rewards and replay, and changes only:
+
+* the existing ``upright`` reward into a sign-aware upward reward (no new term);
+* clamps on the large raw penalty terms;
+* random roll/pitch/yaw at reset;
+* a 4-second sustained bad-orientation termination.
+
+Reference: IsaacLab S10 ``upward_env_cfg.py`` for the upward/recovery state
+machine only. Its speed ranges, reward set and terrain config are not copied.
+"""
+
+import math
+
+from .base import ActorSpec, TaskSpec, register
+
+ACTOR = ActorSpec(
+  name="s10-v4-upward-53",
+  terms=(3, 3, 12, 12, 4, 16, 3),
+  frame_terms=(
+    {"name": "base_ang_vel", "dim": 3, "scale": 1.0},
+    {"name": "projected_gravity", "dim": 3, "scale": 1.0},
+    {"name": "joint_pos_rel", "dim": 12, "scale": 1.0},
+    {"name": "joint_vel_rel", "dim": 12, "scale": 0.05},
+    {"name": "wheel_vel_rel", "dim": 4, "scale": 0.05},
+    {"name": "last_action", "dim": 16, "scale": 1.0},
+    {"name": "command", "dim": 3, "scale": 1.0},
+  ),
+  wheel_vel=True,
+)
+
+
+def env_cfg(play: bool, history: int, seed: int):
+  from mjlab.managers import RewardTermCfg, TerminationTermCfg
+  from mjlab.managers.scene_entity_config import SceneEntityCfg
+
+  from s10_supermimic.core.upward import (
+    BAD_ORIENTATION_WEIGHT, UPWARD_REWARD_CLAMPS, BadOrientationSustained,
+    bad_orientation_penalty, clamped_action_rate, clamped_dof_pos_limits,
+    clamped_joint_torques, clamped_leg_posture, clamped_rolling_cost,
+    clamped_standstill, reset_root_state_upward, upright_upward,
+    upward_curriculum,
+  )
+  from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
+
+  cfg = s10_velocity_terrain_env_cfg(
+    play=play, history=history, seed=seed, revision=4, actor_wheel_vel=True)
+  # A 10 s episode lets each env complete 20 episodes per curriculum step
+  # within a 20k-iteration run. The 4 s bad-orientation termination still
+  # leaves 6 s after a successful recovery. Play/evaluation keeps the normal
+  # terrain-course horizon.
+  if not play:
+    cfg.episode_length_s = 10.0
+
+  # Random orientation at reset. Training starts at +/-60 degrees and advances
+  # per environment through upward_curriculum(); play/eval uses the full sphere
+  # so the checkpoint is tested on the final distribution.
+  tilt = math.pi if play else math.radians(60.0)
+  pose = cfg.events["reset_base"].params["pose_range"]
+  pose.update(
+    z=(0.0, 0.2),
+    roll=(-tilt, tilt),
+    pitch=(-tilt, tilt),
+    yaw=(-math.pi, math.pi),
+  )
+  cfg.events["reset_base"].func = reset_root_state_upward
+
+  # Existing upright term becomes the normalized upward reward. Max rate 2.0,
+  # i.e. 0.04 per 20 ms step, the same ratio to tracking (2.0 each) as IsaacLab.
+  upright = cfg.rewards["upright"]
+  upright_asset = upright.params.get(
+    "asset_cfg", SceneEntityCfg("robot", body_names=("base_link",)))
+  upright.func = upright_upward
+  upright.weight = 2.0
+  upright.params = {"asset_cfg": upright_asset}
+
+  # IsaacLab-style recovery shaping: no reward for termination itself, but a
+  # continuous per-step penalty while the orientation is bad. Weight -40 keeps
+  # the same 20x bad/upward ratio as IsaacLab (-80 vs max 4; here -40 vs max 2).
+  cfg.rewards["terminated"].weight = 0.0
+  cfg.rewards["bad_orientation_penalty"] = RewardTermCfg(
+    func=bad_orientation_penalty,
+    weight=BAD_ORIENTATION_WEIGHT,
+    params={"asset_cfg": SceneEntityCfg("robot")},
+  )
+
+  # Clamp the large raw penalties. The weights and the rest of the v4 reward
+  # set are unchanged; the caps only prevent recovery poses from dominating.
+  leg = cfg.rewards["leg_posture"]
+  leg.params = {"asset_cfg": leg.params["asset_cfg"],
+                "max_value": UPWARD_REWARD_CLAMPS["leg_posture"]}
+  leg.func = clamped_leg_posture
+
+  standstill = cfg.rewards["standstill"]
+  standstill.params = {"wheel_cfg": standstill.params["wheel_cfg"],
+                       "max_value": UPWARD_REWARD_CLAMPS["standstill"]}
+  standstill.func = clamped_standstill
+
+  cfg.rewards["action_rate_l2"].func = clamped_action_rate
+  cfg.rewards["action_rate_l2"].params = {
+    "max_value": UPWARD_REWARD_CLAMPS["action_rate_l2"]}
+
+  cfg.rewards["joint_torques_l2"].func = clamped_joint_torques
+  cfg.rewards["joint_torques_l2"].params = {
+    "max_value": UPWARD_REWARD_CLAMPS["joint_torques_l2"]}
+
+  limits = cfg.rewards["dof_pos_limits"]
+  limits.params = {"asset_cfg": limits.params["asset_cfg"],
+                   "max_value": UPWARD_REWARD_CLAMPS["dof_pos_limits"]}
+  limits.func = clamped_dof_pos_limits
+
+  for name in ("unnecessary_wheel_lift", "rolling_leg_motion"):
+    term = cfg.rewards[name]
+    term.func = clamped_rolling_cost
+    term.params = {"term": name, "max_value": UPWARD_REWARD_CLAMPS[name]}
+
+  # Replace v4's immediate 70-degree fall with the upward recovery state
+  # machine: bad orientation is tolerated for 4 s, then the episode ends.
+  cfg.terminations["fell_over"] = TerminationTermCfg(
+    func=BadOrientationSustained,
+    params={"duration_s": 4.0, "asset_cfg": SceneEntityCfg("robot")},
+  )
+
+  # Each environment advances its own terrain/tilt level from its last 20
+  # recovery outcomes (16/20 promotes); there is no global batch gate.
+  if not play:
+    cfg.curriculum["course"].func = upward_curriculum
+  return cfg
+
+
+TASK = register(TaskSpec(
+  name="s10-v4-upward",
+  description="S10 v4 recovery task: upward reward, monotonic tilt severity, per-env 20-episode curriculum, horizon 64",
+  terrain_revision=4,
+  actor=ACTOR,
+  env_cfg=env_cfg,
+  defaults={
+    "terrain": True,
+    "terrain_revision": 4,
+    "num_envs": 4096,
+    "iterations": 30000,
+    "num_steps_per_env": 1,
+    "buffer_len": 512,
+    "wm_window": 8,
+    "wm_batch": 2048,
+    "wm_updates": 4,
+    "imag_batch": 2048,
+    "horizon": 64,
+    "warmup": 50,
+    "action_noise": 0.22,
+    "critic_epochs": 5,
+    "critic_minibatches": 4,
+    "save_interval": 500,
+    "log_interval": 10,
+    "command_curriculum": True,
+    "command_curriculum_start": 0,
+    "prioritized_replay": True,
+    "archive_size": 100000,
+    "actor_ramp_iterations": 1000,
+    "actor_hold_iterations": 0,
+    "lr_decay_iterations": 30000,
+    "lr_final_frac": 0.03,
+    "history": 5,
+    "wm_hidden": [512, 512, 512],
+    "actor_hidden": [256, 128, 64],
+    "critic_hidden": [256, 128, 64],
+    "compile_imagination": False,
+  },
+))

s10_supermimic/tasks/task_v7.py

@@ -0,0 +1,78 @@
+"""S10 terrain revision 7: safe low-speed all-terrain locomotion.
+
+V7 keeps the 53-value v4 Actor frame (measured wheel speed is required for
+wheel slip/stall feedback), a stair-pass-gated command stage ending at
+1.5 m/s, 0.8 m/s and 1.2 rad/s with a 0.05 per-axis deadband, a conservative
+initialization (small x/y/yaw only, no random roll/pitch), a hipx 0.05 /
+hipy +/-25 deg / knee 1.0 default stance, always-on wheel/body projection
+penalties (per-wheel outer-1/4 and paired-leg 5/6 with a 0.5 s grace), reduced
+knee/hipy posture stiffness on stairs, and the v5 stair curriculum plus v7
+non-wheel collision head (50 N, wheels/knees excluded). Only ``imag_batch`` is
+reduced to 2048 to fit a 16 GB GPU.
+"""
+
+from .base import ActorSpec, TaskSpec, register
+
+ACTOR = ActorSpec(
+  name="s10-v7-53",
+  terms=(3, 3, 12, 12, 4, 16, 3),
+  frame_terms=(
+    {"name": "base_ang_vel", "dim": 3, "scale": 1.0},
+    {"name": "projected_gravity", "dim": 3, "scale": 1.0},
+    {"name": "joint_pos_rel", "dim": 12, "scale": 1.0},
+    {"name": "joint_vel_rel", "dim": 12, "scale": 0.05},
+    {"name": "wheel_vel_rel", "dim": 4, "scale": 0.05},
+    {"name": "last_action", "dim": 16, "scale": 1.0},
+    {"name": "command", "dim": 3, "scale": 1.0},
+  ),
+  wheel_vel=True,
+)
+
+
+def env_cfg(play: bool, history: int, seed: int):
+  from s10_supermimic.env.terrain import s10_velocity_terrain_env_cfg
+  return s10_velocity_terrain_env_cfg(
+    play=play, history=history, seed=seed, revision=7,
+    actor_wheel_vel=ACTOR.wheel_vel)
+
+
+TASK = register(TaskSpec(
+  name="s10-v7",
+  description="S10 terrain v7 safe all-terrain task: 53-value Actor frame, wheel speed, stair-gated low-speed commands, stair curriculum + collision",
+  terrain_revision=7,
+  actor=ACTOR,
+  env_cfg=env_cfg,
+  defaults={
+    "terrain": True,
+    "terrain_revision": 7,
+    "num_envs": 4096,
+    "iterations": 40000,
+    "num_steps_per_env": 1,
+    "buffer_len": 512,
+    "wm_window": 8,
+    "wm_batch": 2048,
+    "wm_updates": 4,
+    # 8192 is the full-size imagination batch; 2048 keeps the 16 GB footprint.
+    "imag_batch": 2048,
+    "horizon": 32,
+    "warmup": 300,
+    "critic_epochs": 5,
+    "critic_minibatches": 4,
+    "save_interval": 500,
+    "log_interval": 50,
+    # V7 uses stair-pass-gated stages, capped at its safe envelope.
+    "command_curriculum": True,
+    "command_curriculum_start": 0,
+    "prioritized_replay": True,
+    "archive_size": 100000,
+    "actor_ramp_iterations": 1000,
+    "actor_hold_iterations": 20000,
+    "lr_decay_iterations": 20000,
+    "lr_final_frac": 0.2,
+    "history": 5,
+    "wm_hidden": [512, 512, 512],
+    "actor_hidden": [256, 128, 64],
+    "critic_hidden": [256, 128, 64],
+    "compile_imagination": True,
+  },
+))