diff --git a/examples/inverted_pendulum/src/genelab_inverted_pendulum/tasks.py b/examples/inverted_pendulum/src/genelab_inverted_pendulum/tasks.py index 1c8ef49c..9d9d13b8 100644 --- a/examples/inverted_pendulum/src/genelab_inverted_pendulum/tasks.py +++ b/examples/inverted_pendulum/src/genelab_inverted_pendulum/tasks.py @@ -134,6 +134,12 @@ def register() -> None: InvertedPendulumTask, description="PPO inverted-pendulum balancing (single pole).", cfg_type=TaskCfg, + examples=[ + f"genelab play {INVERTED_PENDULUM_TASK_ID} --vis", + f"genelab play {INVERTED_PENDULUM_TASK_ID} --agent trained " + "--checkpoint PATH/model.pt", + f"genelab train {INVERTED_PENDULUM_TASK_ID} --num_envs 4096 --max_iterations 100", + ], ) if DOUBLE_PENDULUM_TASK_ID not in TASKS: register_task( @@ -141,4 +147,10 @@ def register() -> None: DoubleInvertedPendulumTask, description="PPO double-inverted-pendulum balancing (two stacked poles).", cfg_type=TaskCfg, + examples=[ + f"genelab play {DOUBLE_PENDULUM_TASK_ID} --vis", + f"genelab play {DOUBLE_PENDULUM_TASK_ID} --agent trained " + "--checkpoint PATH/model.pt", + f"genelab train {DOUBLE_PENDULUM_TASK_ID} --num_envs 4096 --max_iterations 200", + ], ) diff --git a/examples/unitree/src/genelab_unitree/tasks.py b/examples/unitree/src/genelab_unitree/tasks.py index 5a7cb2d0..06d8af2b 100644 --- a/examples/unitree/src/genelab_unitree/tasks.py +++ b/examples/unitree/src/genelab_unitree/tasks.py @@ -125,11 +125,28 @@ def register() -> None: G1VelocityTask, description="PPO velocity tracking for Unitree G1 (flat).", cfg_type=TaskCfg, + examples=[ + f"genelab play {VELOCITY_TASK_ID}", + f"genelab play {VELOCITY_TASK_ID} --agent random", + f"genelab play {VELOCITY_TASK_ID} --agent trained --checkpoint PATH/model.pt", + f"genelab train {VELOCITY_TASK_ID} --num_envs 4096", + f"genelab train {VELOCITY_TASK_ID} --num_envs 4096 --gpus 2", + ], ) if TRACKING_TASK_ID not in TASKS: register_task( TRACKING_TASK_ID, G1TrackingTask, - description="PPO motion imitation for Unitree G1 (flat); requires --env.commands.motion.motion_file.", + description=( + "PPO motion imitation for Unitree G1 on flat ground. " + "Requires a motion clip; see examples." + ), cfg_type=TaskCfg, + examples=[ + f"genelab play {TRACKING_TASK_ID} --env.commands.motion.motion_file PATH/clip.npy", + f"genelab train {TRACKING_TASK_ID} " + "--env.commands.motion.motion_file PATH/clip.npy --num_envs 4096", + f"genelab train {TRACKING_TASK_ID} " + "--env.commands.motion.motion_file PATH/clip.npy --num_envs 4096 --gpus 2", + ], ) diff --git a/src/genelab/cli/__init__.py b/src/genelab/cli/__init__.py index a88b792c..051c497a 100644 --- a/src/genelab/cli/__init__.py +++ b/src/genelab/cli/__init__.py @@ -20,7 +20,8 @@ from genelab.cli._interactive import pick_task_interactively from genelab.cli._render import ( render_cache, - render_hint, + render_entry_info, + render_main_help, render_registry, ) from genelab.cli._scaffold import create_project_skeleton @@ -67,23 +68,56 @@ class _RegistryKindArg(str, Enum): ) +_RUN_FLAGS_HELP: Final[str] = """\ +Shorthand flags rewritten into env overrides: + +\b + -v, --vis Enable the Genesis viewer (env.scene.vis=true). + --gpu Use the GPU backend (env.scene.gpu=true). + --steps N Run for N steps (env.scene.steps=N). + --dt SECONDS Override the sim timestep (env.scene.dt=SECONDS). + --a.b.c VALUE Set any dotted cfg path. + +Runner flags (used when an RL runner is engaged): + +\b + --num_envs N Parallel environments. + --agent KIND one of: zero, random, trained (play only). + --checkpoint PATH Resume from a checkpoint. + --seed N RNG seed. + --log_dir PATH Override the log directory. + --max_iterations N Cap training iterations (train only). + --gpus N Distributed training across N GPUs (train only). + +Use `genelab info TASK` to see the full overridable path list for a task. +""" + + +_PLAY_HELP: Final[str] = "Run a registered task.\n\n" + _RUN_FLAGS_HELP +_TRAIN_HELP: Final[str] = "Train a registered task when a runner exists.\n\n" + _RUN_FLAGS_HELP + + app = typer.Typer( name="genelab", - help="Run registered GeneLab tasks.", + help=( + "GeneLab — Genesis robot lab CLI.\n\n" + "Run `genelab` with no arguments for a landing page with quickstart commands " + "and a count of registered robots, envs, and tasks." + ), no_args_is_help=False, add_completion=False, pretty_exceptions_enable=False, - rich_markup_mode=None, + rich_markup_mode="rich", ) project_app = typer.Typer( name="project", - help="Create and manage GeneLab projects.", + help="Create and manage GeneLab extension projects.", no_args_is_help=True, pretty_exceptions_enable=False, - rich_markup_mode=None, + rich_markup_mode="rich", ) -app.add_typer(project_app, name="project") +app.add_typer(project_app, name="project", rich_help_panel="Project") def _version_callback(value: bool) -> None: @@ -109,33 +143,43 @@ def root_callback( typer.Option( "--import", metavar="MODULE", - help="Import a downstream extension module before dispatch (can be repeated).", + help="Import an extension module before dispatch (repeatable).", ), ] = None, no_entry_points: Annotated[ bool, typer.Option( "--no-entry-points", - help="Skip installed extensions from the genelab.extensions entry point group.", + help="Skip installed entry points from the genelab.extensions group.", ), ] = False, ) -> None: _ = version # consumed by the eager callback - ctx.obj = _RootState( + state = _RootState( extension_modules=list(extension_modules or []), no_entry_points=no_entry_points, ) + ctx.obj = state if ctx.invoked_subcommand is None: - render_hint() + _load_extensions(state) + render_main_help() -@app.command("cache", help="Create project-local simulation cache directories.") +@app.command( + "cache", + help="Create project-local simulation cache directories.", + rich_help_panel="Utilities", +) def cache_cmd() -> None: ensure_project_cache() render_cache(CACHE_DIR) -@app.command("list", help="List registered robots, environments, or tasks.") +@app.command( + "list", + help="List registered robots, envs, or tasks.", + rich_help_panel="Registry", +) def list_cmd( ctx: typer.Context, kind: Annotated[ @@ -147,9 +191,29 @@ def list_cmd( render_registry(kind.value) +@app.command( + "info", + help="Show detail for one registered task, env, or robot.", + rich_help_panel="Registry", +) +def info_cmd( + ctx: typer.Context, + name: Annotated[ + str, + typer.Argument( + metavar="NAME", + help="Registered task, env, or robot name.", + ), + ], +) -> None: + _load_extensions(_state(ctx)) + render_entry_info(name) + + @app.command( "play", - help="Run a registered task.", + help=_PLAY_HELP, + rich_help_panel="Runtime", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}, ) def play_cmd(ctx: typer.Context) -> None: @@ -163,7 +227,8 @@ def play_cmd(ctx: typer.Context) -> None: @app.command( "train", - help="Train a registered task when a runner exists.", + help=_TRAIN_HELP, + rich_help_panel="Runtime", context_settings={"allow_extra_args": True, "ignore_unknown_options": True}, ) def train_cmd(ctx: typer.Context) -> None: @@ -185,7 +250,7 @@ def project_new_cmd( typer.Option( "--path", "-p", - help="Parent directory where the project directory is created.", + help="Parent directory under which the project directory is created.", ), ] = Path("."), package: Annotated[ diff --git a/src/genelab/cli/_render.py b/src/genelab/cli/_render.py index 4602d9f0..a9e787d2 100644 --- a/src/genelab/cli/_render.py +++ b/src/genelab/cli/_render.py @@ -1,23 +1,21 @@ """Rich-based rendering helpers for the GeneLab CLI.""" -from dataclasses import asdict, is_dataclass +from collections.abc import Iterator +from dataclasses import asdict, fields, is_dataclass from pathlib import Path from typing import Literal, Protocol, runtime_checkable -from rich.console import Console +from rich.console import Console, Group, RenderableType from rich.panel import Panel +from rich.table import Table +from rich.text import Text from rich.theme import Theme -from genelab.registry import ENVS, ROBOTS, TASKS +from genelab.registry import ENVS, ROBOTS, TASKS, Registry, RegistryEntry RegistryKind = Literal["robots", "envs", "tasks"] -@runtime_checkable -class _TaskLike(Protocol): - cfg: object - - @runtime_checkable class _CfgBacked(Protocol): cfg: object @@ -33,17 +31,63 @@ class _CfgBacked(Protocol): "error": "bold red", "ok": "bold green", "step": "cyan", + "panel.title": "bold cyan", + "panel.border": "cyan", + "table.column": "bold white", } ) -console: Console = Console(theme=_theme, highlight=False, soft_wrap=True) -err_console: Console = Console(stderr=True, theme=_theme, highlight=False, soft_wrap=True) +console: Console = Console(theme=_theme, highlight=False) +err_console: Console = Console(stderr=True, theme=_theme, highlight=False) + + +def render_main_help() -> None: + """Print the landing page shown when ``genelab`` is invoked with no subcommand.""" + robot_count = len(ROBOTS.entries()) + env_count = len(ENVS.entries()) + task_count = len(TASKS.entries()) -def render_hint() -> None: + console.print() + console.print("[entry.name]GeneLab[/] — Genesis robot lab CLI") + console.print() + console.print("[registry.kind]Quickstart[/]") console.print( - "Registered GeneLab tasks. Try [hint]genelab list tasks[/] or import an extension package." + " [step]genelab cache[/] " + "[entry.desc]Create project-local cache directories.[/]" ) + console.print(" [step]genelab list tasks[/] [entry.desc]Show registered tasks.[/]") + console.print( + " [step]genelab info NAME[/] [entry.desc]Show detail for a registered name.[/]" + ) + console.print( + " [step]genelab play TASK[/] [entry.desc]Run a task; add --vis for a viewer.[/]" + ) + console.print( + " [step]genelab project new NAME[/] [entry.desc]Scaffold a new extension package.[/]" + ) + console.print() + console.print("[registry.kind]Commands[/]") + console.print(" [entry.name]Registry[/] list, info") + console.print(" [entry.name]Runtime[/] play, train") + console.print(" [entry.name]Project[/] project new") + console.print(" [entry.name]Utilities[/] cache") + console.print() + console.print( + f"[entry.detail]Registered:[/] {robot_count} robots, {env_count} envs, {task_count} tasks" + ) + if robot_count == 0 and env_count == 0 and task_count == 0: + console.print( + "[hint]No extensions loaded.[/] Try " + "[step]genelab --import MODULE ...[/] " + "or install an extension package." + ) + else: + console.print( + "[entry.detail]Drill in with [step]genelab info NAME[/] for fields and examples.[/]" + ) + console.print("[entry.detail]Use [step]genelab --help[/] for the full command reference.[/]") + console.print() def render_registry(kind: RegistryKind) -> None: @@ -55,11 +99,35 @@ def render_registry(kind: RegistryKind) -> None: return for entry in entries: details = _entry_details(kind, entry.name) + if entry.examples: + suffix = f" ({len(entry.examples)} example{'s' if len(entry.examples) != 1 else ''})" + details = f"{details}{suffix}" if details else suffix.strip() console.print(f" - [entry.name]{entry.name}[/]: [entry.desc]{entry.description}[/]") if details: console.print(f" [entry.detail]{details}[/]") +def render_entry_info(name: str) -> None: + """Render a detail panel for a registered task, env, or robot. + + Looks up ``name`` in the task / env / robot registries in that order. Raises + ``SystemExit`` with the combined available-name listing when nothing matches. + """ + + candidates: tuple[tuple[str, Registry[object]], ...] = ( + ("task", TASKS), + ("env", ENVS), + ("robot", ROBOTS), + ) + for kind, registry in candidates: + if name in registry: + _render_entry_panel(kind, registry.entry(name)) + return + available = sorted({*TASKS.names(), *ENVS.names(), *ROBOTS.names()}) + listing = ", ".join(available) if available else "" + raise SystemExit(f"unknown name {name!r}; available: {listing}") + + def render_cache(cache_dir: Path) -> None: console.print(f"Using project cache at [ok]{cache_dir}[/]") @@ -72,19 +140,125 @@ def render_project_created(target: Path, task_id: str) -> None: f"[step]uv run genelab play[/] {task_id}" ) console.print(f"Created GeneLab extension project at [ok]{target}[/]") - console.print(Panel(body, title="Next steps", border_style="entry.detail")) + console.print(Panel(body, title="Next steps", border_style="panel.border")) def render_error(message: str) -> None: err_console.print(f"[error]error:[/] {message}") +def iter_overridable_paths(value: object, prefix: str = "") -> Iterator[tuple[str, str, str]]: + """Yield ``(dotted_path, type_name, default_repr)`` for each dataclass field. + + Walks nested dataclasses depth-first. Fields whose name starts with ``_`` are + skipped. The runtime type of the current value is used for the type column, + so opaque ``object``-typed fields (e.g. ``TaskCfg.env``) still surface their + real downstream dataclass. + """ + + if not is_dataclass(value) or isinstance(value, type): + return + for f in fields(value): + if f.name.startswith("_"): + continue + current = getattr(value, f.name) + path = f.name if not prefix else f"{prefix}.{f.name}" + if is_dataclass(current) and not isinstance(current, type): + yield from iter_overridable_paths(current, path) + continue + yield (path, _short_type_name(current), _format_default(current)) + + +def _render_entry_panel(kind: str, entry: RegistryEntry[object]) -> None: + rows: list[RenderableType] = [Text(entry.description, style="entry.desc")] + + instance: object | None = None + try: + instance = entry.factory() + except Exception: # pragma: no cover — factories with side effects + instance = None + + if kind == "task" and isinstance(instance, _CfgBacked): + cfg = instance.cfg + env_name = getattr(cfg, "env_name", None) + robot_name = getattr(cfg, "robot_name", None) + trainable = getattr(cfg, "trainable", None) + meta_bits: list[str] = [] + if env_name is not None: + meta_bits.append(f"env=[entry.name]{env_name}[/]") + if robot_name is not None: + meta_bits.append(f"robot=[entry.name]{robot_name}[/]") + if trainable is not None: + meta_bits.append(f"trainable=[entry.name]{trainable}[/]") + if meta_bits: + rows.append(Text("")) + rows.append(Text.from_markup(", ".join(meta_bits))) + + cfg_value = instance.cfg if isinstance(instance, _CfgBacked) else instance + if cfg_value is not None and is_dataclass(cfg_value) and not isinstance(cfg_value, type): + table = _overridable_paths_table(cfg_value) + if table is not None: + rows.append(Text("")) + rows.append(Text("Overridable cfg paths", style="entry.name")) + rows.append(table) + elif entry.cfg_type is None: + rows.append(Text("")) + rows.append(Text("(no cfg introspection available)", style="entry.detail")) + + if entry.examples: + rows.append(Text("")) + rows.append(Text("Examples", style="entry.name")) + for example in entry.examples: + rows.append(Text.from_markup(f" [step]{example}[/]")) + + title = f"[panel.title]{kind.capitalize()}:[/] [entry.name]{entry.name}[/]" + console.print(Panel(Group(*rows), title=title, border_style="panel.border")) + + +def _overridable_paths_table(cfg: object) -> Table | None: + paths = list(iter_overridable_paths(cfg)) + if not paths: + return None + table = Table.grid(padding=(0, 2)) + table.add_column("path", style="entry.name") + table.add_column("type", style="table.column") + table.add_column("default", style="entry.desc") + table.add_row("[table.column]Path[/]", "[table.column]Type[/]", "[table.column]Default[/]") + for dotted, type_name, default_repr in paths: + table.add_row(dotted, type_name, default_repr) + return table + + +def _short_type_name(value: object) -> str: + if value is None: + return "None" + return type(value).__name__ + + +def _format_default(value: object) -> str: + if value is None: + return "None" + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (int, float)): + return repr(value) + if isinstance(value, str): + return repr(value) + if isinstance(value, Path): + return str(value) + if isinstance(value, (list, tuple)): + if not value: + return "()" if isinstance(value, tuple) else "[]" + return f"<{type(value).__name__} of {len(value)}>" + return type(value).__name__ + + def _entry_details(kind: RegistryKind, name: str) -> str: try: value = {"robots": ROBOTS, "envs": ENVS, "tasks": TASKS}[kind].get(name) except Exception: return "" - if kind == "tasks" and isinstance(value, _TaskLike): + if kind == "tasks" and isinstance(value, _CfgBacked): cfg = value.cfg env_name = getattr(cfg, "env_name", None) robot_name = getattr(cfg, "robot_name", None) diff --git a/src/genelab/cli/_scaffold.py b/src/genelab/cli/_scaffold.py index 29201fcc..2b83cd66 100644 --- a/src/genelab/cli/_scaffold.py +++ b/src/genelab/cli/_scaffold.py @@ -262,6 +262,10 @@ def register() -> None: ExampleTask, description="Task provided by this GeneLab extension project.", cfg_type=TaskCfg, + examples=[ + f"genelab play {{TASK_ID}}", + f"genelab play {{TASK_ID}} --vis --steps 200", + ], ) """ ) diff --git a/src/genelab/registry.py b/src/genelab/registry.py index 73700d39..be1c7893 100644 --- a/src/genelab/registry.py +++ b/src/genelab/registry.py @@ -1,7 +1,7 @@ """Small registries for robots, environments, tasks, and downstream extensions.""" -from collections.abc import Callable, Iterable -from dataclasses import dataclass +from collections.abc import Callable, Iterable, Sequence +from dataclasses import dataclass, field import importlib from importlib import metadata from pathlib import Path @@ -17,6 +17,7 @@ class RegistryEntry[T]: description: str factory: Callable[[], T] cfg_type: type[object] | None = None + examples: tuple[str, ...] = field(default_factory=tuple) class Registry[T]: @@ -33,11 +34,16 @@ def register( *, description: str, cfg_type: type[object] | None = None, + examples: Sequence[str] = (), ) -> RegistryEntry[T]: if name in self._entries: raise ValueError(f"{self.kind} already registered: {name}") entry = RegistryEntry( - name=name, description=description, factory=factory, cfg_type=cfg_type + name=name, + description=description, + factory=factory, + cfg_type=cfg_type, + examples=tuple(examples), ) self._entries[name] = entry return entry @@ -83,8 +89,15 @@ def register_robot[T]( *, description: str, cfg_type: type[object] | None = None, + examples: Sequence[str] = (), ) -> RegistryEntry[T]: - return ROBOTS.register(name, factory, description=description, cfg_type=cfg_type) + return ROBOTS.register( + name, + factory, + description=description, + cfg_type=cfg_type, + examples=examples, + ) def register_env[T]( @@ -93,8 +106,15 @@ def register_env[T]( *, description: str, cfg_type: type[object] | None = None, + examples: Sequence[str] = (), ) -> RegistryEntry[T]: - return ENVS.register(name, factory, description=description, cfg_type=cfg_type) + return ENVS.register( + name, + factory, + description=description, + cfg_type=cfg_type, + examples=examples, + ) def register_task[T]( @@ -103,8 +123,15 @@ def register_task[T]( *, description: str, cfg_type: type[object] | None = None, + examples: Sequence[str] = (), ) -> RegistryEntry[T]: - return TASKS.register(name, factory, description=description, cfg_type=cfg_type) + return TASKS.register( + name, + factory, + description=description, + cfg_type=cfg_type, + examples=examples, + ) def load_builtin_registries() -> None: diff --git a/tests/fake_extension.py b/tests/fake_extension.py index 8f4f2f2d..8f7bb237 100644 --- a/tests/fake_extension.py +++ b/tests/fake_extension.py @@ -59,4 +59,8 @@ def register() -> None: FakeTask, description="Task from a fake external package.", cfg_type=TaskCfg, + examples=[ + "genelab play External-Fake-Task-v0", + "genelab play External-Fake-Task-v0 --steps 7", + ], ) diff --git a/tests/test_cli.py b/tests/test_cli.py index f1fc74e4..b341ca01 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -1,3 +1,4 @@ +import re from collections.abc import Sequence from pathlib import Path @@ -35,6 +36,12 @@ type FloatArray = NDArray[np.floating] +_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") + + +def _strip_ansi(text: str) -> str: + return _ANSI_RE.sub("", text) + class _FakeLink: def __init__(self, q_start: int, idx: int | None = None) -> None: @@ -175,6 +182,96 @@ def test_cli_outputs_registered_hint(capsys: pytest.CaptureFixture[str]) -> None assert "genelab list tasks" in capsys.readouterr().out +def test_bare_invocation_prints_landing_page(capsys: pytest.CaptureFixture[str]) -> None: + main(["--no-entry-points"]) + + out = capsys.readouterr().out + assert "GeneLab" in out + assert "Quickstart" in out + assert "Commands" in out + assert "Registered:" in out + assert "genelab list tasks" in out + assert "genelab info NAME" in out + + +def test_play_help_documents_runner_keys(capsys: pytest.CaptureFixture[str]) -> None: + from genelab.cli import RUNNER_KEYS + + main(["play", "--help"]) + + out = _strip_ansi(capsys.readouterr().out) + for key in RUNNER_KEYS: + assert key in out, f"runner key {key!r} missing from `play --help`" + + +def test_play_help_documents_short_flag_grammar(capsys: pytest.CaptureFixture[str]) -> None: + main(["play", "--help"]) + + out = _strip_ansi(capsys.readouterr().out) + assert "--vis" in out + assert "--gpu" in out + assert "--steps" in out + assert "env.scene" in out + + +def test_register_task_accepts_examples_kwarg() -> None: + from genelab.registry import Registry + + isolated: Registry[object] = Registry("test-task") + entry = isolated.register( + "Examples-Roundtrip-v0", + lambda: None, + description="Examples round-trip test entry.", + examples=["genelab play Examples-Roundtrip-v0"], + ) + + assert entry.examples == ("genelab play Examples-Roundtrip-v0",) + + +def test_register_task_examples_default_to_empty_tuple() -> None: + from genelab.registry import Registry + + isolated: Registry[object] = Registry("test-task") + entry = isolated.register( + "Defaults-v0", + lambda: None, + description="Default examples test entry.", + ) + + assert entry.examples == () + + +def test_info_renders_examples_and_overrides(capsys: pytest.CaptureFixture[str]) -> None: + main( + [ + "--no-entry-points", + "--import", + "tests.fake_extension", + "info", + "External-Fake-Task-v0", + ] + ) + + out = capsys.readouterr().out + assert "External-Fake-Task-v0" in out + assert "Task from a fake external package." in out + assert "genelab play External-Fake-Task-v0" in out + assert "--steps 7" in out + # cfg introspection surfaces the scene fields that overrides walk through. + assert "env.scene.steps" in out + + +def test_info_unknown_name_errors(capsys: pytest.CaptureFixture[str]) -> None: + try: + main(["--no-entry-points", "info", "definitely-not-a-registered-name"]) + except SystemExit as exc: + assert "definitely-not-a-registered-name" in str(exc) + assert "available" in str(exc) + else: + raise AssertionError("expected info to exit on unknown name") + capsys.readouterr() # drain any output + + def test_core_does_not_register_example_tasks_by_default( capsys: pytest.CaptureFixture[str], ) -> None: