File size: 1,315 Bytes
7399b6f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | """Task registry, in the shape ManiSkill uses: one class per task, registered by name.
@register_task("grape_box")
class GrapeBoxTask(YamTask):
...
`make("grape_box")` then builds it. Importing `bimanual.yam_tasks.tasks` pulls in every task
module, which is what fills this registry -- the same trick as ManiSkill's `envs/tasks/__init__`.
"""
from __future__ import annotations
REGISTRY: dict[str, type] = {}
def register_task(name: str, **meta):
"""Class decorator: put a task class in the registry under `name`."""
def deco(cls):
if name in REGISTRY:
raise ValueError(f"task {name!r} is already registered by {REGISTRY[name].__name__}")
cls.task_name = name
cls.meta = meta
REGISTRY[name] = cls
return cls
return deco
def make(name: str, **kwargs):
if name not in REGISTRY:
raise SystemExit(f"unknown task {name!r}. Registered: {sorted(REGISTRY)}")
return REGISTRY[name](**kwargs)
def list_tasks():
rows = []
for n, cls in sorted(REGISTRY.items()):
rows.append({"name": n, "class": cls.__name__,
"title": getattr(cls, "title", ""),
"tags": ",".join(getattr(cls, "tags", [])),
"skill": getattr(cls, "skill", "")})
return rows
|