diff --git a/README.md b/README.md index a3ef4083e..2d289ae8b 100644 --- a/README.md +++ b/README.md @@ -26,7 +26,7 @@ Run a pre-trained motion mimic policy on the Unitree G1 humanoid: ```bash uv run --extra rl scripts/tracking/rl/play.py \ - --task Tracking-Flat-G1-Play-v0 \ + --task Mjlab-Tracking-Flat-G1-Play-v0 \ --wandb-run-path gcbc_researchers/mjlab_alpha/rfdej55h ``` @@ -34,7 +34,7 @@ You can train this exact motion mimic policy using the following command: ```bash MUJOCO_GL=egl uv run --extra rl scripts/tracking/rl/train.py \ - --task Tracking-Flat-G1-v0 \ + --task Mjlab-Tracking-Flat-G1-v0 \ --registry-name gcbc_researchers/csv_to_npz/lafan_cartwheel \ --num_envs 4096 ``` diff --git a/scripts/list_envs.py b/scripts/list_envs.py new file mode 100644 index 000000000..f34ad31e5 --- /dev/null +++ b/scripts/list_envs.py @@ -0,0 +1,39 @@ +"""Script to list MJLab environments.""" + +import gymnasium as gym +from prettytable import PrettyTable + +import mjlab.tasks # noqa: F401 to register environments + + +def main(): + """Print all environments registered whose id contains `Mjlab-`.""" + prefix_substring = "Mjlab-" + + table = PrettyTable(["#", "Task ID", "Entry Point", "env_cfg_entry_point"]) + table.title = "Available Environments in Mjlab" + table.align["Task ID"] = "l" + table.align["Entry Point"] = "l" + table.align["env_cfg_entry_point"] = "l" + + idx = 0 + for spec in gym.registry.values(): + try: + if prefix_substring in spec.id: + env_cfg_ep = spec.kwargs.get("env_cfg_entry_point", "") + table.add_row([idx + 1, spec.id, spec.entry_point, env_cfg_ep]) + idx += 1 + except Exception: + continue + + print(table) + if idx == 0: + print(f"[INFO] No tasks matched filter: '{prefix_substring}'") + return idx + + +if __name__ == "__main__": + try: + main() + except Exception as e: + raise RuntimeError(f"Error listing environments: {e}") from e diff --git a/src/mjlab/tasks/locomotion/velocity/config/go1/__init__.py b/src/mjlab/tasks/locomotion/velocity/config/go1/__init__.py index b9930f7d3..757d3cd93 100644 --- a/src/mjlab/tasks/locomotion/velocity/config/go1/__init__.py +++ b/src/mjlab/tasks/locomotion/velocity/config/go1/__init__.py @@ -1,7 +1,7 @@ import gymnasium as gym gym.register( - id="Velocity-Flat-Unitree-Go1-v0", + id="Mjlab-Velocity-Flat-Unitree-Go1-v0", entry_point="mjlab.envs:ManagerBasedRlEnv", disable_env_checker=True, kwargs={ diff --git a/src/mjlab/tasks/tracking/config/g1/__init__.py b/src/mjlab/tasks/tracking/config/g1/__init__.py index 16cf9a4f3..f4ab2836a 100644 --- a/src/mjlab/tasks/tracking/config/g1/__init__.py +++ b/src/mjlab/tasks/tracking/config/g1/__init__.py @@ -1,7 +1,7 @@ import gymnasium as gym gym.register( - id="Tracking-Flat-G1-v0", + id="Mjlab-Tracking-Flat-G1-v0", entry_point="mjlab.envs:ManagerBasedRlEnv", disable_env_checker=True, kwargs={ @@ -11,7 +11,7 @@ ) gym.register( - id="Tracking-Flat-G1-Play-v0", + id="Mjlab-Tracking-Flat-G1-Play-v0", entry_point="mjlab.envs:ManagerBasedRlEnv", disable_env_checker=True, kwargs={ diff --git a/src/mjlab/tasks/tracking/config/t1/__init__.py b/src/mjlab/tasks/tracking/config/t1/__init__.py index 68551b315..985b2a499 100644 --- a/src/mjlab/tasks/tracking/config/t1/__init__.py +++ b/src/mjlab/tasks/tracking/config/t1/__init__.py @@ -1,7 +1,7 @@ import gymnasium as gym gym.register( - id="Tracking-Flat-T1-v0", + id="Mjlab-Tracking-Flat-T1-v0", entry_point="mjlab.envs:ManagerBasedRlEnv", disable_env_checker=True, kwargs={ @@ -11,7 +11,7 @@ ) gym.register( - id="Tracking-Flat-T1-Play-v0", + id="Mjlab-Tracking-Flat-T1-Play-v0", entry_point="mjlab.envs:ManagerBasedRlEnv", disable_env_checker=True, kwargs={