96 lines
2.6 KiB
Python
96 lines
2.6 KiB
Python
"""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)
|