Initial commit

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)