From 6564b4859bff455177f1c028f3fb2e3058d3e778 Mon Sep 17 00:00:00 2001 From: acrlw <13927622+acrlw@users.noreply.github.com> Date: Thu, 13 Aug 2026 11:14:34 +0800 Subject: [PATCH 1/7] feat(rl): add motion policy viewer --- docs/requirements.txt | 3 +- ...n.learning.rl.motion_policy_evaluation.rst | 15 + .../embodichain/embodichain.learning.rl.rst | 2 + .../embodichain.learning.rl.runtime.rst | 29 ++ docs/source/conf.py | 2 + docs/source/guides/cli.md | 70 +++ docs/source/guides/index.rst | 1 + .../source/guides/motion_policy_evaluation.md | 359 +++++++++++++ docs/source/tutorial/rl.rst | 23 +- embodichain/__main__.py | 5 + embodichain/learning/rl/evaluation.py | 45 +- .../rl/motion_policy_evaluation/__init__.py | 50 ++ .../rl/motion_policy_evaluation/bridge.py | 251 +++++++++ .../rl/motion_policy_evaluation/checkpoint.py | 60 +++ .../rl/motion_policy_evaluation/cli.py | 398 +++++++++++++++ .../rl/motion_policy_evaluation/manifest.py | 205 ++++++++ .../motion_policy_evaluation/native_task.py | 481 ++++++++++++++++++ .../rl/motion_policy_evaluation/profile.py | 148 ++++++ .../rl/motion_policy_evaluation/report.py | 78 +++ embodichain/learning/rl/runtime.py | 350 +++++++++++++ embodichain/learning/rl/train.py | 236 ++++----- .../motion_policy_evaluation/README.md | 130 +++++ .../anymal_c/__init__.py | 30 ++ .../anymal_c/profile.py | 295 +++++++++++ .../motion_policy_evaluation/eval_policy.py | 72 +++ .../prepare_resources.py | 243 +++++++++ .../rl/motion_policy_evaluation/__init__.py | 19 + .../test_anymal_c_example.py | 255 ++++++++++ .../motion_policy_evaluation/test_bridge.py | 159 ++++++ .../test_checkpoint.py | 30 ++ .../rl/motion_policy_evaluation/test_cli.py | 143 ++++++ .../motion_policy_evaluation/test_manifest.py | 49 ++ .../test_native_task.py | 251 +++++++++ .../motion_policy_evaluation/test_profile.py | 54 ++ .../test_train_integration.py | 43 ++ tests/learning/test_runtime.py | 166 ++++++ tests/test_main.py | 1 + 37 files changed, 4596 insertions(+), 155 deletions(-) create mode 100644 docs/source/api_reference/embodichain/embodichain.learning.rl.motion_policy_evaluation.rst create mode 100644 docs/source/api_reference/embodichain/embodichain.learning.rl.runtime.rst create mode 100644 docs/source/guides/motion_policy_evaluation.md create mode 100644 embodichain/learning/rl/motion_policy_evaluation/__init__.py create mode 100644 embodichain/learning/rl/motion_policy_evaluation/bridge.py create mode 100644 embodichain/learning/rl/motion_policy_evaluation/checkpoint.py create mode 100644 embodichain/learning/rl/motion_policy_evaluation/cli.py create mode 100644 embodichain/learning/rl/motion_policy_evaluation/manifest.py create mode 100644 embodichain/learning/rl/motion_policy_evaluation/native_task.py create mode 100644 embodichain/learning/rl/motion_policy_evaluation/profile.py create mode 100644 embodichain/learning/rl/motion_policy_evaluation/report.py create mode 100644 embodichain/learning/rl/runtime.py create mode 100644 examples/learning/motion_policy_evaluation/README.md create mode 100644 examples/learning/motion_policy_evaluation/anymal_c/__init__.py create mode 100644 examples/learning/motion_policy_evaluation/anymal_c/profile.py create mode 100644 examples/learning/motion_policy_evaluation/eval_policy.py create mode 100644 examples/learning/motion_policy_evaluation/prepare_resources.py create mode 100644 tests/learning/rl/motion_policy_evaluation/__init__.py create mode 100644 tests/learning/rl/motion_policy_evaluation/test_anymal_c_example.py create mode 100644 tests/learning/rl/motion_policy_evaluation/test_bridge.py create mode 100644 tests/learning/rl/motion_policy_evaluation/test_checkpoint.py create mode 100644 tests/learning/rl/motion_policy_evaluation/test_cli.py create mode 100644 tests/learning/rl/motion_policy_evaluation/test_manifest.py create mode 100644 tests/learning/rl/motion_policy_evaluation/test_native_task.py create mode 100644 tests/learning/rl/motion_policy_evaluation/test_profile.py create mode 100644 tests/learning/rl/motion_policy_evaluation/test_train_integration.py create mode 100644 tests/learning/test_runtime.py diff --git a/docs/requirements.txt b/docs/requirements.txt index 87db2c491..b4f8f14a9 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -8,4 +8,5 @@ sphinx-autosummary-accessors sphinxcontrib-bibtex sphinx-design sphinx_autodoc_typehints -pypandoc_binary \ No newline at end of file +sphinxcontrib-mermaid +pypandoc_binary diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.motion_policy_evaluation.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.motion_policy_evaluation.rst new file mode 100644 index 000000000..03d6beaf2 --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.learning.rl.motion_policy_evaluation.rst @@ -0,0 +1,15 @@ +embodichain.learning.rl.motion_policy_evaluation +================================================ + +.. automodule:: embodichain.learning.rl.motion_policy_evaluation + :members: + :undoc-members: + :show-inheritance: + +Native Task +----------- + +.. automodule:: embodichain.learning.rl.motion_policy_evaluation.native_task + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.rst index bf1b5b01e..1393f7505 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rl.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rl.rst @@ -18,6 +18,8 @@ collection logic, policy/model builders, and training entry points. buffer collector models + motion_policy_evaluation + runtime train utils diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.runtime.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.runtime.rst new file mode 100644 index 000000000..10d68dbef --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.learning.rl.runtime.rst @@ -0,0 +1,29 @@ +embodichain.learning.rl.runtime +=============================== + +.. automodule:: embodichain.learning.rl.runtime + + .. rubric:: Module Attributes + + .. autosummary:: + + GymPolicyFactory + + .. rubric:: Functions + + .. autosummary:: + + build_gym_environment + build_gym_policy_runtime + build_learning_environment + build_learning_policy_runtime + resolve_config_reference + resolve_torch_device + seed_policy_runtime + + .. rubric:: Classes + + .. autosummary:: + + GymEnvironmentRuntime + PolicyRuntime diff --git a/docs/source/conf.py b/docs/source/conf.py index 8065112b5..b7c1f4aa5 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -40,6 +40,7 @@ "sphinx.ext.viewcode", "sphinx_autodoc_typehints", # optional, shows type hints "sphinx_design", + "sphinxcontrib.mermaid", "myst_parser", # if you prefer Markdown pages "sphinx_copybutton", ] @@ -59,6 +60,7 @@ # If using MyST and writing .md API stubs: myst_enable_extensions = ["colon_fence", "deflist", "html_admonition"] +myst_fence_as_directive = ["mermaid"] templates_path = ["_templates"] diff --git a/docs/source/guides/cli.md b/docs/source/guides/cli.md index f009d64ab..30eb51bd1 100644 --- a/docs/source/guides/cli.md +++ b/docs/source/guides/cli.md @@ -372,6 +372,76 @@ See the Profiling section under Run Env for report format. Outputs are written t --- +## Motion Policy Evaluation + +Run visual evaluation for a simulator training run. EmbodiChain restores the +training Environment, observation, Policy, action manager, and task reset: + +```bash +embodichain eval-motion-policy outputs/my_policy_ \ + --checkpoint best \ + --device cuda \ + --sim-device gpu \ + --viewer +``` + +Evaluate an older EmbodiChain `.pt` checkpoint: + +```bash +embodichain eval-motion-policy \ + --checkpoint /path/to/policy.pt \ + --config /path/to/train.yaml \ + --gym-config /path/to/gym.yaml \ + --viewer +``` + +Evaluate a checkpoint from an external Policy project with a registered Motion +Profile: + +```bash +embodichain eval-motion-policy \ + --profile my-joint-policy \ + --checkpoint /path/to/policy.pt \ + --config /path/to/train.yaml \ + --resource-root /path/to/task-resources \ + --viewer +``` + +Viewer runs continuously when `--episodes`, `--control-steps`, and `--duration` +are omitted. An EmbodiChain task Headless run completes one episode. An external +Motion Profile Headless run defaults to 100 control steps per episode. + +### Arguments + +| Argument | Default | Description | +|---|---|---| +| ``RUN`` | *(optional)* | Training run containing ``run-manifest.json`` | +| ``--profile`` | manifest or EmbodiChain task | Registered external Motion Profile ID | +| ``--checkpoint`` | ``best`` with RUN | ``best``, ``latest``, a path within RUN, or an explicit `.pt` path | +| ``--config`` | *(optional)* | Training config for an explicit checkpoint | +| ``--gym-config`` | *(optional)* | Task config for an explicit checkpoint | +| ``--resource-root`` | *(optional)* | Task assets and local reference data passed to the Profile | +| ``--episodes`` | EmbodiChain task Headless: ``1`` | Completed EmbodiChain task episodes, or independent Motion Profile runs | +| ``--control-steps`` | Viewer continuous | Exact Policy action count | +| ``--duration`` | *(optional)* | Seconds converted upward to integer control steps | +| ``--command`` | Profile | External Motion Profile command values | +| ``--device`` | Training config or ``cpu`` | PyTorch inference device | +| ``--sim-device`` | Inference device or ``cpu`` | Simulation device | +| ``--physics-backend`` | Policy Spec | External Motion Profile backend override | +| ``--renderer`` | ``hybrid`` | ``raster``, ``hybrid``, ``fastrt``, or ``offlinert`` | +| ``--gpu-id`` | ``0`` | GPU index | +| ``--scene-config`` | Profile: ``standard`` | External Motion Profile scene style or YAML path | +| ``--termination-behavior`` | EmbodiChain task: ``auto_reset`` | EmbodiChain task: ``pause`` or ``auto_reset``; Profile also supports ``continue`` | +| ``--viewer`` | Headless | Open the DexSim Viewer | +| ``--cache-dir`` | DexSim cache | External Motion Profile resource cache | +| ``--offline`` | ``False`` | Resolve external Motion Profile resources from cache | +| ``--output`` | RUN or checkpoint evaluations | Evaluation output parent directory | + +See {doc}`motion_policy_evaluation` for EmbodiChain training runs, external +Motion Profiles, run manifests, reports, and complete examples. + +--- + ## Annotate Grasp Launch the browser-based grasp-region annotation tool. diff --git a/docs/source/guides/index.rst b/docs/source/guides/index.rst index 9044a589a..3c36b2137 100644 --- a/docs/source/guides/index.rst +++ b/docs/source/guides/index.rst @@ -13,4 +13,5 @@ Practical guides for common tasks in EmbodiChain. add_robot preview_asset run_env + motion_policy_evaluation cli diff --git a/docs/source/guides/motion_policy_evaluation.md b/docs/source/guides/motion_policy_evaluation.md new file mode 100644 index 000000000..70ab271cf --- /dev/null +++ b/docs/source/guides/motion_policy_evaluation.md @@ -0,0 +1,359 @@ +# Visual Motion Policy Evaluation + +An EmbodiChain simulator RL task can open its trained `.pt` checkpoint in the +original task Environment and DexSim Viewer. Evaluation restores the Policy and +Environment from the training configuration, loads the checkpoint, and reuses +the task's observation, action processing, reset logic, objects, and termination +conditions. Lightweight learning tasks such as PointMass use the same command +for Headless evaluation. + +## Training runs + +`train-rl` writes a `[new] run-manifest.json` file to each output directory. The +manifest indexes the saved configurations and best/latest checkpoints: + +```text +outputs/_/ +├── checkpoints/ +│ └── policy_*.pt +├── configs/ +│ ├── train.yaml +│ └── gym.yaml +└── run-manifest.json +``` + +A simulator RL task manifest has the following structure: + +```json +{ + "schema_version": 1, + "motion_profile": null, + "configs": { + "train": "configs/train.yaml", + "gym": "configs/gym.yaml" + }, + "checkpoints": { + "best": "checkpoints/cart_pole_grpo_best.pt", + "latest": "checkpoints/cart_pole_grpo_step_.pt" + } +} +``` + +| Field | Description | +|---|---| +| `schema_version` | Manifest schema version; currently `1` | +| `motion_profile` | `null` for an EmbodiChain training run; an external Profile ID when configured | +| `configs.train` | Snapshot of the training configuration | +| `configs.gym` | Snapshot of the Gym configuration for a simulator task | +| `checkpoints.best` | Relative path to the best checkpoint; `null` when no best checkpoint was produced | +| `checkpoints.latest` | Relative path to the checkpoint saved at the end of training | + +Every path is relative to the run directory. `--checkpoint best` selects the +best checkpoint and falls back to `latest` when the manifest value is `null`. +`--checkpoint latest` selects the latest checkpoint directly. A relative path +to another checkpoint in the run is also accepted. The resolved checkpoint is +recorded in `evaluation.json`. + +Open the Viewer for a training run: + +```bash +embodichain eval-motion-policy outputs/_ \ + --checkpoint best \ + --device cuda:0 \ + --sim-device gpu \ + --viewer +``` + +The Viewer runs until its window closes when neither `--control-steps` nor +`--duration` is set. A Headless run uses an exact integer number of control +steps: + +```bash +embodichain eval-motion-policy outputs/_ \ + --checkpoint latest \ + --device cuda:0 \ + --sim-device gpu \ + --control-steps 500 +``` + +## CartPole example + +Train a Policy with the repository's CartPole GRPO configuration: + +```bash +embodichain train-rl \ + --config embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml +``` + +Pass the new run directory to the evaluation command: + +```bash +embodichain eval-motion-policy \ + outputs/cart_pole_grpo_ \ + --checkpoint best \ + --viewer +``` + +The CLI finds the `.pt`, training configuration, and Gym configuration through +`run-manifest.json`. It recreates the CartPole Environment with `num_envs=1`, +rebuilds the `actor_only` Policy, loads its weights, and opens the Viewer. The +report is written under `/evaluations/`. + +## Existing checkpoints + +For a training output created before run manifests were introduced, pass the +checkpoint and its training configuration explicitly: + +```bash +embodichain eval-motion-policy \ + --checkpoint /path/to/policy.pt \ + --config /path/to/train.yaml \ + --gym-config /path/to/gym.yaml \ + --device cuda:0 \ + --sim-device gpu \ + --viewer +``` + +When `trainer.gym_config` already points to the task configuration, +`--gym-config` can be omitted. + +## Execution pipeline + +```mermaid +flowchart LR + Config["Training configuration"] --> Runtime["Shared RL Runtime"] + Checkpoint[".pt checkpoint"] --> Policy["Policy + weights"] + Runtime --> Env["Original EmbodiChain Environment"] + Runtime --> Policy + Env -->|observation| Adapter["EmbodiChain Policy Adapter"] + Policy --> Adapter + Adapter -->|raw Policy action| TaskEnv["EmbodiChain Environment Bridge"] + TaskEnv -->|original action manager| Env + Env -->|env.step| TaskEnv + Evaluator["DexSim MotionPolicyEvaluator"] --> Adapter + Evaluator --> TaskEnv + TaskEnv --> Report["evaluation.json"] +``` + +Each control cycle follows this sequence: + +```mermaid +sequenceDiagram + autonumber + participant Eval as MotionPolicyEvaluator + participant Adapter as EmbodiChainTaskPolicyAdapter + participant TaskEnv as EmbodiChainTaskEnvironment + participant Policy as Policy + participant Env as Original Environment + + Eval->>TaskEnv: reset() + TaskEnv->>Env: reset() + Env-->>TaskEnv: observation and task state + TaskEnv-->>Eval: EvaluationFrame + Eval->>Adapter: reset(frame) + loop Integer control steps + Eval->>Adapter: infer(frame) + Adapter->>Policy: prepare observation and run deterministic inference + Policy-->>Adapter: action tensor + Adapter-->>Eval: PolicyOutput + Eval->>TaskEnv: step(action) + TaskEnv->>Env: apply ActionManager and call env.step() + Env-->>TaskEnv: observation, reward, terminated, truncated, info + TaskEnv-->>Eval: EnvironmentStep + end +``` + +| Component | Responsibility | +|---|---| +| `embodichain.learning.rl.runtime` `[new]` | Builds the Environment and Policy shared by training and evaluation | +| `embodichain.learning.rl.evaluation` `[updated]` | Provides shared observation flattening, deterministic inference, and action conversion | +| `EmbodiChainTaskPolicyAdapter` `[new]` | Reads the original observation from the frame, calls the shared inference path, and returns the raw action | +| `EmbodiChainTaskEnvironment` `[new]` | Calls the original Environment reset, action manager, and step methods, then exposes task results | +| `MotionPolicyEvaluator` | Schedules the Adapter and Environment, counts control steps, applies termination behavior, and updates the Viewer | + +This path loads `.pt` checkpoints through their training-time PyTorch Policy +definition. A standalone DexSim Policy project can use ONNX through its own +Adapter. + +## Viewer and Headless modes + +A simulator RL training run uses `trainer.gym_config` to recreate its original +Environment with `num_envs=1`. An external Policy uses a Motion Profile to +provide the robot assets, Adapter, and Environment. The following Viewer paths +have been validated: + +| Input | Viewer scene | Validated tasks | +|---|---|---| +| EmbodiChain training run | Robot, task objects, and scene restored from the Gym configuration | CartPole GRPO, PushCube PPO | +| External Policy | Robot and scene created by the Motion Profile | ANYmal-C velocity TorchScript `.pt` | + +A lightweight learning task uses `trainer.learning_env` to recreate its tensor +environment for Headless evaluation. PointMass stores its positions, +velocities, targets, and obstacles in PyTorch tensors, so its results are +reported through terminal metrics and `evaluation.json`. + +### Policy loading + +Evaluation rebuilds the Policy from its `policy` configuration, loads the Policy +state dictionary from the checkpoint, and runs deterministic inference. PPO, +GRPO, and APG select the training update algorithm; evaluation follows the +configured Policy type and weights. + +| Training algorithm | Policy type | Viewer validation | Headless validation | +|---|---|---|---| +| PPO | `actor_critic` | PushCube | PushCube, PointMass | +| GRPO | `actor_only` | CartPole | CartPole | +| APG | `actor_only` | — | PointMass | + +A custom EmbodiChain Policy registers its model builder and keeps the matching +network configuration in the training run. An external checkpoint uses a Motion +Profile to describe its model, robot assets, observation, and action processing. +A robot-state Policy can use the Motion Policy Kit default Environment. A Policy +with task objects supplies a complete task Environment. For example, +`UR5CubeResetEnvironment` creates a UR5, table, cube, and target, while +`AllegroReorientCubeEnvironment` creates an Allegro Hand, an object cube, and a +target pose. + +## Termination and duration + +| Option | Behavior | +|---|---| +| `--episodes N` | Completes `N` task episodes | +| `--control-steps N` | Executes exactly `N` Policy actions | +| `--duration SECONDS` | Converts the duration to an integer number of Policy control steps before execution | +| `--termination-behavior auto_reset` | Restores the original Environment after termination; used by the Viewer by default | +| `--termination-behavior pause` | Keeps the terminal task state visible | + +`--control-steps` and `--duration` are mutually exclusive. Simulation time is +calculated from integer step counters using `physics_dt` and +`sim_steps_per_control`. + +## Rendering + +The default renderer is `hybrid`: + +```bash +embodichain eval-motion-policy --viewer --renderer hybrid +``` + +For an EmbodiChain training task, the original Environment provides its robot, +task objects, and scene parameters. An external Motion Profile can use +`--scene-config standard`, `classic`, or a custom YAML file to configure the +ground, lighting, and post-processing. + +## Output + +Each run writes an `evaluation.json` report containing: + +- run, checkpoint, and configuration paths; +- inference and simulation devices, renderer, versions, and commits; +- termination reason, integer simulation steps, and control steps; +- episode reward, episode length, success, and task metrics. + +A training run writes reports under `/evaluations/`. An explicit +checkpoint writes them under the `evaluations/` directory next to that +checkpoint. `--output` selects another output directory. + +## External Policy example + +`examples/learning/motion_policy_evaluation/` contains a runnable ANYmal-C +velocity example. Its public TorchScript `.pt` accepts `vx`, `vy`, and `yaw`, +uses a 48-dimensional observation, and outputs 12 joint actions. The example +includes resource preparation, a Motion Profile, a complete Adapter, tests, and +run commands. + +```text +examples/learning/motion_policy_evaluation/ +├── README.md +├── prepare_resources.py +├── eval_policy.py +└── anymal_c/ + ├── __init__.py + └── profile.py +``` + +Prepare the upstream model and robot resources: + +```bash +python examples/learning/motion_policy_evaluation/prepare_resources.py +``` + +The script fetches `mjw_anymal.pt`, its Policy configuration and license, and +the ANYmal-C URDF and meshes from newton-assets commit +`261cd1f429619d8ef4f546bd788ab9dea906b5e1`. It verifies the expected resource +digests and prints each download and Git operation. Re-running the command +continues an existing checkout. The resulting cache has this structure: + +```text +~/.cache/embodichain/examples/anymal_c_velocity/ +└── upstream/ + └── anybotics_anymal_c/ + ├── rl_policies/ + │ ├── mjw_anymal.pt + │ └── anymal.yaml + ├── urdf/anymal.urdf + └── meshes/... +``` + +Open the Viewer: + +```bash +python examples/learning/motion_policy_evaluation/eval_policy.py \ + --viewer \ + --renderer hybrid +``` + +`eval_policy.py` imports the adjacent `anymal_c/profile.py`, registers its +Profile in the current process, and fills the checkpoint and robot resource +paths from the default cache. Run it directly from the repository root. Options +such as `--device`, `--sim-device`, `--control-steps`, and `--scene-config` are +forwarded to EmbodiChain. + +Use W/S to adjust `vx`, A/D to adjust `vy`, Q/E to adjust `yaw`, M to zero all +three commands, and R to reset the task. + +The external Policy data path is: + +```mermaid +flowchart LR + CLI[eval-motion-policy] --> Profile[build_profile] + Profile --> Spec[Policy Spec
robot, control parameters, frequency] + Spec --> Setup[Adapter.setup
load TorchScript and JointMap] + Setup --> State[RobotState] + Command[vx + vy + yaw] --> Obs[48-dimensional observation] + State --> Obs + Obs --> Actor[TorchScript actor] + Actor --> Action[scale + default position] + Action --> Sim[advance the Environment] +``` + +The Adapter follows the upstream Policy definition when concatenating body +linear velocity, body angular velocity, projected gravity, the three-dimensional +command, joint position, joint velocity, and previous action. The model contains +its observation normalizer. The actor output is converted to joint position +targets with `default_position + 0.5 * action`. Simulation runs at 200 Hz, and +the Policy runs every four simulation steps at 50 Hz. + +To integrate another external Policy, copy this example and replace: + +1. the model and robot resource sources, fixed revisions, paths, and digests; +2. the initial pose, joint control parameters, and control frequency in `build_profile()`; +3. model loading and joint order in `Adapter.setup()`; +4. observation construction, normalization, network forward pass, and action processing in `Adapter.infer()`; +5. `PROFILE_ID` and the Profile name used by `eval_policy.py`. + +Task data can be resolved to local paths through Policy Spec `resources`, or an +Adapter can call an existing project reader from `__init__()` or `setup()`. +Python code can provide changing runtime data through `EvaluationFrame.inputs`. + +The DexSim Motion Policy Kit documentation describes the default Environment, +task Environment interface, Adapter API, task resources, and Python API. + +## Command summary + +| Scenario | Command | +|---|---| +| New training run | `embodichain eval-motion-policy --checkpoint best --viewer` | +| Existing `.pt` | `embodichain eval-motion-policy --checkpoint --config --viewer` | +| Headless smoke test | `embodichain eval-motion-policy --control-steps 20` | +| External ANYmal-C `.pt` | `python examples/learning/motion_policy_evaluation/eval_policy.py --viewer` | diff --git a/docs/source/tutorial/rl.rst b/docs/source/tutorial/rl.rst index 54ce4590f..f91216172 100644 --- a/docs/source/tutorial/rl.rst +++ b/docs/source/tutorial/rl.rst @@ -82,7 +82,7 @@ APG and PPO. Launch either config with the same CLI: embodichain train-rl --config embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml Configuration Sections ---------------------- +---------------------- Runtime Settings ^^^^^^^^^^^^^^^^ @@ -134,7 +134,7 @@ Example: } Policy Configuration -^^^^^^^^^^^^^^^^^^^ +^^^^^^^^^^^^^^^^^^^^ The ``policy`` section defines the neural network policy: @@ -264,9 +264,23 @@ All outputs are written to ``./outputs/_/``: - **logs/**: TensorBoard logs - **checkpoints/**: Model checkpoints +- **configs/**: Training config and referenced gym config snapshots +- **run-manifest.json**: training configs and best/latest checkpoint index used by ``eval-motion-policy`` + +A simulator training run can be opened directly in its original task Environment: + +.. code-block:: bash + + embodichain eval-motion-policy outputs/_ --viewer + +An external Policy project can record its Motion Profile ID in +``trainer.motion_profile``. + +See :doc:`../guides/motion_policy_evaluation` for EmbodiChain ``.pt`` training +runs and the external Motion Profile example. Training Process -~~~~~~~~~~~~~~~ +~~~~~~~~~~~~~~~~ The training process follows this sequence: @@ -326,7 +340,7 @@ Available Algorithms - **GRPO**: Group Relative Policy Optimization (no Critic, step-wise returns, masked group normalization). Use ``actor_only`` policy. Set ``kl_coef=0`` for from-scratch training (CartPole, dense reward); ``kl_coef=0.02`` for VLA/LLM fine-tuning. Adding a New Algorithm ---------------------- +---------------------- To add a new algorithm: @@ -464,3 +478,4 @@ See Also - :doc:`basic_env` — Creating basic Gymnasium environments - :doc:`modular_env` — Advanced modular environments with managers - :doc:`/resources/task/index` — List of available RL task environments +- :doc:`/guides/motion_policy_evaluation` — Visual evaluation of EmbodiChain ``.pt`` checkpoints and external Motion Profiles diff --git a/embodichain/__main__.py b/embodichain/__main__.py index 3b897a3fa..584dd05e8 100644 --- a/embodichain/__main__.py +++ b/embodichain/__main__.py @@ -81,6 +81,11 @@ class Command: target="embodichain.learning.rl.train:cli", help="Train an RL agent from a JSON or YAML config.", ), + Command( + name="eval-motion-policy", + target="embodichain.learning.rl.motion_policy_evaluation.cli:cli", + help="Visualize a policy through DexSim Motion Policy Kit.", + ), Command( name="annotate-grasp", target="embodichain.toolkits.graspkit.scripts.annotate_grasp:cli", diff --git a/embodichain/learning/rl/evaluation.py b/embodichain/learning/rl/evaluation.py index 974a2ab7c..3a54715b2 100644 --- a/embodichain/learning/rl/evaluation.py +++ b/embodichain/learning/rl/evaluation.py @@ -30,15 +30,43 @@ flatten_dict_observation, ) -__all__ = ["evaluate_episodes"] +__all__ = [ + "convert_policy_action_for_env", + "evaluate_episodes", + "infer_policy_action", + "prepare_policy_observation", +] -def _flat_observation(observation: Any, device: torch.device) -> torch.Tensor: +def prepare_policy_observation( + observation: Any, + device: torch.device | str, +) -> torch.Tensor: + """Flatten one Environment observation in the training input order.""" + device = torch.device(device) tensor_dict = dict_to_tensordict(observation, device) return flatten_dict_observation(tensor_dict) -def _action_for_env(env: Any, action: torch.Tensor) -> Any: +def infer_policy_action( + policy: torch.nn.Module, + observation: Any, + *, + device: torch.device | str, + num_envs: int, +) -> torch.Tensor: + """Run the same deterministic Policy call used by RL evaluation.""" + device = torch.device(device) + policy_input = TensorDict( + {"obs": prepare_policy_observation(observation, device)}, + batch_size=[num_envs], + device=device, + ) + return policy.get_action(policy_input, deterministic=True)["action"] + + +def convert_policy_action_for_env(env: Any, action: torch.Tensor) -> Any: + """Convert a flat Policy action to the task Environment input layout.""" action_manager = getattr(env, "action_manager", None) if action_manager is None and hasattr(env, "get_wrapper_attr"): try: @@ -106,15 +134,14 @@ def evaluate_episodes( try: observation, _ = env.reset(seed=seed) while len(returns) < num_episodes: - flat_observation = _flat_observation(observation, device) - policy_input = TensorDict( - {"obs": flat_observation}, - batch_size=[num_envs], + action = infer_policy_action( + policy, + observation, device=device, + num_envs=num_envs, ) - policy_output = policy.get_action(policy_input, deterministic=True) observation, reward, terminated, truncated, info = env.step( - _action_for_env(env, policy_output["action"]) + convert_policy_action_for_env(env, action) ) reward = torch.as_tensor(reward, device=device).reshape(num_envs) done = ( diff --git a/embodichain/learning/rl/motion_policy_evaluation/__init__.py b/embodichain/learning/rl/motion_policy_evaluation/__init__.py new file mode 100644 index 000000000..d68087a27 --- /dev/null +++ b/embodichain/learning/rl/motion_policy_evaluation/__init__.py @@ -0,0 +1,50 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Visual motion-policy evaluation through DexSim Motion Policy Kit.""" + +from __future__ import annotations + +from .bridge import ( + MotionEvaluationResult, + create_motion_profile_evaluator, + evaluate_motion_profile, +) +from .checkpoint import load_policy_state_dict +from .manifest import RunManifest, write_run_manifest +from .profile import ( + MotionProfile, + MotionProfileRequest, + build_motion_profile, + get_motion_profile_names, + register_motion_profile, +) +from .report import write_evaluation_report + +__all__ = [ + "MotionEvaluationResult", + "MotionProfile", + "MotionProfileRequest", + "RunManifest", + "build_motion_profile", + "create_motion_profile_evaluator", + "evaluate_motion_profile", + "get_motion_profile_names", + "load_policy_state_dict", + "register_motion_profile", + "write_run_manifest", + "write_evaluation_report", +] diff --git a/embodichain/learning/rl/motion_policy_evaluation/bridge.py b/embodichain/learning/rl/motion_policy_evaluation/bridge.py new file mode 100644 index 000000000..1299212a4 --- /dev/null +++ b/embodichain/learning/rl/motion_policy_evaluation/bridge.py @@ -0,0 +1,251 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Run an EmbodiChain Motion Profile through DexSim Motion Policy Kit.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Any + +from dexsim.kit.motion_policy import ( + MotionPolicyEvaluator, + PolicyAdapter, + PolicySpec, + ResolvedPolicy, + ResourceResolver, + RunOptions, + create_motion_policy_evaluator, + load_scene_config, + parse_policy_spec, + policy_spec_to_dict, + resolve_policy_spec, + run_motion_policy, + scene_config_to_dict, +) +from dexsim.kit.motion_policy.environment import PolicyEnvironment +from dexsim.kit.motion_policy.evaluator import InputProvider + +from .profile import MotionProfile + +__all__ = [ + "MotionEvaluationResult", + "create_motion_profile_evaluator", + "evaluate_motion_profile", +] + + +@dataclass(frozen=True) +class MotionEvaluationResult: + """Normalized inputs and per-episode motion evaluation results.""" + + profile: MotionProfile + policy_spec: Mapping[str, Any] + scene_config: Mapping[str, Any] + episodes: tuple[Mapping[str, Any], ...] + summary: Mapping[str, Any] + viewer: bool + + def __post_init__(self) -> None: + object.__setattr__( + self, "policy_spec", MappingProxyType(dict(self.policy_spec)) + ) + object.__setattr__( + self, + "scene_config", + MappingProxyType(dict(self.scene_config)), + ) + object.__setattr__( + self, + "episodes", + tuple(MappingProxyType(dict(value)) for value in self.episodes), + ) + object.__setattr__(self, "summary", MappingProxyType(dict(self.summary))) + + +def evaluate_motion_profile( + profile: MotionProfile, + *, + episodes: int = 1, + viewer: bool = False, + control_steps: int | None = None, + duration: float | None = None, + command: tuple[float, ...] | None = None, + scene_config: str | Path = "standard", + physics_backend: str | None = None, + simulation_device: str = "cpu", + renderer: str = "hybrid", + gpu_id: int = 0, + termination_behavior: str | None = None, + cache_dir: str | Path | None = None, + offline: bool = False, + input_provider: InputProvider | None = None, + environment: PolicyEnvironment | None = None, +) -> MotionEvaluationResult: + """Resolve one Motion Profile and run its visual evaluation. + + Args: + profile: Provider-built profile containing the DexSim Policy Spec. + episodes: Number of independent runs. + viewer: Open the DexSim Viewer. + control_steps: Exact number of applied policy commands per run. + duration: Convenience duration converted by DexSim to policy steps. + command: Optional task command override. + scene_config: Built-in scene style or custom YAML path. + physics_backend: Optional DexSim physics backend override. + simulation_device: ``cpu`` or ``gpu``. + renderer: DexSim renderer. + gpu_id: Selected GPU index. + termination_behavior: Policy termination handling override. + cache_dir: Motion Policy Kit resource cache. + offline: Use resources already available in the cache. + input_provider: Build per-frame task inputs for the Adapter. + environment: Prebuilt task environment owned by the Evaluator. + + Returns: + Normalized inputs, episode results, and aggregate metrics. + """ + if episodes <= 0: + raise ValueError("episodes must be positive") + if viewer and episodes != 1: + raise ValueError("Viewer evaluation supports one episode") + if environment is not None and episodes != 1: + raise ValueError("A prebuilt environment supports one episode") + parsed, resolved = _resolve_profile(profile, cache_dir, offline) + resolved_scene = load_scene_config(scene_config) + options = RunOptions( + physics_backend=physics_backend, + simulation_device=simulation_device, + renderer=renderer, + gpu_id=gpu_id, + headless=not viewer, + control_steps=control_steps, + duration=duration, + command=command, + termination_behavior=termination_behavior, + scene_config=resolved_scene, + ) + results = tuple( + _episode( + index, + run_motion_policy( + resolved, + options, + environment=environment, + input_provider=input_provider, + ), + ) + for index in range(episodes) + ) + return MotionEvaluationResult( + profile=profile, + policy_spec=policy_spec_to_dict(parsed), + scene_config=scene_config_to_dict(resolved_scene), + episodes=results, + summary=_summary(results), + viewer=viewer, + ) + + +def create_motion_profile_evaluator( + profile: MotionProfile, + options: RunOptions | None = None, + *, + cache_dir: str | Path | None = None, + offline: bool = False, + adapter: PolicyAdapter | None = None, + environment: PolicyEnvironment | None = None, +) -> MotionPolicyEvaluator: + """Create a DexSim Evaluator for an EmbodiChain Motion Profile. + + Args: + profile: Provider-built profile containing the DexSim Policy Spec. + options: DexSim evaluation options. + cache_dir: Motion Policy Kit resource cache. + offline: Use resources already available in the cache. + adapter: Prebuilt policy adapter owned by the returned Evaluator. + environment: Prebuilt task environment owned by the returned Evaluator. + + Returns: + Configured evaluator ready for ``reset()``, ``step()``, or ``run()``. + """ + _parsed, resolved = _resolve_profile(profile, cache_dir, offline) + return create_motion_policy_evaluator( + resolved, + options, + adapter=adapter, + environment=environment, + ) + + +def _resolve_profile( + profile: MotionProfile, + cache_dir: str | Path | None, + offline: bool, +) -> tuple[PolicySpec, ResolvedPolicy]: + parsed = parse_policy_spec(profile.policy_spec) + resolved = resolve_policy_spec( + parsed, + ResourceResolver( + None if cache_dir is None else Path(cache_dir), + offline=offline, + ), + ) + return parsed, resolved + + +def _episode(index: int, result: Any) -> dict[str, Any]: + return { + "index": index, + "reason": str(result.reason), + "simulation_time": float(result.simulation_time), + "simulation_steps": int(result.simulation_steps), + "control_steps": int(result.control_steps), + "physics_backend": str(result.physics_backend), + "requested_duration": ( + None + if result.requested_duration is None + else float(result.requested_duration) + ), + "effective_duration": float(result.effective_duration), + "metrics": {name: float(value) for name, value in result.metrics.items()}, + } + + +def _summary(episodes: tuple[Mapping[str, Any], ...]) -> dict[str, Any]: + count = len(episodes) + metric_names = set.intersection(*(set(episode["metrics"]) for episode in episodes)) + metrics = { + name: sum(episode["metrics"][name] for episode in episodes) / count + for name in sorted(metric_names) + } + result: dict[str, Any] = { + "episodes": count, + "avg_simulation_time": sum(episode["simulation_time"] for episode in episodes) + / count, + "avg_control_steps": sum(episode["control_steps"] for episode in episodes) + / count, + "avg_effective_duration": sum( + episode["effective_duration"] for episode in episodes + ) + / count, + } + if metrics: + result["metrics"] = metrics + return result diff --git a/embodichain/learning/rl/motion_policy_evaluation/checkpoint.py b/embodichain/learning/rl/motion_policy_evaluation/checkpoint.py new file mode 100644 index 000000000..817eddb4c --- /dev/null +++ b/embodichain/learning/rl/motion_policy_evaluation/checkpoint.py @@ -0,0 +1,60 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Load the policy weights stored in an EmbodiChain checkpoint.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import torch + +__all__ = ["load_policy_state_dict"] + + +def load_policy_state_dict( + checkpoint: str | Path, + *, + map_location: str | torch.device = "cpu", +) -> Mapping[str, Any]: + """Load the ``policy`` state mapping from an EmbodiChain ``.pt`` file. + + Args: + checkpoint: EmbodiChain training checkpoint. + map_location: Device passed to :func:`torch.load`. + + Returns: + Policy state mapping ready for ``module.load_state_dict()``. + + Raises: + FileNotFoundError: If the checkpoint does not exist. + TypeError: If the checkpoint or policy payload is not a mapping. + ValueError: If the checkpoint has no policy weights. + """ + path = Path(checkpoint).expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(f"Policy checkpoint does not exist: {path}") + payload = torch.load(path, map_location=map_location, weights_only=True) + if not isinstance(payload, Mapping): + raise TypeError("Policy checkpoint root must be a mapping") + state = payload.get("policy") + if not isinstance(state, Mapping): + raise TypeError("Policy checkpoint field 'policy' must be a mapping") + if not state: + raise ValueError("Policy checkpoint field 'policy' is empty") + return state diff --git a/embodichain/learning/rl/motion_policy_evaluation/cli.py b/embodichain/learning/rl/motion_policy_evaluation/cli.py new file mode 100644 index 000000000..0ef445f9f --- /dev/null +++ b/embodichain/learning/rl/motion_policy_evaluation/cli.py @@ -0,0 +1,398 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Command line for visual evaluation of EmbodiChain policy checkpoints.""" + +from __future__ import annotations + +import argparse +import sys +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch + +from embodichain import __version__ +from embodichain.lab.gym.utils.registration import ( + discover_task_packages, + execute_init_hooks, +) +from embodichain.learning.rl.runtime import ( + build_gym_policy_runtime, + build_learning_policy_runtime, + resolve_torch_device, + seed_policy_runtime, +) +from embodichain.utils.utility import load_config + +from .bridge import MotionEvaluationResult, evaluate_motion_profile +from .checkpoint import load_policy_state_dict +from .manifest import RunManifest +from .native_task import NativeTaskEvaluationResult, evaluate_native_task +from .profile import MotionProfileRequest, build_motion_profile +from .report import write_evaluation_report + +__all__ = ["cli", "parse_args", "run"] + + +@dataclass(frozen=True) +class MotionInput: + """Resolved checkpoint, configs, and Profile selection.""" + + checkpoint: Path + profile: str | None + configs: Mapping[str, Path] + run: Path | None + requested_checkpoint: str + selected_checkpoint: str + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + """Parse ``embodichain eval-motion-policy`` arguments.""" + parser = argparse.ArgumentParser( + prog="embodichain eval-motion-policy", + description="Open a DexSim Viewer for an EmbodiChain policy checkpoint.", + ) + parser.add_argument("run", nargs="?", help="EmbodiChain training run directory.") + parser.add_argument("--profile", help="Registered Motion Profile name.") + parser.add_argument( + "--checkpoint", + help="Checkpoint path, or best/latest when RUN is supplied.", + ) + parser.add_argument("--config", help="Training config for an explicit checkpoint.") + parser.add_argument("--gym-config", help="Task config for an explicit checkpoint.") + parser.add_argument("--resource-root", help="Profile-specific local resource root.") + parser.add_argument("--episodes", type=int) + count = parser.add_mutually_exclusive_group() + count.add_argument("--control-steps", type=int) + count.add_argument("--duration", type=float) + parser.add_argument("--command", nargs="+", type=float) + parser.add_argument("--device", help="PyTorch inference device.") + parser.add_argument("--sim-device", choices=("cpu", "gpu")) + parser.add_argument("--seed", type=int) + parser.add_argument("--physics-backend", choices=("default",)) + parser.add_argument( + "--renderer", + choices=("raster", "hybrid", "fastrt", "offlinert"), + default="hybrid", + ) + parser.add_argument("--gpu-id", type=int, default=0) + parser.add_argument("--scene-config") + parser.add_argument( + "--termination-behavior", + choices=("pause", "continue", "auto_reset"), + ) + parser.add_argument("--viewer", action="store_true") + parser.add_argument("--cache-dir") + parser.add_argument("--offline", action="store_true") + parser.add_argument("--output", help="Evaluation output parent directory.") + return parser.parse_args(argv) + + +def run(args: argparse.Namespace) -> Path: + """Resolve inputs, run the evaluation, and write ``evaluation.json``.""" + resolved = _resolve_input(args) + discover_task_packages() + execute_init_hooks() + if resolved.profile is None: + return _run_native_task(args, resolved) + + device = torch.device(args.device or "cpu") + profile = build_motion_profile( + resolved.profile, + MotionProfileRequest( + checkpoint=resolved.checkpoint, + device=device, + configs=resolved.configs, + resource_root=( + None if args.resource_root is None else Path(args.resource_root) + ), + renderer=args.renderer, + ), + ) + for warning in profile.warnings: + print(f"Warning: {warning}", file=sys.stderr) + result = evaluate_motion_profile( + profile, + episodes=args.episodes or 1, + viewer=args.viewer, + control_steps=args.control_steps, + duration=args.duration, + command=None if args.command is None else tuple(args.command), + scene_config=args.scene_config or "standard", + physics_backend=args.physics_backend, + simulation_device=args.sim_device or "cpu", + renderer=args.renderer, + gpu_id=args.gpu_id, + termination_behavior=args.termination_behavior, + cache_dir=args.cache_dir, + offline=args.offline, + ) + return write_evaluation_report( + _output_parent(args.output, resolved), + _profile_report(result, resolved, device), + ) + + +def cli(argv: Sequence[str] | None = None) -> None: + """Run motion-policy evaluation from the unified EmbodiChain CLI.""" + try: + report = run(parse_args(argv)) + except ( + FileNotFoundError, + ImportError, + KeyError, + RuntimeError, + TypeError, + ValueError, + ) as error: + raise SystemExit(f"eval-motion-policy: {error}") from error + print(f"Evaluation report: {report}") + + +def _resolve_input(args: argparse.Namespace) -> MotionInput: + if args.run is not None: + manifest = RunManifest.load(args.run) + requested = args.checkpoint or "best" + if requested in {"best", "latest"}: + selected, checkpoint = manifest.select_checkpoint(requested) + else: + selected = "explicit" + candidate = Path(requested).expanduser() + checkpoint = ( + candidate.resolve() + if candidate.is_absolute() + else (manifest.root / candidate).resolve() + ) + profile = args.profile or manifest.motion_profile + return MotionInput( + checkpoint=checkpoint, + profile=profile, + configs=manifest.configs, + run=manifest.root, + requested_checkpoint=requested, + selected_checkpoint=selected, + ) + if args.checkpoint is None: + raise ValueError("--checkpoint is required without RUN") + checkpoint = Path(args.checkpoint).expanduser().resolve() + configs = {} + if args.config is not None: + configs["train"] = Path(args.config).expanduser().resolve() + if args.gym_config is not None: + configs["gym"] = Path(args.gym_config).expanduser().resolve() + if args.profile is None and "train" not in configs: + raise ValueError("--config is required for a native EmbodiChain checkpoint") + return MotionInput( + checkpoint=checkpoint, + profile=args.profile, + configs=configs, + run=None, + requested_checkpoint=args.checkpoint, + selected_checkpoint="explicit", + ) + + +def _run_native_task(args: argparse.Namespace, resolved: MotionInput) -> Path: + _validate_native_options(args) + train_config = resolved.configs.get("train") + if train_config is None: + raise ValueError("Training config is required for native task evaluation") + config = load_config(train_config) + config["trainer"] = dict(config["trainer"]) + gym_config = resolved.configs.get("gym") + if gym_config is not None: + config["trainer"]["gym_config"] = str(gym_config) + trainer = config["trainer"] + device = resolve_torch_device(args.device or trainer.get("device", "cpu")) + if args.sim_device == "gpu": + simulation_device = resolve_torch_device(f"cuda:{args.gpu_id}") + elif args.sim_device == "cpu": + simulation_device = torch.device("cpu") + else: + simulation_device = device + seed = int( + args.seed + if args.seed is not None + else trainer.get("eval_seed", int(trainer.get("seed", 1)) + 10_000) + ) + seed_policy_runtime(seed, device) + uses_simulator = "learning_env" not in trainer + if uses_simulator: + runtime = build_gym_policy_runtime( + config, + device=device, + simulation_device=simulation_device, + num_envs=1, + headless=not args.viewer, + renderer=args.renderer, + gpu_id=args.gpu_id, + config_dir=train_config.parent, + ) + else: + runtime = build_learning_policy_runtime( + config, + device=device, + num_envs=1, + ) + try: + try: + runtime.policy.load_state_dict( + load_policy_state_dict(resolved.checkpoint, map_location="cpu") + ) + except Exception: + runtime.env.close() + raise + + episodes = args.episodes + if ( + episodes is None + and not args.viewer + and args.control_steps is None + and args.duration is None + ): + episodes = 1 + result = evaluate_native_task( + runtime, + seed=seed, + viewer=args.viewer, + episodes=episodes, + control_steps=args.control_steps, + duration=args.duration, + termination_behavior=args.termination_behavior or "auto_reset", + ) + finally: + if uses_simulator: + from embodichain.lab.sim.sim_manager import SimulationManager + + SimulationManager.flush_cleanup_queue() + return write_evaluation_report( + _output_parent(args.output, resolved), + _native_report( + result, + resolved, + device=device, + simulation_device=simulation_device, + seed=seed, + renderer=args.renderer, + ), + ) + + +def _validate_native_options(args: argparse.Namespace) -> None: + """Reject options that belong to external Motion Profiles.""" + values = { + "--resource-root": args.resource_root, + "--command": args.command, + "--physics-backend": args.physics_backend, + "--scene-config": args.scene_config, + "--cache-dir": args.cache_dir, + "--offline": args.offline, + } + selected = [name for name, value in values.items() if value not in (None, False)] + if selected: + raise ValueError(f"{', '.join(selected)} requires --profile") + + +def _output_parent(configured: str | None, resolved: MotionInput) -> Path: + if configured is not None: + return Path(configured) + if resolved.run is not None: + return resolved.run / "evaluations" + return resolved.checkpoint.parent / "evaluations" + + +def _profile_report( + result: MotionEvaluationResult, + resolved: MotionInput, + device: torch.device, +) -> dict[str, Any]: + import dexsim + + return { + "mode": "viewer" if result.viewer else "headless", + "inputs": { + "run": resolved.run, + "checkpoint": { + "path": resolved.checkpoint, + "requested": resolved.requested_checkpoint, + "selected": resolved.selected_checkpoint, + }, + "profile": { + "id": result.profile.profile_id, + "provider_version": result.profile.provider_version, + "provenance": result.profile.provenance, + "warnings": result.profile.warnings, + }, + "configs": resolved.configs, + "policy_spec": result.policy_spec, + "scene_config": result.scene_config, + "inference_device": str(device), + "embodichain_version": __version__, + "dexsim_version": getattr(dexsim, "__version__", None), + "dexsim_commit": getattr(dexsim, "__commit_id__", None), + }, + "result": { + "episodes": result.episodes, + "summary": result.summary, + }, + } + + +def _native_report( + result: NativeTaskEvaluationResult, + resolved: MotionInput, + *, + device: torch.device, + simulation_device: torch.device, + seed: int, + renderer: str, +) -> dict[str, Any]: + import dexsim + + return { + "mode": "viewer" if result.viewer else "headless", + "inputs": { + "run": resolved.run, + "checkpoint": { + "path": resolved.checkpoint, + "requested": resolved.requested_checkpoint, + "selected": resolved.selected_checkpoint, + }, + "configs": resolved.configs, + "task_id": result.task_id, + "seed": seed, + "inference_device": str(device), + "simulation_device": str(simulation_device), + "renderer": renderer, + "embodichain_version": __version__, + "dexsim_version": getattr(dexsim, "__version__", None), + "dexsim_commit": getattr(dexsim, "__commit_id__", None), + }, + "result": { + "reason": result.reason, + "simulation_time": result.simulation_time, + "simulation_steps": result.simulation_steps, + "control_steps": result.control_steps, + "physics_backend": result.physics_backend, + "requested_duration": result.requested_duration, + "effective_duration": result.effective_duration, + "episodes": result.episodes, + "metrics": result.metrics, + }, + } diff --git a/embodichain/learning/rl/motion_policy_evaluation/manifest.py b/embodichain/learning/rl/motion_policy_evaluation/manifest.py new file mode 100644 index 000000000..19b9c5f7f --- /dev/null +++ b/embodichain/learning/rl/motion_policy_evaluation/manifest.py @@ -0,0 +1,205 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Minimal index connecting a training run to motion-policy evaluation.""" + +from __future__ import annotations + +import json +import shutil +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Any + +__all__ = ["RUN_MANIFEST_NAME", "RunManifest", "write_run_manifest"] + +RUN_MANIFEST_NAME = "run-manifest.json" + + +@dataclass(frozen=True) +class RunManifest: + """Resolved paths from one EmbodiChain training run.""" + + root: Path + motion_profile: str | None + configs: Mapping[str, Path] + checkpoints: Mapping[str, Path | None] + + def __post_init__(self) -> None: + object.__setattr__(self, "root", Path(self.root).resolve()) + object.__setattr__(self, "configs", MappingProxyType(dict(self.configs))) + object.__setattr__( + self, + "checkpoints", + MappingProxyType(dict(self.checkpoints)), + ) + + @classmethod + def load(cls, run: str | Path) -> RunManifest: + """Load ``run-manifest.json`` and resolve its referenced files. + + Args: + run: EmbodiChain training run directory. + + Returns: + Resolved manifest. + """ + root = Path(run).expanduser().resolve() + path = root / RUN_MANIFEST_NAME + if not path.is_file(): + raise FileNotFoundError(f"Run manifest does not exist: {path}") + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, Mapping) or value.get("schema_version") != 1: + raise ValueError(f"Unsupported run manifest: {path}") + configs = _resolve_group(root, value.get("configs"), "configs") + checkpoints = _resolve_group( + root, + value.get("checkpoints"), + "checkpoints", + allow_none=True, + ) + profile = value.get("motion_profile") + if profile is not None and not isinstance(profile, str): + raise TypeError("Run manifest motion_profile must be a string or null") + return cls(root, profile, configs, checkpoints) + + def select_checkpoint(self, requested: str = "best") -> tuple[str, Path]: + """Select ``best`` or ``latest`` and return its resolved path. + + Args: + requested: Checkpoint role. + + Returns: + Selected role and checkpoint path. ``best`` uses ``latest`` when + the training run has no best checkpoint. + """ + if requested not in {"best", "latest"}: + raise ValueError("checkpoint role must be best or latest") + selected = requested + checkpoint = self.checkpoints.get(selected) + if checkpoint is None and requested == "best": + selected = "latest" + checkpoint = self.checkpoints.get(selected) + if checkpoint is None: + raise FileNotFoundError( + f"Run manifest has no {requested} checkpoint: {self.root}" + ) + return selected, checkpoint + + +def write_run_manifest( + run: str | Path, + *, + train_config: str | Path, + latest_checkpoint: str | Path, + best_checkpoint: str | Path | None = None, + gym_config: str | Path | None = None, + motion_profile: str | None = None, +) -> Path: + """Snapshot training configs and write the minimal run manifest. + + Args: + run: Training run directory containing the checkpoints. + train_config: Training config used for the run. + latest_checkpoint: Final saved checkpoint. + best_checkpoint: Best checkpoint when evaluation selected one. + gym_config: Referenced task config when the trainer uses one. + motion_profile: Default Motion Profile for visual evaluation. + + Returns: + Written manifest path. + """ + root = Path(run).expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + config_dir = root / "configs" + config_dir.mkdir(exist_ok=True) + configs = { + "train": _snapshot_config(train_config, config_dir, "train"), + } + if gym_config is not None: + configs["gym"] = _snapshot_config(gym_config, config_dir, "gym") + checkpoints = { + "best": _relative_file(root, best_checkpoint), + "latest": _relative_file(root, latest_checkpoint), + } + value: dict[str, Any] = { + "schema_version": 1, + "motion_profile": motion_profile, + "configs": configs, + "checkpoints": checkpoints, + } + path = root / RUN_MANIFEST_NAME + path.write_text( + json.dumps(value, indent=2, sort_keys=True, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + return path + + +def _snapshot_config(source: str | Path, target: Path, name: str) -> str: + path = Path(source).expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(f"Training config does not exist: {path}") + suffix = path.suffix.lower() if path.suffix else ".yaml" + destination = target / f"{name}{suffix}" + shutil.copyfile(path, destination) + return destination.relative_to(target.parent).as_posix() + + +def _relative_file(root: Path, value: str | Path | None) -> str | None: + if value is None: + return None + path = Path(value).expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(f"Training checkpoint does not exist: {path}") + try: + return path.relative_to(root).as_posix() + except ValueError as error: + raise ValueError(f"Training checkpoint is outside its run: {path}") from error + + +def _resolve_group( + root: Path, + value: object, + field: str, + *, + allow_none: bool = False, +) -> dict[str, Path | None]: + if not isinstance(value, Mapping): + raise TypeError(f"Run manifest {field} must be a mapping") + result: dict[str, Path | None] = {} + for name, reference in value.items(): + if reference is None and allow_none: + result[str(name)] = None + continue + if not isinstance(reference, str) or not reference: + raise TypeError(f"Run manifest {field}.{name} must be a path") + relative = Path(reference) + if relative.is_absolute(): + raise ValueError(f"Run manifest {field}.{name} must be relative") + path = (root / relative).resolve() + try: + path.relative_to(root) + except ValueError as error: + raise ValueError( + f"Run manifest {field}.{name} escapes the run directory" + ) from error + if not path.is_file(): + raise FileNotFoundError(f"Run manifest file does not exist: {path}") + result[str(name)] = path + return result diff --git a/embodichain/learning/rl/motion_policy_evaluation/native_task.py b/embodichain/learning/rl/motion_policy_evaluation/native_task.py new file mode 100644 index 000000000..f976b58c6 --- /dev/null +++ b/embodichain/learning/rl/motion_policy_evaluation/native_task.py @@ -0,0 +1,481 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Connect an EmbodiChain RL task to DexSim Motion Policy Evaluator.""" + +from __future__ import annotations + +import math +import time +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch + +from dexsim.kit.motion_policy import ( + EvaluationFrame, + PolicyContext, + PolicyOutput, + RunOptions, + create_motion_policy_evaluator, +) +from dexsim.kit.motion_policy.types import EnvironmentStep + +from embodichain.learning.rl.evaluation import ( + convert_policy_action_for_env, + infer_policy_action, +) +from embodichain.learning.rl.runtime import PolicyRuntime + +__all__ = [ + "EmbodiChainTaskEnvironment", + "EmbodiChainTaskPolicyAdapter", + "NativeTaskEvaluationResult", + "evaluate_native_task", +] + +_MISSING = object() + + +@dataclass(frozen=True) +class NativeTaskEvaluationResult: + """Result of evaluating one Policy in its original EmbodiChain task.""" + + task_id: str + reason: str + simulation_time: float + simulation_steps: int + control_steps: int + physics_backend: str + effective_duration: float + requested_duration: float | None + episodes: tuple[Mapping[str, float | int | bool | str], ...] + metrics: Mapping[str, float] + viewer: bool + + +class EmbodiChainTaskPolicyAdapter: + """Run an EmbodiChain Policy from the task observation in each frame.""" + + def __init__(self, policy: torch.nn.Module, device: torch.device, num_envs: int): + self.policy = policy + self.device = device + self.num_envs = num_envs + self._previous_training = policy.training + + def setup(self, context: PolicyContext) -> None: + """Select deterministic inference for this evaluation.""" + del context + self.policy.eval() + + def reset(self, frame: EvaluationFrame) -> None: + """Validate that the Environment supplied the next observation.""" + if frame.observation is None: + raise RuntimeError("EmbodiChain task frame has no observation") + + @torch.no_grad() + def infer(self, frame: EvaluationFrame) -> PolicyOutput: + """Run the same observation and deterministic Policy path as RL evaluation.""" + if frame.observation is None: + raise RuntimeError("EmbodiChain task frame has no observation") + action = infer_policy_action( + self.policy, + frame.observation, + device=self.device, + num_envs=self.num_envs, + ) + return PolicyOutput(action=action) + + def metrics(self) -> dict[str, float]: + """Return Policy-side metrics.""" + return {} + + def close(self) -> None: + """Restore the Policy mode used before evaluation.""" + self.policy.train(self._previous_training) + + +class EmbodiChainTaskEnvironment: + """Expose one original EmbodiChain RL Environment to the Evaluator.""" + + def __init__( + self, + env: Any, + *, + seed: int, + viewer: bool, + ) -> None: + if int(env.num_envs) != 1: + raise ValueError("Visual task evaluation requires num_envs=1") + self.env = env + self._base_env = getattr(env, "unwrapped", env) + self._viewer = viewer + if viewer: + world = self._world() + if world is None or not world.is_window_initialized(): + raise ValueError( + "Viewer evaluation requires an EmbodiChain simulator task " + "with an initialized window" + ) + self._seed = seed + self._first_reset = True + self._control_step = 0 + self._frame: EvaluationFrame | None = None + self._episode_return = 0.0 + self._episode_length = 0 + self._episodes: list[dict[str, float | int | bool | str]] = [] + self._closed = False + self._previous_no_auto_reset = getattr( + self._base_env, + "_demo_no_auto_reset", + _MISSING, + ) + self._base_env._demo_no_auto_reset = True + self._policy_context = _policy_context_from_env(self._base_env) + + @property + def policy_context(self) -> PolicyContext: + """Return the timing used by the original task Environment.""" + return self._policy_context + + @property + def physics_backend(self) -> str: + """Return the backend selected by the original task Environment.""" + return "default" + + @property + def viewer_is_open(self) -> bool: + """Return whether the original task Viewer remains open.""" + if not self._viewer: + return False + world = self._world() + return bool(world is not None and world.is_window_initialized()) + + @property + def current_frame(self) -> EvaluationFrame: + """Return the latest observation and task state.""" + if self._frame is None: + raise RuntimeError("Environment has not been reset") + return self._frame + + @property + def episodes(self) -> tuple[Mapping[str, float | int | bool | str], ...]: + """Return completed episode summaries.""" + return tuple(dict(episode) for episode in self._episodes) + + def open_viewer(self, title: str) -> None: + """Apply the evaluation title to the task Viewer.""" + if not self._viewer: + return + world = self._world() + if world is not None and world.is_window_initialized(): + world.get_windows().set_window_title(title) + + def reset(self) -> EvaluationFrame: + """Run the task's original reset and return its observation.""" + kwargs = {"seed": self._seed} if self._first_reset else {} + observation, info = self.env.reset(**kwargs) + self._first_reset = False + self._control_step = 0 + self._episode_return = 0.0 + self._episode_length = 0 + self._frame = self._make_frame(observation, {"info": info}) + return self._frame + + def poll(self) -> str | None: + """Report when the native Viewer is closed or Escape is pressed.""" + if not self._viewer: + return None + world = self._world() + if world is None or not world.is_window_initialized(): + return "viewer closed" + from dexsim.types import InputKey + + native = world.get_windows().native() + if native.key_state(InputKey.SCANCODE_ESCAPE): + return "viewer closed" + return None + + def step(self, action: object) -> EnvironmentStep: + """Apply one raw Policy action through the task's original action path.""" + if not isinstance(action, torch.Tensor): + raise TypeError("EmbodiChain Policy action must be a torch.Tensor") + started = time.perf_counter() + env_action = convert_policy_action_for_env(self.env, action) + observation, reward, terminated, truncated, info = self.env.step(env_action) + reward_value = _single_float(reward, "reward") + terminated_value = _single_bool(terminated, "terminated") + truncated_value = _single_bool(truncated, "truncated") + self._control_step += 1 + self._episode_return += reward_value + self._episode_length += 1 + task_state = { + "reward": reward, + "terminated": terminated, + "truncated": truncated, + "info": info, + } + self._frame = self._make_frame(observation, task_state) + reason = _termination_reason(info, terminated_value, truncated_value) + metrics = _step_metrics(info, reward_value) + if reason is not None: + success = _info_bool(info, "success") + self._episodes.append( + { + "index": len(self._episodes), + "reason": reason, + "reward": self._episode_return, + "length": self._episode_length, + "success": success, + } + ) + if self._viewer: + remaining = self._policy_context.policy_dt - (time.perf_counter() - started) + if remaining > 0.0: + time.sleep(remaining) + return EnvironmentStep( + frame=self._frame, + termination_reason=reason, + metrics=metrics, + ) + + def metrics(self) -> dict[str, float]: + """Aggregate completed task episodes.""" + if not self._episodes: + return {} + return { + "eval/avg_reward": float( + np.mean([float(episode["reward"]) for episode in self._episodes]) + ), + "eval/avg_length": float( + np.mean([float(episode["length"]) for episode in self._episodes]) + ), + "eval/success_rate": float( + np.mean([bool(episode["success"]) for episode in self._episodes]) + ), + } + + def wait_for_reset_or_close(self) -> str: + """Keep a paused Viewer responsive until it is closed.""" + while self.viewer_is_open: + event = self.poll() + if event is not None: + return event + world = self._world() + if world is not None: + world.update(0.0) + time.sleep(0.01) + return "viewer closed" + + def close(self) -> None: + """Close the original task Environment.""" + if self._closed: + return + if self._previous_no_auto_reset is _MISSING: + delattr(self._base_env, "_demo_no_auto_reset") + else: + self._base_env._demo_no_auto_reset = self._previous_no_auto_reset + if getattr(self._base_env, "sim", None) is not None: + self._base_env.close(exit_process=False) + else: + self.env.close() + self._closed = True + + def _make_frame( + self, + observation: object, + task_state: Mapping[str, object], + ) -> EvaluationFrame: + simulation_step = ( + self._control_step * self._policy_context.sim_steps_per_control + ) + return EvaluationFrame( + control_step=self._control_step, + policy_time=self._control_step * self._policy_context.policy_dt, + simulation_step=simulation_step, + simulation_time=simulation_step * self._policy_context.physics_dt, + observation=observation, + task_state=task_state, + ) + + def _world(self) -> Any | None: + sim = getattr(self._base_env, "sim", None) + return None if sim is None else sim.get_world() + + +def evaluate_native_task( + runtime: PolicyRuntime, + *, + seed: int, + viewer: bool, + episodes: int | None, + control_steps: int | None, + duration: float | None, + termination_behavior: str = "auto_reset", +) -> NativeTaskEvaluationResult: + """Evaluate an EmbodiChain Policy in the task used for training.""" + if episodes is not None and episodes <= 0: + raise ValueError("episodes must be positive") + if control_steps is not None and control_steps <= 0: + raise ValueError("control_steps must be positive") + if duration is not None and (duration <= 0.0 or not math.isfinite(duration)): + raise ValueError("duration must be finite and positive") + if control_steps is not None and duration is not None: + raise ValueError("control_steps and duration are mutually exclusive") + if termination_behavior == "continue": + raise ValueError("Native task evaluation supports pause or auto_reset") + + try: + environment = EmbodiChainTaskEnvironment( + runtime.env, + seed=seed, + viewer=viewer, + ) + except Exception: + runtime.env.close() + raise + adapter = EmbodiChainTaskPolicyAdapter(runtime.policy, runtime.device, num_envs=1) + if duration is not None: + control_steps = math.ceil( + duration / environment.policy_context.policy_dt - 1e-12 + ) + total_steps = 0 + reason = "viewer closed" if viewer else "episode target reached" + options = RunOptions( + headless=not viewer, + termination_behavior=( + "continue" if termination_behavior == "auto_reset" else "pause" + ), + ) + with create_motion_policy_evaluator( + options=options, + adapter=adapter, + environment=environment, + title=f"{runtime.env_id} - EmbodiChain", + ) as evaluator: + evaluator.reset() + while True: + if control_steps is not None and total_steps >= control_steps: + reason = "control steps reached" + break + if episodes is not None and len(environment.episodes) >= episodes: + reason = "episode target reached" + break + completed_before = len(environment.episodes) + result = evaluator.step() + if result.advanced: + total_steps += 1 + if len(environment.episodes) > completed_before: + if episodes is not None and len(environment.episodes) >= episodes: + reason = "episode target reached" + break + if termination_behavior == "auto_reset": + evaluator.reset() + continue + if result.reason is not None and not result.reset_performed: + reason = result.reason + break + episode_results = environment.episodes + metrics = environment.metrics() + context = environment.policy_context + backend = environment.physics_backend + + simulation_steps = total_steps * context.sim_steps_per_control + return NativeTaskEvaluationResult( + task_id=runtime.env_id, + reason=reason, + simulation_time=simulation_steps * context.physics_dt, + simulation_steps=simulation_steps, + control_steps=total_steps, + physics_backend=backend, + effective_duration=total_steps * context.policy_dt, + requested_duration=duration, + episodes=episode_results, + metrics=metrics, + viewer=viewer, + ) + + +def _single_float(value: object, name: str) -> float: + tensor = torch.as_tensor(value).reshape(-1) + if tensor.numel() != 1: + raise ValueError(f"Native task {name} must contain one value") + return float(tensor.item()) + + +def _single_bool(value: object, name: str) -> bool: + tensor = torch.as_tensor(value, dtype=torch.bool).reshape(-1) + if tensor.numel() != 1: + raise ValueError(f"Native task {name} must contain one value") + return bool(tensor.item()) + + +def _info_bool(info: object, name: str) -> bool: + if not isinstance(info, Mapping) or name not in info: + return False + return _single_bool(info[name], f"info.{name}") + + +def _termination_reason( + info: object, + terminated: bool, + truncated: bool, +) -> str | None: + if _info_bool(info, "success"): + return "success" + if _info_bool(info, "fail"): + return "failure" + if truncated: + return "time limit" + if terminated: + return "terminated" + return None + + +def _step_metrics(info: object, reward: float) -> dict[str, float]: + result = {"reward": reward} + if not isinstance(info, Mapping): + return result + metrics = info.get("metrics") + if not isinstance(metrics, Mapping): + return result + for name, value in metrics.items(): + tensor = torch.as_tensor(value).reshape(-1) + if tensor.numel() == 1: + result[str(name)] = float(tensor.item()) + return result + + +def _policy_context_from_env(env: Any) -> PolicyContext: + """Read timing from a simulator task or lightweight learning task.""" + if hasattr(env, "physics_dt") and hasattr(env, "step_dt"): + return PolicyContext( + robot=None, + physics_dt=float(env.physics_dt), + sim_steps_per_control=int(env.cfg.sim_steps_per_control), + policy_dt=float(env.step_dt), + ) + if not hasattr(env, "dt"): + raise ValueError("Lightweight task evaluation requires an env.dt value") + policy_dt = float(env.dt) + return PolicyContext( + robot=None, + physics_dt=policy_dt, + sim_steps_per_control=1, + policy_dt=policy_dt, + ) diff --git a/embodichain/learning/rl/motion_policy_evaluation/profile.py b/embodichain/learning/rl/motion_policy_evaluation/profile.py new file mode 100644 index 000000000..ed107c885 --- /dev/null +++ b/embodichain/learning/rl/motion_policy_evaluation/profile.py @@ -0,0 +1,148 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Motion Profile providers for task-specific model and control semantics.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from dataclasses import dataclass, field +from pathlib import Path +from types import MappingProxyType +from typing import Any + +import torch + +__all__ = [ + "MotionProfile", + "MotionProfileRequest", + "build_motion_profile", + "get_motion_profile_names", + "register_motion_profile", +] + + +@dataclass(frozen=True) +class MotionProfileRequest: + """Checkpoint, configs, and runtime choices supplied to a provider.""" + + checkpoint: Path + device: torch.device + configs: Mapping[str, Path] = field(default_factory=dict) + resource_root: Path | None = None + renderer: str = "hybrid" + + def __post_init__(self) -> None: + checkpoint = Path(self.checkpoint).expanduser().resolve() + if not checkpoint.is_file(): + raise FileNotFoundError(f"Motion checkpoint does not exist: {checkpoint}") + configs = { + name: Path(path).expanduser().resolve() + for name, path in self.configs.items() + } + for name, path in configs.items(): + if not path.is_file(): + raise FileNotFoundError( + f"Motion config {name!r} does not exist: {path}" + ) + root = ( + None + if self.resource_root is None + else Path(self.resource_root).expanduser().resolve() + ) + object.__setattr__(self, "checkpoint", checkpoint) + object.__setattr__(self, "configs", MappingProxyType(configs)) + object.__setattr__(self, "resource_root", root) + + +@dataclass(frozen=True) +class MotionProfile: + """DexSim Policy Spec and report metadata built by one provider.""" + + profile_id: str + checkpoint: Path + policy_spec: Mapping[str, Any] + provider_version: int = 1 + provenance: Mapping[str, Any] = field(default_factory=dict) + warnings: tuple[str, ...] = () + + def __post_init__(self) -> None: + checkpoint = Path(self.checkpoint).expanduser().resolve() + if not checkpoint.is_file(): + raise FileNotFoundError(f"Motion checkpoint does not exist: {checkpoint}") + object.__setattr__(self, "checkpoint", checkpoint) + object.__setattr__( + self, + "policy_spec", + MappingProxyType(dict(self.policy_spec)), + ) + object.__setattr__( + self, + "provenance", + MappingProxyType(dict(self.provenance)), + ) + object.__setattr__(self, "warnings", tuple(self.warnings)) + + +MotionProfileProvider = Callable[[MotionProfileRequest], MotionProfile] +_PROVIDERS: dict[str, MotionProfileProvider] = {} + + +def register_motion_profile(name: str, provider: MotionProfileProvider) -> None: + """Register a Motion Profile provider under its CLI name. + + Args: + name: Stable profile name. + provider: Callable that builds one :class:`MotionProfile`. + """ + if not name: + raise ValueError("Motion profile name must not be empty") + if name in _PROVIDERS: + raise ValueError(f"Motion profile is already registered: {name}") + _PROVIDERS[name] = provider + + +def get_motion_profile_names() -> tuple[str, ...]: + """Return the registered profile names.""" + return tuple(sorted(_PROVIDERS)) + + +def build_motion_profile( + name: str, + request: MotionProfileRequest, +) -> MotionProfile: + """Build one profile with its registered provider. + + Args: + name: Registered profile name. + request: Checkpoint, configs, and runtime choices. + + Returns: + Provider-built Motion Profile. + """ + try: + provider = _PROVIDERS[name] + except KeyError: + available = ", ".join(sorted(_PROVIDERS)) or "none" + raise ValueError( + f"Unknown motion profile {name!r}; available: {available}" + ) from None + profile = provider(request) + if profile.profile_id != name: + raise ValueError( + f"Motion provider {name!r} returned profile {profile.profile_id!r}" + ) + return profile diff --git a/embodichain/learning/rl/motion_policy_evaluation/report.py b/embodichain/learning/rl/motion_policy_evaluation/report.py new file mode 100644 index 000000000..02060fea0 --- /dev/null +++ b/embodichain/learning/rl/motion_policy_evaluation/report.py @@ -0,0 +1,78 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Write one timestamped motion-policy evaluation report.""" + +from __future__ import annotations + +import json +import math +from collections.abc import Mapping, Sequence +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +__all__ = ["write_evaluation_report"] + + +def write_evaluation_report( + parent: str | Path, + payload: Mapping[str, Any], +) -> Path: + """Write ``evaluation.json`` under a new timestamped directory. + + Args: + parent: Output parent directory. + payload: Evaluation inputs and results. + + Returns: + Written report path. + """ + output = Path(parent).expanduser().resolve() + output.mkdir(parents=True, exist_ok=True) + stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") + directory = output / f"{stamp}-motion-policy" + directory.mkdir() + report = { + "schema_version": 1, + "created_at": datetime.now(timezone.utc).isoformat(), + **dict(payload), + } + path = directory / "evaluation.json" + path.write_text( + json.dumps( + _json_value(report), + indent=2, + sort_keys=True, + ensure_ascii=False, + allow_nan=False, + ) + + "\n", + encoding="utf-8", + ) + return path + + +def _json_value(value: Any) -> Any: + if isinstance(value, Path): + return str(value) + if isinstance(value, float): + return value if math.isfinite(value) else None + if isinstance(value, Mapping): + return {str(name): _json_value(item) for name, item in value.items()} + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return [_json_value(item) for item in value] + return value diff --git a/embodichain/learning/rl/runtime.py b/embodichain/learning/rl/runtime.py new file mode 100644 index 000000000..d70fb07f3 --- /dev/null +++ b/embodichain/learning/rl/runtime.py @@ -0,0 +1,350 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Shared environment and Policy construction for RL training and evaluation.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from embodichain.lab.gym.utils.gym_utils import config_to_cfg, get_manager_modules +from embodichain.lab.gym.utils.profiler import EnvProfilerCfg +from embodichain.lab.gym.utils.registration import build_env +from embodichain.lab.sim import SimulationManagerCfg +from embodichain.lab.sim.cfg import RenderCfg +from embodichain.learning.rl.env import build_learning_env +from embodichain.learning.rl.models import build_mlp_from_cfg, build_policy +from embodichain.learning.rl.utils import dict_to_tensordict, flatten_dict_observation +from embodichain.utils.utility import load_config + +__all__ = [ + "GymEnvironmentRuntime", + "GymPolicyFactory", + "PolicyRuntime", + "build_gym_environment", + "build_gym_policy_runtime", + "build_learning_environment", + "build_learning_policy_runtime", + "resolve_config_reference", + "resolve_torch_device", + "seed_policy_runtime", +] + +GymPolicyFactory = Callable[[int, int, torch.device], torch.nn.Module] +"""Build an evaluation Policy after the task dimensions are known.""" + + +@dataclass(frozen=True) +class GymEnvironmentRuntime: + """A simulator task reconstructed from one training configuration.""" + + env: Any + env_id: str + env_cfg: Any + gym_config: dict[str, Any] + gym_config_path: Path + + +@dataclass(frozen=True) +class PolicyRuntime: + """An Environment and Policy reconstructed from one training configuration.""" + + env: Any + policy: torch.nn.Module + device: torch.device + env_id: str + env_cfg: Any | None = None + gym_config: dict[str, Any] | None = None + gym_config_path: Path | None = None + + +def resolve_torch_device(value: str | torch.device) -> torch.device: + """Validate and activate one CPU or CUDA device.""" + if not isinstance(value, (str, torch.device)): + raise TypeError("device must be a string or torch.device") + try: + device = torch.device(value) + except RuntimeError as error: + raise ValueError(f"Failed to parse device {value!r}: {error}") from error + + if device.type == "cuda": + if not torch.cuda.is_available(): + raise ValueError( + "CUDA was requested but torch.cuda.is_available() is False." + ) + index = device.index + if index is None: + index = torch.cuda.current_device() + if index < 0 or index >= torch.cuda.device_count(): + raise ValueError( + f"CUDA device index {index} is out of range " + f"(available devices: {torch.cuda.device_count()})." + ) + torch.cuda.set_device(index) + return torch.device(f"cuda:{index}") + if device.type != "cpu": + raise ValueError(f"Unsupported device type: {device.type}") + return torch.device("cpu") + + +def seed_policy_runtime(seed: int, device: torch.device) -> None: + """Seed NumPy and PyTorch for one Policy runtime.""" + np.random.seed(seed) + torch.manual_seed(seed) + torch.backends.cudnn.deterministic = True + if device.type == "cuda": + torch.cuda.manual_seed_all(seed) + + +def resolve_config_reference( + value: str | Path, + *, + base_dir: str | Path | None = None, +) -> Path: + """Resolve a referenced config relative to its containing config file.""" + path = Path(value).expanduser() + if path.is_absolute(): + return path + if base_dir is not None: + candidate = Path(base_dir).expanduser().resolve() / path + if candidate.exists(): + return candidate + return path + + +def build_learning_environment( + config: dict[str, Any], + *, + device: torch.device, + num_envs: int, +) -> tuple[str, Any]: + """Build the lightweight Environment declared by a training config.""" + env_block = config["trainer"]["learning_env"] + if isinstance(env_block, str): + env_name = env_block + env_cfg: dict[str, Any] = {} + else: + env_name = env_block["name"] + env_cfg = dict(env_block.get("cfg", {})) + return str(env_name), build_learning_env( + env_name, + num_envs=num_envs, + device=device, + **env_cfg, + ) + + +def build_learning_policy_runtime( + config: dict[str, Any], + *, + device: torch.device, + num_envs: int, +) -> PolicyRuntime: + """Build a lightweight Environment and its configured Policy.""" + env_name, env = build_learning_environment( + config, + num_envs=num_envs, + device=device, + ) + try: + observation_dim = int(env.single_observation_space.shape[-1]) + action_dim = int(env.single_action_space.shape[-1]) + policy_block = config["policy"] + actor_cfg = policy_block.get("actor") + critic_cfg = policy_block.get("critic") + actor = ( + build_mlp_from_cfg(actor_cfg, observation_dim, action_dim) + if actor_cfg is not None + else None + ) + critic = ( + build_mlp_from_cfg(critic_cfg, observation_dim, 1) + if critic_cfg is not None + else None + ) + policy = build_policy( + policy_block, + env.single_observation_space, + env.single_action_space, + device, + actor=actor, + critic=critic, + ) + if "initial_log_std" in policy_block and hasattr(policy, "log_std"): + with torch.no_grad(): + policy.log_std.fill_(float(policy_block["initial_log_std"])) + except Exception: + env.close() + raise + return PolicyRuntime(env, policy, device, env_name) + + +def build_gym_environment( + config: dict[str, Any], + *, + simulation_device: torch.device, + num_envs: int | None, + headless: bool, + renderer: str, + gpu_id: int, + config_dir: str | Path | None = None, + profiler: EnvProfilerCfg | None = None, +) -> GymEnvironmentRuntime: + """Build the simulator Environment declared by a training config.""" + trainer_cfg = config["trainer"] + gym_config_path = resolve_config_reference( + trainer_cfg["gym_config"], + base_dir=config_dir, + ) + gym_config = load_config(gym_config_path) + env_cfg = config_to_cfg(gym_config, manager_modules=get_manager_modules()) + if num_envs is not None: + env_cfg.num_envs = int(num_envs) + if env_cfg.sim_cfg is None: + env_cfg.sim_cfg = SimulationManagerCfg() + env_cfg.sim_cfg.sim_device = simulation_device + env_cfg.sim_cfg.headless = headless + env_cfg.sim_cfg.render_cfg = RenderCfg(renderer=renderer) + env_cfg.sim_cfg.gpu_id = ( + simulation_device.index + if simulation_device.type == "cuda" and simulation_device.index is not None + else gpu_id + ) + env_cfg.profiler = profiler + env = build_env(gym_config["id"], base_env_cfg=env_cfg) + return GymEnvironmentRuntime( + env=env, + env_id=str(gym_config["id"]), + env_cfg=env_cfg, + gym_config=gym_config, + gym_config_path=gym_config_path.resolve(), + ) + + +def build_gym_policy_runtime( + config: dict[str, Any], + *, + device: torch.device, + num_envs: int | None, + headless: bool, + renderer: str, + gpu_id: int, + config_dir: str | Path | None = None, + profiler: EnvProfilerCfg | None = None, + policy_factory: GymPolicyFactory | None = None, + simulation_device: torch.device | None = None, +) -> PolicyRuntime: + """Build a simulator task and the Policy declared by its training config.""" + task = build_gym_environment( + config, + simulation_device=simulation_device or device, + num_envs=num_envs, + headless=headless, + renderer=renderer, + gpu_id=gpu_id, + config_dir=config_dir, + profiler=profiler, + ) + env = task.env + try: + sample_observation, _ = env.reset() + sample_observation_td = dict_to_tensordict(sample_observation, device) + observation_dim = int(flatten_dict_observation(sample_observation_td).shape[-1]) + action_manager = env.get_wrapper_attr("action_manager") + environment_action_dim = ( + action_manager.total_action_dim + if action_manager is not None + else len(env.get_wrapper_attr("active_joint_ids")) + ) + if policy_factory is None: + policy = _build_gym_policy( + config["policy"], + env=env, + device=device, + observation_dim=observation_dim, + action_dim=environment_action_dim, + ) + else: + policy = policy_factory(observation_dim, environment_action_dim, device) + if not isinstance(policy, torch.nn.Module): + raise TypeError("policy_factory must return a torch.nn.Module.") + except Exception: + env.close() + raise + return PolicyRuntime( + env=env, + policy=policy, + device=device, + env_id=task.env_id, + env_cfg=task.env_cfg, + gym_config=task.gym_config, + gym_config_path=task.gym_config_path, + ) + + +def _build_gym_policy( + policy_block: dict[str, Any], + *, + env: Any, + device: torch.device, + observation_dim: int, + action_dim: int, +) -> torch.nn.Module: + configured_action_dim = int(policy_block.get("action_dim", action_dim)) + if configured_action_dim != action_dim: + raise ValueError( + f"Configured policy.action_dim={configured_action_dim} does not match " + f"env action dim {action_dim}." + ) + policy_name = str(policy_block["name"]).lower() + if policy_name == "actor_critic": + actor_cfg = policy_block.get("actor") + critic_cfg = policy_block.get("critic") + if actor_cfg is None or critic_cfg is None: + raise ValueError( + "ActorCritic requires policy.actor and policy.critic definitions." + ) + return build_policy( + policy_block, + env.flattened_observation_space, + env.action_space, + device, + actor=build_mlp_from_cfg(actor_cfg, observation_dim, action_dim), + critic=build_mlp_from_cfg(critic_cfg, observation_dim, 1), + ) + if policy_name == "actor_only": + actor_cfg = policy_block.get("actor") + if actor_cfg is None: + raise ValueError("ActorOnly requires a policy.actor definition.") + return build_policy( + policy_block, + env.flattened_observation_space, + env.action_space, + device, + actor=build_mlp_from_cfg(actor_cfg, observation_dim, action_dim), + ) + return build_policy( + policy_block, + env.observation_space, + env.action_space, + device, + ) diff --git a/embodichain/learning/rl/train.py b/embodichain/learning/rl/train.py index 41c44a3f1..8b1da5441 100644 --- a/embodichain/learning/rl/train.py +++ b/embodichain/learning/rl/train.py @@ -26,10 +26,11 @@ import torch import wandb from torch.utils.tensorboard import SummaryWriter -from copy import deepcopy -from embodichain.learning.rl.models import build_policy, get_registered_policy_names -from embodichain.learning.rl.models import build_mlp_from_cfg +from embodichain.learning.rl.models import get_registered_policy_names +from embodichain.learning.rl.motion_policy_evaluation.manifest import ( + write_run_manifest, +) from embodichain.learning.rl.algo import ( RolloutKind, build_algo, @@ -40,21 +41,21 @@ DifferentiableTrainerCfg, ) from embodichain.learning.rl.env import build_learning_env +from embodichain.learning.rl.runtime import ( + build_gym_environment, + build_gym_policy_runtime, + build_learning_policy_runtime, +) from embodichain.learning.rl.routing import get_trainer_class -from embodichain.learning.rl.utils import dict_to_tensordict, flatten_dict_observation from embodichain.learning.rl.utils.trainer import Trainer from embodichain.utils import logger from embodichain.lab.gym.utils.registration import ( - build_env, discover_task_packages, execute_init_hooks, ) -from embodichain.lab.gym.utils.gym_utils import config_to_cfg, get_manager_modules from embodichain.lab.gym.utils.profiler import EnvProfilerCfg from embodichain.utils.utility import load_config from embodichain.utils.module_utils import find_function_from_modules -from embodichain.lab.sim import SimulationManagerCfg -from embodichain.lab.sim.cfg import RenderCfg from embodichain.lab.gym.envs.managers.cfg import EventCfg @@ -114,51 +115,19 @@ def _resolve_profile_output( return str(output.with_name(f"{output.stem}_rank{rank}{output.suffix}")) -def _build_learning_policy( - policy_block: dict, - env, - device: torch.device, -): - obs_dim = int(env.single_observation_space.shape[-1]) - action_dim = int(env.single_action_space.shape[-1]) - policy_name = policy_block["name"].lower() - actor_cfg = policy_block.get("actor") - critic_cfg = policy_block.get("critic") - actor = ( - build_mlp_from_cfg(actor_cfg, obs_dim, action_dim) - if actor_cfg is not None - else None - ) - critic = ( - build_mlp_from_cfg(critic_cfg, obs_dim, 1) if critic_cfg is not None else None - ) - policy = build_policy( - policy_block, - env.single_observation_space, - env.single_action_space, - device, - actor=actor, - critic=critic, - ) - if "initial_log_std" in policy_block and hasattr(policy, "log_std"): - with torch.no_grad(): - policy.log_std.fill_(float(policy_block["initial_log_std"])) - return policy - - def _train_learning_env( cfg_data: dict, *, + config_path: str | Path, distributed: bool | None, profile: bool = False, -): +) -> dict[str, object]: """Train a lightweight registered environment through the unified CLI.""" if profile: raise ValueError( "--profile requires trainer.gym_config; learning_env is unsupported." ) trainer_cfg = cfg_data["trainer"] - policy_block = cfg_data["policy"] algorithm_block = cfg_data["algorithm"] distributed = ( bool(trainer_cfg.get("distributed", False)) @@ -182,24 +151,24 @@ def _train_learning_env( np.random.seed(seed) torch.manual_seed(seed) - env_block = trainer_cfg["learning_env"] - if isinstance(env_block, str): - env_name = env_block - env_cfg = {} - else: - env_name = env_block["name"] - env_cfg = dict(env_block.get("cfg", {})) num_envs = int(trainer_cfg.get("num_envs", 64)) - env = build_learning_env( - env_name, + runtime = build_learning_policy_runtime( + cfg_data, num_envs=num_envs, device=device, - **env_cfg, ) + env = runtime.env + policy = runtime.policy + env_name = runtime.env_id enable_eval = bool(trainer_cfg.get("enable_eval", False)) eval_env = None if enable_eval: + env_block = trainer_cfg["learning_env"] + if isinstance(env_block, str): + env_cfg = {} + else: + env_cfg = dict(env_block.get("cfg", {})) eval_env = build_learning_env( env_name, num_envs=int(trainer_cfg.get("num_eval_envs", 16)), @@ -207,7 +176,6 @@ def _train_learning_env( **env_cfg, ) - policy = _build_learning_policy(policy_block, env, device) algorithm = build_algo( algorithm_block["name"], dict(algorithm_block.get("cfg", {})), @@ -291,7 +259,14 @@ def _train_learning_env( total_timesteps = int(trainer_cfg.get("total_timesteps", default_steps)) trainer.train(total_timesteps) trainer.save_checkpoint() - return trainer.get_summary() + summary = trainer.get_summary() + _write_motion_run_manifest( + run_base, + config_path, + trainer_cfg, + summary, + ) + return summary finally: writer.close() if use_wandb: @@ -307,7 +282,7 @@ def train_from_config( *, profile: bool = False, profile_output: str | None = None, -): +) -> dict[str, object] | None: """Run training from a config file path. Args: @@ -316,6 +291,9 @@ def train_from_config( If None, use trainer.distributed from config. profile: Enable gym ``EnvProfiler`` on the training environment. profile_output: Optional JSON dump path for the profiling report. + + Returns: + The lightweight trainer summary, or ``None`` for simulator training. """ if profile_output is not None and not profile: raise ValueError("--profile_output requires --profile.") @@ -326,6 +304,7 @@ def train_from_config( if "learning_env" in trainer_cfg: return _train_learning_env( cfg_data, + config_path=config_path, distributed=distributed, profile=profile, ) @@ -441,32 +420,11 @@ def train_from_config( if use_wandb and rank == 0: wandb.init(project=wandb_project_name, name=exp_name, config=cfg_data) - gym_config_path = Path(trainer_cfg["gym_config"]) if rank == 0: logger.log_info(f"Current working directory: {Path.cwd()}") - gym_config_data = load_config(str(gym_config_path)) - gym_env_cfg = config_to_cfg(gym_config_data, manager_modules=get_manager_modules()) - if num_envs is not None: - gym_env_cfg.num_envs = int(num_envs) - - # Ensure sim configuration mirrors runtime overrides - if gym_env_cfg.sim_cfg is None: - gym_env_cfg.sim_cfg = SimulationManagerCfg() - if device.type == "cuda": - gpu_index = device.index - if gpu_index is None: - gpu_index = torch.cuda.current_device() - gym_env_cfg.sim_cfg.sim_device = torch.device(f"cuda:{gpu_index}") - if hasattr(gym_env_cfg.sim_cfg, "gpu_id"): - gym_env_cfg.sim_cfg.gpu_id = gpu_index - else: - gym_env_cfg.sim_cfg.sim_device = torch.device("cpu") - gym_env_cfg.sim_cfg.headless = headless - gym_env_cfg.sim_cfg.render_cfg = RenderCfg(renderer=renderer) - gym_env_cfg.sim_cfg.gpu_id = gpu_id - if profile: - gym_env_cfg.profiler = EnvProfilerCfg( + profiler = ( + EnvProfilerCfg( enable_time=True, output_path=_resolve_profile_output( profile_output, @@ -474,83 +432,50 @@ def train_from_config( world_size=world_size, ), ) + if profile + else None + ) + runtime = build_gym_policy_runtime( + cfg_data, + device=device, + num_envs=num_envs, + headless=headless, + renderer=renderer, + gpu_id=gpu_id, + config_dir=Path(config_path).expanduser().resolve().parent, + profiler=profiler, + ) + env = runtime.env + policy = runtime.policy + gym_config_path = runtime.gym_config_path + gym_config_data = runtime.gym_config + gym_env_cfg = runtime.env_cfg + if gym_config_path is None or gym_config_data is None or gym_env_cfg is None: + raise RuntimeError("Simulator Policy runtime is missing task configuration") if rank == 0: logger.log_info( f"Loaded gym_config from {gym_config_path} (env_id={gym_config_data['id']}, num_envs={gym_env_cfg.num_envs}, headless={gym_env_cfg.sim_cfg.headless}, renderer={gym_env_cfg.sim_cfg.render_cfg.renderer}, sim_device={gym_env_cfg.sim_cfg.sim_device})" ) - env = build_env(gym_config_data["id"], base_env_cfg=gym_env_cfg) - sample_obs, _ = env.reset() - sample_obs_td = dict_to_tensordict(sample_obs, device) - obs_dim = flatten_dict_observation(sample_obs_td).shape[-1] - flat_obs_space = env.flattened_observation_space - # Create evaluation environment only if enabled eval_env = None num_eval_envs = trainer_cfg.get("num_eval_envs", 4) if enable_eval and rank == 0: - eval_gym_env_cfg = deepcopy(gym_env_cfg) - eval_gym_env_cfg.num_envs = num_eval_envs - eval_gym_env_cfg.sim_cfg.headless = True - eval_gym_env_cfg.profiler = None - eval_env = build_env(gym_config_data["id"], base_env_cfg=eval_gym_env_cfg) + eval_runtime = build_gym_environment( + cfg_data, + simulation_device=device, + num_envs=int(num_eval_envs), + headless=True, + renderer=renderer, + gpu_id=gpu_id, + config_dir=Path(config_path).expanduser().resolve().parent, + ) + eval_env = eval_runtime.env logger.log_info( f"Evaluation environment created (num_envs={num_eval_envs}, headless=True)" ) - # Build Policy via registry policy_name = policy_block["name"] - env_action_dim = ( - env.get_wrapper_attr("action_manager").total_action_dim - if env.get_wrapper_attr("action_manager") is not None - else len(env.get_wrapper_attr("active_joint_ids")) - ) - action_dim = policy_block.get("action_dim", env_action_dim) - action_dim = int(action_dim) - if action_dim != env_action_dim: - raise ValueError( - f"Configured policy.action_dim={action_dim} does not match env action dim {env_action_dim}." - ) - # Build Policy via registry (actor/critic must be explicitly defined in JSON when using actor_critic/actor_only) - if policy_name.lower() == "actor_critic": - actor_cfg = policy_block.get("actor") - critic_cfg = policy_block.get("critic") - if actor_cfg is None or critic_cfg is None: - raise ValueError( - "ActorCritic requires 'actor' and 'critic' definitions in JSON (policy.actor / policy.critic)." - ) - - actor = build_mlp_from_cfg(actor_cfg, obs_dim, action_dim) - critic = build_mlp_from_cfg(critic_cfg, obs_dim, 1) - - policy = build_policy( - policy_block, - flat_obs_space, - env.action_space, - device, - actor=actor, - critic=critic, - ) - elif policy_name.lower() == "actor_only": - actor_cfg = policy_block.get("actor") - if actor_cfg is None: - raise ValueError( - "ActorOnly requires 'actor' definition in JSON (policy.actor)." - ) - - actor = build_mlp_from_cfg(actor_cfg, obs_dim, action_dim) - - policy = build_policy( - policy_block, - flat_obs_space, - env.action_space, - device, - actor=actor, - ) - else: - policy = build_policy( - policy_block, env.observation_space, env.action_space, device - ) # Build Algorithm via factory algo_name = algo_block["name"].lower() @@ -682,9 +607,38 @@ def train_from_config( torch.distributed.destroy_process_group() if rank == 0: + _write_motion_run_manifest( + run_base, + config_path, + trainer_cfg, + trainer.get_summary(), + gym_config=gym_config_path, + ) logger.log_info("Training finished") +def _write_motion_run_manifest( + run_base: str | Path, + config_path: str | Path, + trainer_cfg: dict, + summary: dict, + *, + gym_config: str | Path | None = None, +) -> Path: + """Write the final checkpoint and config index for motion evaluation.""" + latest = summary.get("latest_checkpoint_path") + if latest is None: + raise RuntimeError("Training finished without a checkpoint") + return write_run_manifest( + run_base, + train_config=config_path, + gym_config=gym_config, + latest_checkpoint=latest, + best_checkpoint=summary.get("best_checkpoint_path"), + motion_profile=trainer_cfg.get("motion_profile"), + ) + + def cli(argv: Sequence[str] | None = None) -> None: """Command-line interface for RL training. diff --git a/examples/learning/motion_policy_evaluation/README.md b/examples/learning/motion_policy_evaluation/README.md new file mode 100644 index 000000000..5dc227685 --- /dev/null +++ b/examples/learning/motion_policy_evaluation/README.md @@ -0,0 +1,130 @@ +# ANYmal-C Velocity Policy Evaluation + +This example connects Newton's public ANYmal-C velocity TorchScript `.pt` to +EmbodiChain and opens it in the DexSim Viewer through Motion Policy Kit. The +model accepts `vx`, `vy`, and `yaw` commands. W/A/S/D and Q/E update these +commands while the Viewer is running. + +The model, configuration, and robot resources come from newton-assets commit +`261cd1f429619d8ef4f546bd788ab9dea906b5e1`. The Policy is distributed under +Apache-2.0, and the robot resources use BSD-3-Clause. The Adapter follows +Newton v1.2.1 +[`example_robot_policy.py`](https://github.com/newton-physics/newton/blob/v1.2.1/newton/examples/robot/example_robot_policy.py) +to reproduce the 48-dimensional observation, TorchScript inference, and joint +target processing. + +## Directory layout + +```text +motion_policy_evaluation/ +├── README.md +├── prepare_resources.py +├── eval_policy.py # Register the local Profile and run the example +└── anymal_c/ + ├── __init__.py # Register newton-anymal-c-velocity + └── profile.py # Policy Spec and AnymalCVelocityAdapter +``` + +Resource preparation creates this local cache: + +```text +~/.cache/embodichain/examples/anymal_c_velocity/ +└── upstream/ + └── anybotics_anymal_c/ + ├── rl_policies/ + │ ├── mjw_anymal.pt + │ ├── anymal.yaml + │ └── LICENSE + ├── urdf/anymal.urdf + ├── meshes/... + └── LICENSE +``` + +## Run the example + +Run these commands from the EmbodiChain repository root. The preparation script +prints the model, asset, checkout, and digest verification progress. Re-running +the command continues an existing Git checkout after an interrupted download. + +```bash +python examples/learning/motion_policy_evaluation/prepare_resources.py +python examples/learning/motion_policy_evaluation/eval_policy.py \ + --viewer \ + --renderer hybrid +``` + +Viewer controls: + +| Key | Command | +|---|---| +| W / S | Increase / decrease `vx` | +| A / D | Increase / decrease `vy` | +| Q / E | Increase / decrease `yaw` | +| M | Set all three commands to zero | +| R | Reset the robot and Policy history | + +The terminal prints the path to `evaluation.json` when the Viewer closes. Run a +Headless smoke test with: + +```bash +python examples/learning/motion_policy_evaluation/eval_policy.py \ + --device cpu \ + --sim-device cpu \ + --control-steps 20 +``` + +`eval_policy.py` reads the checkpoint and robot assets from the default cache, +imports the adjacent `anymal_c/profile.py`, and registers the Profile in the +current process. Run the script directly from the repository root. To use +another cache directory: + +```bash +python examples/learning/motion_policy_evaluation/prepare_resources.py \ + --output /tmp/anymal_c_velocity + +ANYMAL_C_EXAMPLE_CACHE=/tmp/anymal_c_velocity \ + python examples/learning/motion_policy_evaluation/eval_policy.py --viewer +``` + +## Execution pipeline + +```mermaid +flowchart LR + CLI[eval-motion-policy] --> Profile[build_profile] + Profile --> Spec[Policy Spec
assets, control parameters, frequency] + Spec --> Setup[Adapter.setup
load TorchScript and joint mapping] + Setup --> State[read RobotState] + Command[WASD + QE command] --> Obs[build 48-dimensional observation] + State --> Obs + Obs --> Actor[TorchScript actor] + Actor --> Action[map to 12 joint targets] + Action --> Sim[advance the Environment] +``` + +`AnymalCVelocityAdapter` restores the upstream data path: + +| Stage | Processing | +|---|---| +| `setup()` | Build a `JointMap` for the 12 ANYmal-C joints, then load and validate the TorchScript inputs and outputs | +| observation | 3 body linear velocity, 3 body angular velocity, 3 projected gravity, 3 command, 12 joint position, 12 joint velocity, and 12 previous action values | +| command | Read `vx`, `vy`, and `yaw` from `frame.controls["command"]`, with ranges ±1.0, ±0.5, and ±1.0 | +| actor | Pass a `[1, 48]` tensor through the model's normalizer and actor to produce `[1, 12]` | +| action | Apply `default_position + 0.5 * action` | +| control | Run simulation at 200 Hz and infer once every four simulation steps for a 50 Hz Policy rate | + +The Adapter clears the previous action during reset. After each inference call, +it stores the current action for the next observation. + +## Integrate another external Policy + +Copy this directory and replace: + +1. the fixed revisions, paths, and digests for the model and robot resources in `prepare_resources.py`; +2. the initial pose, joint control parameters, simulation step, and `sim_steps_per_control` in `build_profile()`; +3. the model format and training joint order in `Adapter.setup()`; +4. observation construction, normalization, network forward pass, action clipping, scale, and offset in `Adapter.infer()`; +5. `PROFILE_ID` and the Profile name used by `eval_policy.py`. + +An Adapter can call an existing project data reader from `__init__()` or +`setup()`. To let Policy Spec resolve a data file path, declare it under +`policy.resources` and read the resolved path from `AdapterRequest.resources`. diff --git a/examples/learning/motion_policy_evaluation/anymal_c/__init__.py b/examples/learning/motion_policy_evaluation/anymal_c/__init__.py new file mode 100644 index 000000000..e18fe4904 --- /dev/null +++ b/examples/learning/motion_policy_evaluation/anymal_c/__init__.py @@ -0,0 +1,30 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Register the concrete ANYmal-C velocity Motion Profile.""" + +from __future__ import annotations + +from embodichain.learning.rl.motion_policy_evaluation import register_motion_profile + +from .profile import PROFILE_ID, build_profile + +__all__ = ["register"] + + +def register() -> None: + """Register the ANYmal-C velocity Profile for the example process.""" + register_motion_profile(PROFILE_ID, build_profile) diff --git a/examples/learning/motion_policy_evaluation/anymal_c/profile.py b/examples/learning/motion_policy_evaluation/anymal_c/profile.py new file mode 100644 index 000000000..edea62d0b --- /dev/null +++ b/examples/learning/motion_policy_evaluation/anymal_c/profile.py @@ -0,0 +1,295 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""ANYmal-C velocity Profile for Newton's public TorchScript policy.""" + +from __future__ import annotations + +import math +from pathlib import Path + +import numpy as np +import torch +from dexsim.kit.motion_policy import ( + AdapterRequest, + EvaluationFrame, + JointMap, + PolicyContext, + PolicyOutput, + require_finite, +) + +from embodichain.learning.rl.motion_policy_evaluation import ( + MotionProfile, + MotionProfileRequest, +) + +__all__ = ["AnymalCVelocityAdapter", "PROFILE_ID", "build_profile"] + +PROFILE_ID = "newton-anymal-c-velocity" + +_SOURCE_REVISION = "7249270ab41be1c2d4c809aa87536bab3a1a26f4" +_ASSET_REVISION = "261cd1f429619d8ef4f546bd788ab9dea906b5e1" +_ROBOT_PATH = Path("anybotics_anymal_c/urdf/anymal.urdf") +_JOINT_NAMES = ( + "LF_HAA", + "LF_HFE", + "LF_KFE", + "LH_HAA", + "LH_HFE", + "LH_KFE", + "RF_HAA", + "RF_HFE", + "RF_KFE", + "RH_HAA", + "RH_HFE", + "RH_KFE", +) +_DEFAULT_POSITION = np.asarray( + (0.0, 0.4, -0.8, 0.0, -0.4, 0.8, 0.0, 0.4, -0.8, 0.0, -0.4, 0.8), + dtype=np.float32, +) + + +def build_profile(request: MotionProfileRequest) -> MotionProfile: + """Build the Policy Spec for the public ANYmal-C checkpoint. + + Args: + request: Checkpoint, resource checkout, device, and renderer. + + Returns: + A Motion Profile ready for DexSim Motion Policy Kit. + """ + if request.checkpoint.suffix.lower() != ".pt": + raise ValueError("The ANYmal-C example requires a TorchScript .pt file") + if request.resource_root is None: + raise ValueError( + "The ANYmal-C example requires --resource-root from prepare_resources.py" + ) + robot_asset = request.resource_root / _ROBOT_PATH + if not robot_asset.is_file(): + raise FileNotFoundError( + f"ANYmal-C robot asset does not exist: {robot_asset}. " + "Run prepare_resources.py first." + ) + + return MotionProfile( + profile_id=PROFILE_ID, + checkpoint=request.checkpoint, + policy_spec={ + "schema_version": 1, + "kind": "policy", + "id": PROFILE_ID, + "metadata": { + "title": "Newton ANYmal-C velocity policy", + "description": "Public 48-D command locomotion TorchScript policy.", + "status": "example", + "tags": ["external", "quadruped", "velocity", "torchscript"], + }, + "robot": { + "asset": {"path": str(robot_asset)}, + "use_urdf_material": True, + "initial": { + "root_height": 0.76, + "joint_positions": { + "default": 0.0, + "overrides": dict( + zip( + _JOINT_NAMES, + _DEFAULT_POSITION.tolist(), + strict=True, + ) + ), + }, + }, + "control": { + "defaults": { + "stiffness": 300.0, + "damping": 10.0, + "effort_limit": 80.0, + "armature": 0.06, + }, + }, + }, + "policy": { + "models": {"actor": {"path": str(request.checkpoint)}}, + "adapter": { + "type": "python", + "entrypoint": ("anymal_c.profile:AnymalCVelocityAdapter"), + "config": { + "inference_device": str(request.device), + "joint_names": list(_JOINT_NAMES), + }, + }, + }, + "evaluation": { + "initial_command": [0.0, 0.0, 0.0], + "termination": {"behavior": "pause"}, + }, + "runtime": { + "physics_dt": 0.005, + "sim_steps_per_control": 4, + "physics_backend": "default", + "simulation_device": "cpu", + "inference_provider": ( + "cuda" if request.device.type == "cuda" else "cpu" + ), + "renderer": request.renderer, + }, + }, + provenance={ + "source": "Newton ANYmal-C keyboard policy example", + "source_revision": _SOURCE_REVISION, + "source_example": "newton/examples/robot/example_robot_policy.py", + "asset_revision": _ASSET_REVISION, + "model_format": "torchscript", + "observation_size": 48, + "action_size": 12, + }, + ) + + +class AnymalCVelocityAdapter: + """Reproduce the upstream command locomotion observation and action path.""" + + command_enabled = True + command_step = (0.1, 0.05, 0.1) + command_limits = (1.0, 0.5, 1.0) + + def __init__(self, request: AdapterRequest) -> None: + config = dict(request.config) + self.device = torch.device(str(config["inference_device"])) + self.joint_names = tuple(config["joint_names"]) + if self.joint_names != _JOINT_NAMES: + raise ValueError("ANYmal-C joint order is incompatible") + self.checkpoint = request.models["actor"] + self.previous_action = np.zeros(12, dtype=np.float32) + self.joints: JointMap | None = None + self.model: torch.jit.ScriptModule | None = None + + def setup(self, context: PolicyContext) -> None: + """Load the model and bind the runtime joint order.""" + if context.robot is None: + raise RuntimeError("ANYmal-C robot description is required") + self.joints = JointMap.from_joint_names( + context.robot.joint_names, + self.joint_names, + ) + self.model = torch.jit.load( + str(self.checkpoint), + map_location=self.device, + ).eval() + with torch.inference_mode(): + output = self.model(torch.zeros((1, 48), device=self.device)) + if not isinstance(output, torch.Tensor) or tuple(output.shape) != (1, 12): + shape = ( + None if not isinstance(output, torch.Tensor) else tuple(output.shape) + ) + raise ValueError(f"ANYmal-C policy output must be (1, 12), got {shape}") + + def reset(self, frame: EvaluationFrame) -> None: + """Reset the previous action used by the policy observation.""" + self.previous_action.fill(0.0) + + def infer(self, frame: EvaluationFrame) -> PolicyOutput: + """Build one 48-D observation and return 12 joint targets.""" + if frame.robot_state is None: + raise RuntimeError("ANYmal-C robot state is required") + observation = self._build_observation(frame) + tensor = torch.from_numpy(observation).to(self.device).unsqueeze(0) + with torch.inference_mode(): + output = self._model()(tensor) + if not isinstance(output, torch.Tensor) or tuple(output.shape) != (1, 12): + shape = ( + None if not isinstance(output, torch.Tensor) else tuple(output.shape) + ) + raise ValueError(f"ANYmal-C policy output must be (1, 12), got {shape}") + action = require_finite( + "ANYmal-C action", + output[0].detach().cpu().numpy(), + ) + self.previous_action = action.copy() + return PolicyOutput( + action=self._joints().command( + position=_DEFAULT_POSITION + 0.5 * action, + ), + termination_reason=_fall_reason(frame.robot_state.root_pose), + ) + + def metrics(self) -> dict[str, float]: + """Return the metrics produced by this velocity example.""" + return {} + + def close(self) -> None: + """Release the loaded TorchScript model.""" + self.model = None + + def _build_observation(self, frame: EvaluationFrame) -> np.ndarray: + state = frame.robot_state + if state is None: + raise RuntimeError("ANYmal-C robot state is required") + pose = np.asarray(state.root_pose, dtype=np.float32) + velocity = np.asarray(state.root_velocity, dtype=np.float32) + rotation = pose[:3, :3] + qpos = self._joints().to_model(state.qpos) + qvel = self._joints().to_model(state.qvel) + command = require_finite( + "ANYmal-C command", + frame.controls["command"], + ) + if command.shape != (3,): + raise ValueError("ANYmal-C command must contain vx, vy, and yaw rate") + observation = np.concatenate( + ( + rotation.T @ velocity[:3], + rotation.T @ velocity[3:], + rotation.T @ np.asarray((0.0, 0.0, -1.0), dtype=np.float32), + command, + qpos - _DEFAULT_POSITION, + qvel, + self.previous_action, + ), + dtype=np.float32, + ) + if observation.shape != (48,): + raise ValueError( + f"ANYmal-C observation must have shape (48,), got {observation.shape}" + ) + return require_finite("ANYmal-C observation", observation) + + def _joints(self) -> JointMap: + if self.joints is None: + raise RuntimeError("ANYmal-C Adapter is not set up") + return self.joints + + def _model(self) -> torch.jit.ScriptModule: + if self.model is None: + raise RuntimeError("ANYmal-C Adapter is not set up") + return self.model + + +def _fall_reason(root_pose: np.ndarray) -> str | None: + pose = np.asarray(root_pose, dtype=np.float64) + if pose.shape != (4, 4) or not np.all(np.isfinite(pose)): + raise ValueError("ANYmal-C root pose must be a finite 4x4 matrix") + height = float(pose[2, 3]) + tilt = math.acos(float(np.clip(pose[2, 2], -1.0, 1.0))) + reasons = [] + if height < 0.25: + reasons.append(f"base_height_below_minimum: {height:.3f} m") + if tilt > math.pi * 0.4: + reasons.append(f"bad_orientation: {tilt:.3f} rad") + return "; ".join(reasons) or None diff --git a/examples/learning/motion_policy_evaluation/eval_policy.py b/examples/learning/motion_policy_evaluation/eval_policy.py new file mode 100644 index 000000000..1f5f69929 --- /dev/null +++ b/examples/learning/motion_policy_evaluation/eval_policy.py @@ -0,0 +1,72 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Run the public ANYmal-C checkpoint directly from the example directory.""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +__all__ = ["example_arguments", "main"] + +EXAMPLE_ROOT = Path(__file__).resolve().parent +REPOSITORY_ROOT = EXAMPLE_ROOT.parents[2] +DEFAULT_CACHE = Path.home() / ".cache/embodichain/examples/anymal_c_velocity" + + +def example_arguments(argv: list[str]) -> list[str]: + """Add the example Profile, checkpoint, and resource paths. + + Args: + argv: Evaluation options accepted by ``eval-motion-policy``. + + Returns: + Arguments ready for the EmbodiChain evaluation CLI. + """ + cache = Path(os.environ.get("ANYMAL_C_EXAMPLE_CACHE", DEFAULT_CACHE)) + resource_root = cache / "upstream" + checkpoint = resource_root / "anybotics_anymal_c/rl_policies/mjw_anymal.pt" + return [ + "--profile", + "newton-anymal-c-velocity", + "--checkpoint", + str(checkpoint), + "--resource-root", + str(resource_root), + *argv, + ] + + +def main(argv: list[str] | None = None) -> None: + """Register the local Profile and run visual policy evaluation. + + Args: + argv: Evaluation options. Uses command-line arguments when omitted. + """ + if str(REPOSITORY_ROOT) not in sys.path: + sys.path.insert(0, str(REPOSITORY_ROOT)) + + from anymal_c import register + from embodichain.learning.rl.motion_policy_evaluation.cli import cli + + register() + cli(example_arguments(sys.argv[1:] if argv is None else argv)) + + +if __name__ == "__main__": + main() diff --git a/examples/learning/motion_policy_evaluation/prepare_resources.py b/examples/learning/motion_policy_evaluation/prepare_resources.py new file mode 100644 index 000000000..703b41938 --- /dev/null +++ b/examples/learning/motion_policy_evaluation/prepare_resources.py @@ -0,0 +1,243 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Fetch the pinned public policy and robot assets for the ANYmal-C example.""" + +from __future__ import annotations + +import argparse +import hashlib +import os +import subprocess +from pathlib import Path + +__all__ = ["main", "prepare_resources"] + +UPSTREAM_URL = "https://github.com/newton-physics/newton-assets.git" +UPSTREAM_REVISION = "261cd1f429619d8ef4f546bd788ab9dea906b5e1" +MODEL_RELATIVE_PATH = Path("anybotics_anymal_c/rl_policies/mjw_anymal.pt") +POLICY_CONFIG_RELATIVE_PATH = Path("anybotics_anymal_c/rl_policies/anymal.yaml") +POLICY_LICENSE_RELATIVE_PATH = Path("anybotics_anymal_c/rl_policies/LICENSE") +ROBOT_LICENSE_RELATIVE_PATH = Path("anybotics_anymal_c/LICENSE") +ROBOT_RELATIVE_PATH = Path("anybotics_anymal_c/urdf/anymal.urdf") +MESH_RELATIVE_PATH = Path("anybotics_anymal_c/meshes/base.dae") +MODEL_SHA256 = "00765c1c07e497be3825672b05f9cefff9238f2df72fb0bcb5ac9541155b945f" +POLICY_CONFIG_SHA256 = ( + "b5a463ac418c7f40ebe494c7bcf0d8031f021db70a0625dbcc28a718de8ee817" +) +POLICY_LICENSE_SHA256 = ( + "59899c6091b540582ed617e8eeaac4919dc985ccfc35459ee9752b699be5205b" +) +ROBOT_LICENSE_SHA256 = ( + "cef384faae108293b03b5e16a00bc3db8212d44575f69df6296438a3f901700b" +) +ROBOT_SHA256 = "d6bd20292cdd4873ffdeeb6f8ca3f96c4a0096565d78d8b6204f6edf0d19fb83" +MESH_SHA256 = "785bea9b33831f8c741fc0ca070162e73cbf560ea9b03c53abf8978be877fc48" + + +def prepare_resources(output: Path) -> tuple[Path, Path]: + """Fetch and verify the pinned upstream files. + + Args: + output: Cache directory that will contain the sparse Git checkout. + + Returns: + The local checkpoint and resource-root paths. + """ + output = output.expanduser().resolve() + checkout = output / "upstream" + _status("Preparing the ANYmal-C command policy and robot assets") + _prepare_checkout( + checkout, + UPSTREAM_URL, + UPSTREAM_REVISION, + ( + f"/{MODEL_RELATIVE_PATH}", + f"/{POLICY_CONFIG_RELATIVE_PATH}", + f"/{POLICY_LICENSE_RELATIVE_PATH}", + f"/{ROBOT_LICENSE_RELATIVE_PATH}", + "/anybotics_anymal_c/urdf/**", + "/anybotics_anymal_c/meshes/**", + ), + ) + + checkpoint = checkout / MODEL_RELATIVE_PATH + policy_config = checkout / POLICY_CONFIG_RELATIVE_PATH + policy_license = checkout / POLICY_LICENSE_RELATIVE_PATH + robot_license = checkout / ROBOT_LICENSE_RELATIVE_PATH + robot = checkout / ROBOT_RELATIVE_PATH + mesh = checkout / MESH_RELATIVE_PATH + if not checkpoint.is_file(): + raise FileNotFoundError(f"Downloaded checkpoint does not exist: {checkpoint}") + if not robot.is_file(): + raise FileNotFoundError(f"Downloaded robot asset does not exist: {robot}") + _verify_sha256(checkpoint, MODEL_SHA256) + _verify_sha256(policy_config, POLICY_CONFIG_SHA256) + _verify_sha256(policy_license, POLICY_LICENSE_SHA256) + _verify_sha256(robot_license, ROBOT_LICENSE_SHA256) + _verify_sha256(robot, ROBOT_SHA256) + _verify_sha256(mesh, MESH_SHA256) + _status("Resource verification completed") + return checkpoint, checkout + + +def main() -> None: + """Prepare resources and print the paths used by the evaluation command.""" + parser = argparse.ArgumentParser() + parser.add_argument( + "--output", + type=Path, + default=Path.home() / ".cache/embodichain/examples/anymal_c_velocity", + help="Directory used for the pinned upstream checkout", + ) + args = parser.parse_args() + checkpoint, resource_root = prepare_resources(args.output) + print(f"Checkpoint: {checkpoint}") + print(f"Resource root: {resource_root}") + + +def _git( + checkout: Path, + *args: str, + capture_output: bool = False, +) -> subprocess.CompletedProcess[str]: + command = ["git", "-C", str(checkout), *args] + environment = os.environ.copy() + environment["GIT_TERMINAL_PROMPT"] = "0" + try: + return subprocess.run( + command, + check=True, + text=True, + capture_output=capture_output, + env=environment, + timeout=600, + ) + except subprocess.TimeoutExpired as error: + raise RuntimeError( + f"Git command did not finish within 10 minutes: {' '.join(command)}" + ) from error + + +def _git_output(checkout: Path, *args: str) -> str | None: + try: + return _git(checkout, *args, capture_output=True).stdout.strip() + except subprocess.CalledProcessError: + return None + + +def _prepare_checkout( + checkout: Path, + url: str, + revision: str, + includes: tuple[str, ...], +) -> None: + if checkout.exists() and not (checkout / ".git").is_dir(): + raise RuntimeError( + f"Resource path exists but is not a Git checkout: {checkout}" + ) + + if checkout.exists(): + actual_revision = _git_output(checkout, "rev-parse", "HEAD") + if actual_revision == revision: + tracked_changes = _git_output( + checkout, + "status", + "--porcelain", + "--untracked-files=no", + ) + if tracked_changes == "": + _status(f"Using cached revision {revision[:8]} from {checkout}") + return + _status(f"Repairing interrupted checkout in {checkout}") + elif actual_revision is not None: + raise RuntimeError( + f"Existing checkout uses revision {actual_revision}; " + "choose another --output" + ) + else: + _status(f"Resuming interrupted checkout in {checkout}") + + remote_url = _git_output(checkout, "remote", "get-url", "origin") + if remote_url != url: + raise RuntimeError( + f"Incomplete checkout uses an unexpected remote: {remote_url}" + ) + else: + checkout.parent.mkdir(parents=True, exist_ok=True) + checkout.mkdir() + _git(checkout, "init", "--quiet") + _git(checkout, "remote", "add", "origin", url) + + _git(checkout, "sparse-checkout", "init", "--no-cone") + _git(checkout, "sparse-checkout", "set", *includes) + + fetched_revision = _git_output(checkout, "rev-parse", "FETCH_HEAD") + if fetched_revision != revision: + _status(f"Fetching revision {revision[:8]} from {url}") + _git( + checkout, + "fetch", + "--progress", + "--filter=blob:none", + "--depth", + "1", + "origin", + revision, + ) + else: + _status(f"Using previously fetched revision {revision[:8]}") + + _status(f"Checking out required files in {checkout}") + _git( + checkout, + "-c", + "advice.detachedHead=false", + "checkout", + "--progress", + "--force", + "--detach", + "FETCH_HEAD", + ) + actual_revision = _git_output(checkout, "rev-parse", "HEAD") + if actual_revision != revision: + raise RuntimeError( + f"Checkout revision mismatch: expected {revision}, got {actual_revision}" + ) + + +def _status(message: str) -> None: + print(f"[resources] {message}", flush=True) + + +def _verify_sha256(path: Path, expected: str) -> None: + actual = _sha256(path) + if actual != expected: + raise RuntimeError( + f"SHA256 mismatch for {path}: expected {expected}, got {actual}" + ) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +if __name__ == "__main__": + main() diff --git a/tests/learning/rl/motion_policy_evaluation/__init__.py b/tests/learning/rl/motion_policy_evaluation/__init__.py new file mode 100644 index 000000000..82062a260 --- /dev/null +++ b/tests/learning/rl/motion_policy_evaluation/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Tests for EmbodiChain motion-policy evaluation.""" + +from __future__ import annotations diff --git a/tests/learning/rl/motion_policy_evaluation/test_anymal_c_example.py b/tests/learning/rl/motion_policy_evaluation/test_anymal_c_example.py new file mode 100644 index 000000000..68965f64a --- /dev/null +++ b/tests/learning/rl/motion_policy_evaluation/test_anymal_c_example.py @@ -0,0 +1,255 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import importlib.util +import subprocess +from pathlib import Path + +import numpy as np +import torch +from dexsim.kit.motion_policy import ( + AdapterRequest, + EvaluationFrame, + PolicyContext, + RobotDescription, + RobotState, + parse_policy_spec, +) + +from embodichain.learning.rl.motion_policy_evaluation import MotionProfileRequest + +_JOINT_NAMES = ( + "LF_HAA", + "LF_HFE", + "LF_KFE", + "LH_HAA", + "LH_HFE", + "LH_KFE", + "RF_HAA", + "RF_HFE", + "RF_KFE", + "RH_HAA", + "RH_HFE", + "RH_KFE", +) +_DEFAULT_POSITION = np.asarray( + (0.0, 0.4, -0.8, 0.0, -0.4, 0.8, 0.0, 0.4, -0.8, 0.0, -0.4, 0.8), + dtype=np.float32, +) +_COMMAND = np.asarray((0.4, -0.2, 0.6), dtype=np.float32) + + +class _CommandPolicy(torch.nn.Module): + def forward(self, observation: torch.Tensor) -> torch.Tensor: + padding = torch.zeros( + (observation.shape[0], 9), + dtype=observation.dtype, + device=observation.device, + ) + return torch.cat((observation[:, 9:12], padding), dim=1) + + +def _run_git(directory: Path, *args: str) -> str: + result = subprocess.run( + ["git", "-C", str(directory), *args], + check=True, + text=True, + capture_output=True, + ) + return result.stdout.strip() + + +def test_resource_preparation_resumes_interrupted_checkout(tmp_path, capsys): + script = ( + Path(__file__).resolve().parents[4] + / "examples/learning/motion_policy_evaluation/prepare_resources.py" + ) + spec = importlib.util.spec_from_file_location("anymal_c_prepare_resources", script) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + upstream = tmp_path / "upstream" + upstream.mkdir() + _run_git(upstream, "init", "--quiet") + resource = upstream / "resources/robot.urdf" + resource.parent.mkdir() + resource.write_text("", encoding="utf-8") + _run_git(upstream, "add", ".") + _run_git( + upstream, + "-c", + "user.name=Test User", + "-c", + "user.email=test@example.com", + "commit", + "--quiet", + "-m", + "add resource", + ) + revision = _run_git(upstream, "rev-parse", "HEAD") + + checkout = tmp_path / "cache/upstream" + checkout.mkdir(parents=True) + _run_git(checkout, "init", "--quiet") + _run_git(checkout, "remote", "add", "origin", str(upstream)) + _run_git(checkout, "sparse-checkout", "init", "--no-cone") + _run_git(checkout, "sparse-checkout", "set", "/resources/**") + _run_git(checkout, "fetch", "--quiet", "origin", revision) + + module._prepare_checkout( + checkout, + str(upstream), + revision, + ("/resources/**",), + ) + + assert _run_git(checkout, "rev-parse", "HEAD") == revision + assert (checkout / "resources/robot.urdf").is_file() + assert "Resuming interrupted checkout" in capsys.readouterr().out + + module._prepare_checkout( + checkout, + str(upstream), + revision, + ("/resources/**",), + ) + assert "Using cached revision" in capsys.readouterr().out + + (checkout / "resources/robot.urdf").unlink() + module._prepare_checkout( + checkout, + str(upstream), + revision, + ("/resources/**",), + ) + assert (checkout / "resources/robot.urdf").is_file() + assert "Repairing interrupted checkout" in capsys.readouterr().out + + +def test_anymal_c_profile_builds_and_runs_torchscript(tmp_path, monkeypatch): + example_root = ( + Path(__file__).resolve().parents[4] + / "examples/learning/motion_policy_evaluation" + ) + monkeypatch.syspath_prepend(str(example_root)) + from anymal_c.profile import ( + AnymalCVelocityAdapter, + build_profile, + ) + + checkpoint = tmp_path / "mjw_anymal.pt" + traced = torch.jit.trace(_CommandPolicy().eval(), torch.zeros((1, 48))) + torch.jit.save(traced, checkpoint) + + robot_asset = tmp_path / "anybotics_anymal_c/urdf/anymal.urdf" + robot_asset.parent.mkdir(parents=True) + robot_asset.write_text("\n", encoding="utf-8") + + profile = build_profile( + MotionProfileRequest( + checkpoint=checkpoint, + device=torch.device("cpu"), + resource_root=tmp_path, + ) + ) + parse_policy_spec(profile.policy_spec) + config = profile.policy_spec["policy"]["adapter"]["config"] + adapter = AnymalCVelocityAdapter( + AdapterRequest( + asset_path=robot_asset, + models={"actor": checkpoint}, + resources={}, + config=config, + ) + ) + context = PolicyContext( + robot=RobotDescription( + _JOINT_NAMES, + ("base",), + "base", + ), + physics_dt=0.005, + sim_steps_per_control=4, + policy_dt=0.02, + ) + pose = np.eye(4, dtype=np.float32) + pose[2, 3] = 0.76 + frame = EvaluationFrame( + control_step=0, + policy_time=0.0, + simulation_time=0.0, + simulation_step=0, + robot_state=RobotState( + joint_names=_JOINT_NAMES, + qpos=_DEFAULT_POSITION.copy(), + qvel=np.zeros(12, dtype=np.float32), + target_qpos=_DEFAULT_POSITION.copy(), + target_qvel=np.zeros(12, dtype=np.float32), + joint_effort=np.zeros(12, dtype=np.float32), + root_name="base", + root_pose=pose, + root_velocity=np.zeros(6, dtype=np.float32), + link_names=("base",), + link_poses=pose[None, ...], + link_velocities=np.zeros((1, 6), dtype=np.float32), + ), + controls={"command": _COMMAND}, + ) + + adapter.setup(context) + adapter.reset(frame) + output = adapter.infer(frame) + + assert output.action.joint_names == _JOINT_NAMES + expected_position = _DEFAULT_POSITION.copy() + expected_position[:3] += 0.5 * _COMMAND + np.testing.assert_allclose(output.action.position, expected_position) + np.testing.assert_allclose( + adapter.previous_action, + np.concatenate((_COMMAND, np.zeros(9, dtype=np.float32))), + ) + assert adapter.command_enabled + assert adapter.command_limits == (1.0, 0.5, 1.0) + assert output.termination_reason is None + adapter.reset(frame) + np.testing.assert_array_equal(adapter.previous_action, np.zeros(12)) + adapter.close() + + +def test_example_script_supplies_default_resource_paths(tmp_path, monkeypatch): + example_root = ( + Path(__file__).resolve().parents[4] + / "examples/learning/motion_policy_evaluation" + ) + monkeypatch.syspath_prepend(str(example_root)) + monkeypatch.setenv("ANYMAL_C_EXAMPLE_CACHE", str(tmp_path)) + from eval_policy import example_arguments + + arguments = example_arguments(["--control-steps", "5"]) + + assert arguments == [ + "--profile", + "newton-anymal-c-velocity", + "--checkpoint", + str(tmp_path / "upstream/anybotics_anymal_c/rl_policies/mjw_anymal.pt"), + "--resource-root", + str(tmp_path / "upstream"), + "--control-steps", + "5", + ] diff --git a/tests/learning/rl/motion_policy_evaluation/test_bridge.py b/tests/learning/rl/motion_policy_evaluation/test_bridge.py new file mode 100644 index 000000000..754a8467e --- /dev/null +++ b/tests/learning/rl/motion_policy_evaluation/test_bridge.py @@ -0,0 +1,159 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +from embodichain.learning.rl.motion_policy_evaluation.bridge import ( + create_motion_profile_evaluator, + evaluate_motion_profile, +) +from embodichain.learning.rl.motion_policy_evaluation.profile import MotionProfile + + +def test_bridge_forwards_physics_backend_and_exact_control_steps( + tmp_path, + monkeypatch, +): + checkpoint = tmp_path / "policy.pt" + checkpoint.write_bytes(b"checkpoint") + profile = MotionProfile( + profile_id="example", + checkpoint=checkpoint, + policy_spec={"schema_version": 1}, + ) + options = [] + forwarded = [] + monkeypatch.setattr( + "embodichain.learning.rl.motion_policy_evaluation.bridge.parse_policy_spec", + lambda value: "parsed", + ) + monkeypatch.setattr( + "embodichain.learning.rl.motion_policy_evaluation.bridge.resolve_policy_spec", + lambda spec, resolver: "resolved", + ) + monkeypatch.setattr( + "embodichain.learning.rl.motion_policy_evaluation.bridge.policy_spec_to_dict", + lambda value: {"policy_id": "example"}, + ) + monkeypatch.setattr( + "embodichain.learning.rl.motion_policy_evaluation.bridge.scene_config_to_dict", + lambda value: {"style": "standard"}, + ) + + def run_policy(resolved, run_options, **kwargs): + options.append(run_options) + forwarded.append(kwargs) + return SimpleNamespace( + reason="control steps reached", + simulation_time=0.2, + simulation_steps=40, + control_steps=10, + physics_backend="default", + requested_duration=None, + effective_duration=0.2, + metrics={"tracking/error": 0.25}, + ) + + monkeypatch.setattr( + "embodichain.learning.rl.motion_policy_evaluation.bridge.run_motion_policy", + run_policy, + ) + + input_provider = object() + environment = object() + result = evaluate_motion_profile( + profile, + control_steps=10, + physics_backend="default", + input_provider=input_provider, + environment=environment, + ) + + assert options[0].control_steps == 10 + assert options[0].physics_backend == "default" + assert forwarded == [ + { + "environment": environment, + "input_provider": input_provider, + } + ] + assert result.episodes[0]["control_steps"] == 10 + assert result.episodes[0]["effective_duration"] == 0.2 + assert result.summary["metrics"]["tracking/error"] == 0.25 + + +def test_create_profile_evaluator_forwards_the_environment( + tmp_path, + monkeypatch, +): + checkpoint = tmp_path / "policy.pt" + checkpoint.write_bytes(b"checkpoint") + profile = MotionProfile( + profile_id="example", + checkpoint=checkpoint, + policy_spec={"schema_version": 1}, + ) + evaluator = object() + adapter = object() + environment = object() + calls = [] + monkeypatch.setattr( + "embodichain.learning.rl.motion_policy_evaluation.bridge.parse_policy_spec", + lambda value: "parsed", + ) + monkeypatch.setattr( + "embodichain.learning.rl.motion_policy_evaluation.bridge.resolve_policy_spec", + lambda spec, resolver: "resolved", + ) + monkeypatch.setattr( + "embodichain.learning.rl.motion_policy_evaluation.bridge.create_motion_policy_evaluator", + lambda resolved, options, **kwargs: calls.append((resolved, options, kwargs)) + or evaluator, + ) + + result = create_motion_profile_evaluator( + profile, + adapter=adapter, + environment=environment, + ) + + assert result is evaluator + assert calls == [ + ( + "resolved", + None, + {"adapter": adapter, "environment": environment}, + ), + ] + + +def test_prebuilt_environment_requires_one_episode(tmp_path): + checkpoint = tmp_path / "policy.pt" + checkpoint.write_bytes(b"checkpoint") + profile = MotionProfile( + profile_id="example", + checkpoint=checkpoint, + policy_spec={"schema_version": 1}, + ) + + try: + evaluate_motion_profile(profile, episodes=2, environment=object()) + except ValueError as error: + assert str(error) == "A prebuilt environment supports one episode" + else: + raise AssertionError("Expected multi-episode validation") diff --git a/tests/learning/rl/motion_policy_evaluation/test_checkpoint.py b/tests/learning/rl/motion_policy_evaluation/test_checkpoint.py new file mode 100644 index 000000000..19976302a --- /dev/null +++ b/tests/learning/rl/motion_policy_evaluation/test_checkpoint.py @@ -0,0 +1,30 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import torch + +from embodichain.learning.rl.motion_policy_evaluation import load_policy_state_dict + + +def test_load_policy_state_dict_returns_embodichain_policy_weights(tmp_path): + checkpoint = tmp_path / "policy.pt" + torch.save({"policy": {"actor.weight": torch.ones(2, 3)}}, checkpoint) + + state = load_policy_state_dict(checkpoint) + + assert torch.equal(state["actor.weight"], torch.ones(2, 3)) diff --git a/tests/learning/rl/motion_policy_evaluation/test_cli.py b/tests/learning/rl/motion_policy_evaluation/test_cli.py new file mode 100644 index 000000000..62d0e6b0e --- /dev/null +++ b/tests/learning/rl/motion_policy_evaluation/test_cli.py @@ -0,0 +1,143 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import importlib + +import pytest +import torch + +from embodichain.learning.rl.motion_policy_evaluation.cli import ( + _resolve_input, + _validate_native_options, + parse_args, +) +from embodichain.learning.rl.motion_policy_evaluation.manifest import ( + write_run_manifest, +) + + +def test_run_uses_manifest_profile_and_latest_checkpoint(tmp_path): + run = tmp_path / "run" + checkpoint = run / "checkpoints" / "policy.pt" + checkpoint.parent.mkdir(parents=True) + checkpoint.write_bytes(b"checkpoint") + train = tmp_path / "train.yaml" + train.write_text("trainer: {}\n", encoding="utf-8") + write_run_manifest( + run, + train_config=train, + latest_checkpoint=checkpoint, + motion_profile="example-motion", + ) + + resolved = _resolve_input(parse_args((str(run),))) + + assert resolved.profile == "example-motion" + assert resolved.checkpoint == checkpoint + assert resolved.selected_checkpoint == "latest" + + +def test_viewer_defaults_to_hybrid_without_a_time_limit(): + args = parse_args( + ( + "--profile", + "example-motion", + "--checkpoint", + "policy.pt", + "--viewer", + ) + ) + + assert args.renderer == "hybrid" + assert args.physics_backend is None + assert args.control_steps is None + assert args.duration is None + + +def test_run_without_profile_selects_native_task_evaluation(tmp_path, monkeypatch): + cli_module = importlib.import_module( + "embodichain.learning.rl.motion_policy_evaluation.cli" + ) + run = tmp_path / "run" + checkpoint = run / "checkpoints" / "policy.pt" + checkpoint.parent.mkdir(parents=True) + torch.save({"model_state_dict": {}}, checkpoint) + train = tmp_path / "train.yaml" + train.write_text("trainer: {}\n", encoding="utf-8") + write_run_manifest( + run, + train_config=train, + latest_checkpoint=checkpoint, + ) + expected = tmp_path / "evaluation.json" + received = {} + + def fake_native(args, resolved): + received["args"] = args + received["resolved"] = resolved + return expected + + monkeypatch.setattr(cli_module, "discover_task_packages", lambda: None) + monkeypatch.setattr(cli_module, "execute_init_hooks", lambda: None) + monkeypatch.setattr(cli_module, "_run_native_task", fake_native) + + report = cli_module.run(parse_args((str(run), "--viewer"))) + + assert report == expected + assert received["resolved"].profile is None + assert received["resolved"].checkpoint == checkpoint + assert received["args"].viewer is True + + +def test_explicit_native_checkpoint_requires_training_config(tmp_path): + checkpoint = tmp_path / "policy.pt" + checkpoint.write_bytes(b"checkpoint") + + args = parse_args(("--checkpoint", str(checkpoint))) + + with pytest.raises( + ValueError, + match="--config is required for a native EmbodiChain checkpoint", + ): + _resolve_input(args) + + +@pytest.mark.parametrize( + "arguments, option", + [ + (("--command", "0.5"), "--command"), + (("--physics-backend", "default"), "--physics-backend"), + (("--scene-config", "classic"), "--scene-config"), + (("--cache-dir", "cache"), "--cache-dir"), + (("--offline",), "--offline"), + (("--resource-root", "resources"), "--resource-root"), + ], +) +def test_native_task_rejects_profile_only_options(arguments, option): + args = parse_args( + ( + "--checkpoint", + "policy.pt", + "--config", + "train.yaml", + *arguments, + ) + ) + + with pytest.raises(ValueError, match=option): + _validate_native_options(args) diff --git a/tests/learning/rl/motion_policy_evaluation/test_manifest.py b/tests/learning/rl/motion_policy_evaluation/test_manifest.py new file mode 100644 index 000000000..812a6ac2c --- /dev/null +++ b/tests/learning/rl/motion_policy_evaluation/test_manifest.py @@ -0,0 +1,49 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from embodichain.learning.rl.motion_policy_evaluation import ( + RunManifest, + write_run_manifest, +) + + +def test_manifest_snapshots_configs_and_selects_latest_fallback(tmp_path): + run = tmp_path / "run" + checkpoint = run / "checkpoints" / "policy.pt" + checkpoint.parent.mkdir(parents=True) + checkpoint.write_bytes(b"checkpoint") + train = tmp_path / "train.yaml" + train.write_text("trainer: {}\n", encoding="utf-8") + gym = tmp_path / "gym.yaml" + gym.write_text("id: Example\n", encoding="utf-8") + + write_run_manifest( + run, + train_config=train, + gym_config=gym, + latest_checkpoint=checkpoint, + motion_profile="example-motion", + ) + manifest = RunManifest.load(run) + selected, path = manifest.select_checkpoint("best") + + assert manifest.motion_profile == "example-motion" + assert manifest.configs["train"] == run / "configs" / "train.yaml" + assert manifest.configs["gym"] == run / "configs" / "gym.yaml" + assert selected == "latest" + assert path == checkpoint diff --git a/tests/learning/rl/motion_policy_evaluation/test_native_task.py b/tests/learning/rl/motion_policy_evaluation/test_native_task.py new file mode 100644 index 000000000..e836b9093 --- /dev/null +++ b/tests/learning/rl/motion_policy_evaluation/test_native_task.py @@ -0,0 +1,251 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from tensordict import TensorDict + +from embodichain.learning.rl.evaluation import infer_policy_action +from embodichain.learning.rl.motion_policy_evaluation.native_task import ( + EmbodiChainTaskPolicyAdapter, + _policy_context_from_env, + evaluate_native_task, +) +from embodichain.learning.rl.runtime import PolicyRuntime +from dexsim.kit.motion_policy import EvaluationFrame, PolicyContext + + +class Policy(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor([[2.0], [-1.0]])) + + def get_action( + self, + tensordict: TensorDict, + deterministic: bool = False, + ) -> TensorDict: + assert deterministic + tensordict["action"] = tensordict["obs"] @ self.weight + return tensordict + + +class ActionManager: + def __init__(self) -> None: + self.calls = 0 + + def convert_policy_action_to_env_action(self, action: torch.Tensor): + self.calls += 1 + return action + 0.5 + + +class Environment: + num_envs = 1 + physics_dt = 0.005 + step_dt = 0.02 + + def __init__(self) -> None: + self.unwrapped = self + self.cfg = SimpleNamespace(sim_steps_per_control=4) + self.action_manager = ActionManager() + self.actions = [] + self.episode_step = 0 + self.reset_seeds = [] + self.closed = 0 + + def reset(self, seed=None): + self.reset_seeds.append(seed) + self.episode_step = 0 + return self._observation(), {} + + def step(self, action): + self.actions.append(action.clone()) + self.episode_step += 1 + done = self.episode_step == 2 + return ( + self._observation(), + torch.tensor([1.25]), + torch.tensor([done]), + torch.tensor([False]), + {"success": torch.tensor([done])}, + ) + + def close(self): + self.closed += 1 + + def _observation(self): + return { + "policy": torch.tensor( + [[float(self.episode_step), 1.0]], + dtype=torch.float32, + ) + } + + +class SimulatorEnvironment(Environment): + def __init__(self) -> None: + super().__init__() + self.sim = SimpleNamespace(get_world=lambda: None) + self.exit_process_values = [] + + def close(self, *, exit_process=None): + self.closed += 1 + self.exit_process_values.append(exit_process) + + +def test_native_adapter_uses_shared_deterministic_inference_chain(): + policy = Policy() + observation = {"policy": torch.tensor([[0.25, 0.75]])} + expected = infer_policy_action( + policy, + observation, + device="cpu", + num_envs=1, + ) + adapter = EmbodiChainTaskPolicyAdapter(policy, torch.device("cpu"), 1) + adapter.setup(PolicyContext(None, 0.005, 4, 0.02)) + + output = adapter.infer( + EvaluationFrame( + control_step=0, + policy_time=0.0, + simulation_step=0, + simulation_time=0.0, + observation=observation, + ) + ) + + assert torch.equal(output.action, expected) + adapter.close() + + +def test_lightweight_task_timing_uses_one_environment_step(): + context = _policy_context_from_env(SimpleNamespace(dt=0.125)) + + assert context.physics_dt == pytest.approx(0.125) + assert context.sim_steps_per_control == 1 + assert context.policy_dt == pytest.approx(0.125) + + +def test_native_task_runs_original_action_conversion_once_per_step(): + env = Environment() + runtime = PolicyRuntime( + env=env, + policy=Policy(), + device=torch.device("cpu"), + env_id="ExampleTask", + ) + + result = evaluate_native_task( + runtime, + seed=17, + viewer=False, + episodes=2, + control_steps=None, + duration=None, + ) + + assert result.reason == "episode target reached" + assert result.control_steps == 4 + assert result.simulation_steps == 16 + assert result.effective_duration == pytest.approx(0.08) + assert len(result.episodes) == 2 + assert result.metrics == pytest.approx( + { + "eval/avg_reward": 2.5, + "eval/avg_length": 2.0, + "eval/success_rate": 1.0, + } + ) + assert env.action_manager.calls == 4 + assert len(env.actions) == 4 + assert env.reset_seeds == [17, None] + assert env.closed == 1 + + +def test_native_task_control_limit_can_stop_before_episode_end(): + env = Environment() + runtime = PolicyRuntime( + env=env, + policy=Policy(), + device=torch.device("cpu"), + env_id="ExampleTask", + ) + + result = evaluate_native_task( + runtime, + seed=1, + viewer=False, + episodes=None, + control_steps=1, + duration=None, + ) + + assert result.reason == "control steps reached" + assert result.control_steps == 1 + assert result.episodes == () + + +def test_native_task_closes_simulator_without_exiting_process(): + env = SimulatorEnvironment() + runtime = PolicyRuntime( + env=env, + policy=Policy(), + device=torch.device("cpu"), + env_id="ExampleSimulatorTask", + ) + + result = evaluate_native_task( + runtime, + seed=1, + viewer=False, + episodes=None, + control_steps=1, + duration=None, + ) + + assert result.reason == "control steps reached" + assert env.closed == 1 + assert env.exit_process_values == [False] + + +def test_native_task_viewer_requires_a_simulator_environment(): + env = Environment() + runtime = PolicyRuntime( + env=env, + policy=Policy(), + device=torch.device("cpu"), + env_id="ExampleTask", + ) + + with pytest.raises( + ValueError, + match="requires an EmbodiChain simulator task", + ): + evaluate_native_task( + runtime, + seed=1, + viewer=True, + episodes=None, + control_steps=None, + duration=None, + ) + + assert env.closed == 1 diff --git a/tests/learning/rl/motion_policy_evaluation/test_profile.py b/tests/learning/rl/motion_policy_evaluation/test_profile.py new file mode 100644 index 000000000..ba2eac2ce --- /dev/null +++ b/tests/learning/rl/motion_policy_evaluation/test_profile.py @@ -0,0 +1,54 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import torch + +from embodichain.learning.rl.motion_policy_evaluation import ( + MotionProfile, + MotionProfileRequest, + build_motion_profile, + register_motion_profile, +) + + +def test_profile_provider_receives_checkpoint_and_training_configs(tmp_path): + checkpoint = tmp_path / "policy.pt" + checkpoint.write_bytes(b"checkpoint") + train = tmp_path / "train.yaml" + train.write_text("trainer: {}\n", encoding="utf-8") + requests = [] + + def provider(request): + requests.append(request) + return MotionProfile( + profile_id="test-profile-provider", + checkpoint=request.checkpoint, + policy_spec={"schema_version": 1}, + ) + + register_motion_profile("test-profile-provider", provider) + request = MotionProfileRequest( + checkpoint=checkpoint, + device=torch.device("cpu"), + configs={"train": train}, + ) + + profile = build_motion_profile("test-profile-provider", request) + + assert profile.checkpoint == checkpoint + assert requests[0].configs["train"] == train diff --git a/tests/learning/rl/motion_policy_evaluation/test_train_integration.py b/tests/learning/rl/motion_policy_evaluation/test_train_integration.py new file mode 100644 index 000000000..eaaf34da3 --- /dev/null +++ b/tests/learning/rl/motion_policy_evaluation/test_train_integration.py @@ -0,0 +1,43 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from embodichain.learning.rl.motion_policy_evaluation import RunManifest +from embodichain.learning.rl.train import _write_motion_run_manifest + + +def test_training_summary_writes_minimal_motion_manifest(tmp_path): + run = tmp_path / "run" + checkpoint = run / "checkpoints" / "policy.pt" + checkpoint.parent.mkdir(parents=True) + checkpoint.write_bytes(b"checkpoint") + config = tmp_path / "train.yaml" + config.write_text("trainer: {}\n", encoding="utf-8") + + _write_motion_run_manifest( + run, + config, + {"motion_profile": "example-motion"}, + { + "latest_checkpoint_path": str(checkpoint), + "best_checkpoint_path": None, + }, + ) + + manifest = RunManifest.load(run) + assert manifest.motion_profile == "example-motion" + assert manifest.checkpoints["latest"] == checkpoint diff --git a/tests/learning/test_runtime.py b/tests/learning/test_runtime.py new file mode 100644 index 000000000..3dd20cf9b --- /dev/null +++ b/tests/learning/test_runtime.py @@ -0,0 +1,166 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import gymnasium as gym +import torch + +from embodichain.learning.rl.runtime import ( + GymEnvironmentRuntime, + build_gym_environment, + build_gym_policy_runtime, + build_learning_policy_runtime, + resolve_config_reference, +) + + +def _policy_config() -> dict: + network = { + "type": "mlp", + "network_cfg": {"hidden_sizes": [8], "activation": "relu"}, + } + return { + "name": "actor_critic", + "actor": network, + "critic": network, + } + + +class LearningEnvironment: + single_observation_space = gym.spaces.Box(-1.0, 1.0, shape=(3,)) + single_action_space = gym.spaces.Box(-1.0, 1.0, shape=(2,)) + + def __init__(self) -> None: + self.closed = 0 + + def close(self) -> None: + self.closed += 1 + + +class GymEnvironment: + flattened_observation_space = gym.spaces.Box(-1.0, 1.0, shape=(3,)) + observation_space = gym.spaces.Dict( + {"policy": gym.spaces.Box(-1.0, 1.0, shape=(1, 3))} + ) + action_space = gym.spaces.Box(-1.0, 1.0, shape=(1, 2)) + + def __init__(self) -> None: + self.closed = 0 + self.action_manager = SimpleNamespace(total_action_dim=2) + + def reset(self): + return {"policy": torch.zeros(1, 3)}, {} + + def get_wrapper_attr(self, name): + return getattr(self, name) + + def close(self) -> None: + self.closed += 1 + + +def test_config_reference_uses_the_training_config_directory(tmp_path): + config_dir = tmp_path / "configs" + config_dir.mkdir() + gym_config = config_dir / "gym.yaml" + gym_config.write_text("id: Example\n", encoding="utf-8") + + assert resolve_config_reference("gym.yaml", base_dir=config_dir) == gym_config + + +def test_learning_runtime_builds_the_configured_policy(monkeypatch): + env = LearningEnvironment() + monkeypatch.setattr( + "embodichain.learning.rl.runtime.build_learning_environment", + lambda config, device, num_envs: ("PointMass", env), + ) + + runtime = build_learning_policy_runtime( + {"trainer": {"learning_env": "PointMass"}, "policy": _policy_config()}, + device=torch.device("cpu"), + num_envs=4, + ) + + assert runtime.env is env + assert runtime.env_id == "PointMass" + assert runtime.policy.actor[0].in_features == 3 + assert runtime.policy.actor[-1].out_features == 2 + + +def test_gym_environment_applies_runtime_overrides(tmp_path, monkeypatch): + gym_config = tmp_path / "gym.yaml" + gym_config.write_text("id: Example\n", encoding="utf-8") + env_cfg = SimpleNamespace( + num_envs=8, + sim_cfg=None, + profiler=None, + ) + built = GymEnvironment() + monkeypatch.setattr( + "embodichain.learning.rl.runtime.config_to_cfg", + lambda value, manager_modules: env_cfg, + ) + monkeypatch.setattr( + "embodichain.learning.rl.runtime.build_env", + lambda env_id, base_env_cfg: built, + ) + + runtime = build_gym_environment( + {"trainer": {"gym_config": "gym.yaml"}}, + simulation_device=torch.device("cpu"), + num_envs=1, + headless=True, + renderer="hybrid", + gpu_id=0, + config_dir=tmp_path, + ) + + assert runtime.env is built + assert runtime.env_id == "Example" + assert runtime.env_cfg.num_envs == 1 + assert runtime.env_cfg.sim_cfg.sim_device == torch.device("cpu") + assert runtime.env_cfg.sim_cfg.headless is True + assert runtime.env_cfg.sim_cfg.render_cfg.renderer == "hybrid" + + +def test_gym_runtime_uses_the_same_task_spaces_for_policy_build(monkeypatch): + env = GymEnvironment() + task = GymEnvironmentRuntime( + env=env, + env_id="Example", + env_cfg=SimpleNamespace(), + gym_config={"id": "Example"}, + gym_config_path=SimpleNamespace(resolve=lambda: None), + ) + monkeypatch.setattr( + "embodichain.learning.rl.runtime.build_gym_environment", + lambda *args, **kwargs: task, + ) + + runtime = build_gym_policy_runtime( + {"trainer": {"gym_config": "gym.yaml"}, "policy": _policy_config()}, + device=torch.device("cpu"), + num_envs=1, + headless=True, + renderer="hybrid", + gpu_id=0, + ) + + assert runtime.env is env + assert runtime.policy.actor[0].in_features == 3 + assert runtime.policy.actor[-1].out_features == 2 diff --git a/tests/test_main.py b/tests/test_main.py index 5466e7c11..906e00173 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -28,6 +28,7 @@ "benchmark", "data", "decompose-urdf", + "eval-motion-policy", "preview-asset", "preview_lerobot_data", "run-env", From a2a46766f2870b551be1b153cabf66b28d30065e Mon Sep 17 00:00:00 2001 From: acrlw <13927622+acrlw@users.noreply.github.com> Date: Thu, 13 Aug 2026 20:47:23 +0800 Subject: [PATCH 2/7] test(rl): cover paused viewer lifecycle --- .../test_native_task.py | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/tests/learning/rl/motion_policy_evaluation/test_native_task.py b/tests/learning/rl/motion_policy_evaluation/test_native_task.py index e836b9093..8a8ce6a5a 100644 --- a/tests/learning/rl/motion_policy_evaluation/test_native_task.py +++ b/tests/learning/rl/motion_policy_evaluation/test_native_task.py @@ -110,6 +110,44 @@ def close(self, *, exit_process=None): self.exit_process_values.append(exit_process) +class ViewerWindow: + def __init__(self) -> None: + self.titles = [] + + def set_window_title(self, title): + self.titles.append(title) + + def native(self): + return self + + def key_state(self, _key): + return False + + +class ViewerWorld: + def __init__(self) -> None: + self.window = ViewerWindow() + self.open = True + self.update_dts = [] + + def is_window_initialized(self): + return self.open + + def get_windows(self): + return self.window + + def update(self, dt): + self.update_dts.append(dt) + self.open = False + + +class ViewerSimulatorEnvironment(SimulatorEnvironment): + def __init__(self) -> None: + super().__init__() + self.world = ViewerWorld() + self.sim = SimpleNamespace(get_world=lambda: self.world) + + def test_native_adapter_uses_shared_deterministic_inference_chain(): policy = Policy() observation = {"policy": torch.tensor([[0.25, 0.75]])} @@ -226,6 +264,30 @@ def test_native_task_closes_simulator_without_exiting_process(): assert env.exit_process_values == [False] +def test_native_task_pause_waits_for_viewer_close_after_termination(): + env = ViewerSimulatorEnvironment() + runtime = PolicyRuntime( + env=env, + policy=Policy(), + device=torch.device("cpu"), + env_id="ExampleViewerTask", + ) + + result = evaluate_native_task( + runtime, + seed=1, + viewer=True, + episodes=None, + control_steps=None, + duration=None, + termination_behavior="pause", + ) + + assert result.reason == "viewer closed" + assert result.control_steps == 2 + assert env.world.update_dts == [0.0] + + def test_native_task_viewer_requires_a_simulator_environment(): env = Environment() runtime = PolicyRuntime( From 80c9915abbb5f036c561147773ab48c633a5ef40 Mon Sep 17 00:00:00 2001 From: acrlw <13927622+acrlw@users.noreply.github.com> Date: Fri, 14 Aug 2026 09:47:50 +0800 Subject: [PATCH 3/7] fix(rl): preserve task results and runtime lifecycle - Add --original-task to select the recorded EmbodiChain task. - Keep task metrics in the final evaluation result. - Close task resources when evaluator setup or execution fails. - Write the run manifest after training completes successfully. --- .../source/guides/motion_policy_evaluation.md | 5 ++ .../rl/motion_policy_evaluation/cli.py | 12 ++- .../motion_policy_evaluation/native_task.py | 89 ++++++++++++------- embodichain/learning/rl/train.py | 18 ++-- .../rl/motion_policy_evaluation/test_cli.py | 20 +++++ .../test_native_task.py | 69 +++++++++++++- 6 files changed, 169 insertions(+), 44 deletions(-) diff --git a/docs/source/guides/motion_policy_evaluation.md b/docs/source/guides/motion_policy_evaluation.md index 70ab271cf..877462930 100644 --- a/docs/source/guides/motion_policy_evaluation.md +++ b/docs/source/guides/motion_policy_evaluation.md @@ -54,6 +54,10 @@ best checkpoint and falls back to `latest` when the manifest value is `null`. to another checkpoint in the run is also accepted. The resolved checkpoint is recorded in `evaluation.json`. +When the manifest declares a `motion_profile`, the CLI uses that Profile by +default. Add `--original-task` to recreate the EmbodiChain Environment recorded +by the training configurations instead. + Open the Viewer for a training run: ```bash @@ -354,6 +358,7 @@ task Environment interface, Adapter API, task resources, and Python API. | Scenario | Command | |---|---| | New training run | `embodichain eval-motion-policy --checkpoint best --viewer` | +| Original task when a Profile is configured | `embodichain eval-motion-policy --original-task --viewer` | | Existing `.pt` | `embodichain eval-motion-policy --checkpoint --config --viewer` | | Headless smoke test | `embodichain eval-motion-policy --control-steps 20` | | External ANYmal-C `.pt` | `python examples/learning/motion_policy_evaluation/eval_policy.py --viewer` | diff --git a/embodichain/learning/rl/motion_policy_evaluation/cli.py b/embodichain/learning/rl/motion_policy_evaluation/cli.py index 0ef445f9f..91bd1360f 100644 --- a/embodichain/learning/rl/motion_policy_evaluation/cli.py +++ b/embodichain/learning/rl/motion_policy_evaluation/cli.py @@ -69,7 +69,13 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: description="Open a DexSim Viewer for an EmbodiChain policy checkpoint.", ) parser.add_argument("run", nargs="?", help="EmbodiChain training run directory.") - parser.add_argument("--profile", help="Registered Motion Profile name.") + mode = parser.add_mutually_exclusive_group() + mode.add_argument("--profile", help="Registered Motion Profile name.") + mode.add_argument( + "--original-task", + action="store_true", + help="Use the original EmbodiChain task instead of the manifest Profile.", + ) parser.add_argument( "--checkpoint", help="Checkpoint path, or best/latest when RUN is supplied.", @@ -179,7 +185,9 @@ def _resolve_input(args: argparse.Namespace) -> MotionInput: if candidate.is_absolute() else (manifest.root / candidate).resolve() ) - profile = args.profile or manifest.motion_profile + profile = ( + None if args.original_task else args.profile or manifest.motion_profile + ) return MotionInput( checkpoint=checkpoint, profile=profile, diff --git a/embodichain/learning/rl/motion_policy_evaluation/native_task.py b/embodichain/learning/rl/motion_policy_evaluation/native_task.py index f976b58c6..30fd66a6e 100644 --- a/embodichain/learning/rl/motion_policy_evaluation/native_task.py +++ b/embodichain/learning/rl/motion_policy_evaluation/native_task.py @@ -139,6 +139,7 @@ def __init__( self._episode_return = 0.0 self._episode_length = 0 self._episodes: list[dict[str, float | int | bool | str]] = [] + self._reported_metrics: dict[str, float] = {} self._closed = False self._previous_no_auto_reset = getattr( self._base_env, @@ -233,6 +234,7 @@ def step(self, action: object) -> EnvironmentStep: self._frame = self._make_frame(observation, task_state) reason = _termination_reason(info, terminated_value, truncated_value) metrics = _step_metrics(info, reward_value) + self._reported_metrics.update(metrics) if reason is not None: success = _info_bool(info, "success") self._episodes.append( @@ -255,20 +257,29 @@ def step(self, action: object) -> EnvironmentStep: ) def metrics(self) -> dict[str, float]: - """Aggregate completed task episodes.""" - if not self._episodes: - return {} - return { - "eval/avg_reward": float( - np.mean([float(episode["reward"]) for episode in self._episodes]) - ), - "eval/avg_length": float( - np.mean([float(episode["length"]) for episode in self._episodes]) - ), - "eval/success_rate": float( - np.mean([bool(episode["success"]) for episode in self._episodes]) - ), - } + """Return task metrics and completed episode aggregates.""" + result = dict(self._reported_metrics) + if self._episodes: + result.update( + { + "eval/avg_reward": float( + np.mean( + [float(episode["reward"]) for episode in self._episodes] + ) + ), + "eval/avg_length": float( + np.mean( + [float(episode["length"]) for episode in self._episodes] + ) + ), + "eval/success_rate": float( + np.mean( + [bool(episode["success"]) for episode in self._episodes] + ) + ), + } + ) + return result def wait_for_reset_or_close(self) -> str: """Keep a paused Viewer responsive until it is closed.""" @@ -349,25 +360,32 @@ def evaluate_native_task( except Exception: runtime.env.close() raise - adapter = EmbodiChainTaskPolicyAdapter(runtime.policy, runtime.device, num_envs=1) - if duration is not None: - control_steps = math.ceil( - duration / environment.policy_context.policy_dt - 1e-12 + adapter = None + evaluator = None + try: + adapter = EmbodiChainTaskPolicyAdapter( + runtime.policy, + runtime.device, + num_envs=1, + ) + if duration is not None: + control_steps = math.ceil( + duration / environment.policy_context.policy_dt - 1e-12 + ) + total_steps = 0 + reason = "viewer closed" if viewer else "episode target reached" + options = RunOptions( + headless=not viewer, + termination_behavior=( + "continue" if termination_behavior == "auto_reset" else "pause" + ), + ) + evaluator = create_motion_policy_evaluator( + options=options, + adapter=adapter, + environment=environment, + title=f"{runtime.env_id} - EmbodiChain", ) - total_steps = 0 - reason = "viewer closed" if viewer else "episode target reached" - options = RunOptions( - headless=not viewer, - termination_behavior=( - "continue" if termination_behavior == "auto_reset" else "pause" - ), - ) - with create_motion_policy_evaluator( - options=options, - adapter=adapter, - environment=environment, - title=f"{runtime.env_id} - EmbodiChain", - ) as evaluator: evaluator.reset() while True: if control_steps is not None and total_steps >= control_steps: @@ -394,6 +412,13 @@ def evaluate_native_task( metrics = environment.metrics() context = environment.policy_context backend = environment.physics_backend + finally: + if evaluator is None: + if adapter is not None: + adapter.close() + environment.close() + else: + evaluator.close() simulation_steps = total_steps * context.sim_steps_per_control return NativeTaskEvaluationResult( diff --git a/embodichain/learning/rl/train.py b/embodichain/learning/rl/train.py index 8b1da5441..a1b81620e 100644 --- a/embodichain/learning/rl/train.py +++ b/embodichain/learning/rl/train.py @@ -606,15 +606,15 @@ def train_from_config( if distributed and torch.distributed.is_initialized(): torch.distributed.destroy_process_group() - if rank == 0: - _write_motion_run_manifest( - run_base, - config_path, - trainer_cfg, - trainer.get_summary(), - gym_config=gym_config_path, - ) - logger.log_info("Training finished") + if rank == 0: + _write_motion_run_manifest( + run_base, + config_path, + trainer_cfg, + trainer.get_summary(), + gym_config=gym_config_path, + ) + logger.log_info("Training finished") def _write_motion_run_manifest( diff --git a/tests/learning/rl/motion_policy_evaluation/test_cli.py b/tests/learning/rl/motion_policy_evaluation/test_cli.py index 62d0e6b0e..b18c9d35f 100644 --- a/tests/learning/rl/motion_policy_evaluation/test_cli.py +++ b/tests/learning/rl/motion_policy_evaluation/test_cli.py @@ -52,6 +52,26 @@ def test_run_uses_manifest_profile_and_latest_checkpoint(tmp_path): assert resolved.selected_checkpoint == "latest" +def test_original_task_overrides_manifest_profile(tmp_path): + run = tmp_path / "run" + checkpoint = run / "checkpoints" / "policy.pt" + checkpoint.parent.mkdir(parents=True) + checkpoint.write_bytes(b"checkpoint") + train = tmp_path / "train.yaml" + train.write_text("trainer: {}\n", encoding="utf-8") + write_run_manifest( + run, + train_config=train, + latest_checkpoint=checkpoint, + motion_profile="example-motion", + ) + + resolved = _resolve_input(parse_args((str(run), "--original-task"))) + + assert resolved.profile is None + assert resolved.configs["train"] == run / "configs" / "train.yaml" + + def test_viewer_defaults_to_hybrid_without_a_time_limit(): args = parse_args( ( diff --git a/tests/learning/rl/motion_policy_evaluation/test_native_task.py b/tests/learning/rl/motion_policy_evaluation/test_native_task.py index 8a8ce6a5a..82596e504 100644 --- a/tests/learning/rl/motion_policy_evaluation/test_native_task.py +++ b/tests/learning/rl/motion_policy_evaluation/test_native_task.py @@ -84,7 +84,10 @@ def step(self, action): torch.tensor([1.25]), torch.tensor([done]), torch.tensor([False]), - {"success": torch.tensor([done])}, + { + "success": torch.tensor([done]), + "metrics": {"task_progress": torch.tensor([float(self.episode_step)])}, + }, ) def close(self): @@ -207,6 +210,8 @@ def test_native_task_runs_original_action_conversion_once_per_step(): assert len(result.episodes) == 2 assert result.metrics == pytest.approx( { + "reward": 1.25, + "task_progress": 2.0, "eval/avg_reward": 2.5, "eval/avg_length": 2.0, "eval/success_rate": 1.0, @@ -264,6 +269,68 @@ def test_native_task_closes_simulator_without_exiting_process(): assert env.exit_process_values == [False] +def test_native_task_closes_resources_when_evaluator_creation_fails(monkeypatch): + env = Environment() + policy = Policy() + runtime = PolicyRuntime( + env=env, + policy=policy, + device=torch.device("cpu"), + env_id="ExampleTask", + ) + + def fail_evaluator_creation(**_kwargs): + raise RuntimeError("evaluator setup failed") + + monkeypatch.setattr( + "embodichain.learning.rl.motion_policy_evaluation.native_task.create_motion_policy_evaluator", + fail_evaluator_creation, + ) + + with pytest.raises(RuntimeError, match="evaluator setup failed"): + evaluate_native_task( + runtime, + seed=1, + viewer=False, + episodes=1, + control_steps=None, + duration=None, + ) + + assert env.closed == 1 + assert policy.training is True + + +def test_native_task_closes_environment_when_adapter_creation_fails(monkeypatch): + env = Environment() + runtime = PolicyRuntime( + env=env, + policy=Policy(), + device=torch.device("cpu"), + env_id="ExampleTask", + ) + + def fail_adapter_creation(*_args, **_kwargs): + raise RuntimeError("adapter setup failed") + + monkeypatch.setattr( + "embodichain.learning.rl.motion_policy_evaluation.native_task.EmbodiChainTaskPolicyAdapter", + fail_adapter_creation, + ) + + with pytest.raises(RuntimeError, match="adapter setup failed"): + evaluate_native_task( + runtime, + seed=1, + viewer=False, + episodes=1, + control_steps=None, + duration=None, + ) + + assert env.closed == 1 + + def test_native_task_pause_waits_for_viewer_close_after_termination(): env = ViewerSimulatorEnvironment() runtime = PolicyRuntime( From f0d0cbe34ff93e44572e2598b3a0935a5c938b5a Mon Sep 17 00:00:00 2001 From: acrlw <13927622+acrlw@users.noreply.github.com> Date: Fri, 14 Aug 2026 10:12:28 +0800 Subject: [PATCH 4/7] fix(rl): apply config overrides and organize run outputs - Let explicit config paths override run manifest entries. - Store implicit recorder output under each run videos directory. - Build task policy context before applying Viewer reset behavior. - Document config precedence and the generated output layout. --- docs/source/guides/cli.md | 5 +-- .../source/guides/motion_policy_evaluation.md | 22 ++++++++++++ docs/source/tutorial/rl.rst | 6 ++-- .../rl/motion_policy_evaluation/cli.py | 16 ++++++--- .../motion_policy_evaluation/native_task.py | 8 +++-- embodichain/learning/rl/train.py | 20 +++++++++-- .../rl/basic/cart_pole/train_config.json | 5 ++- .../rl/basic/cart_pole/train_config.yaml | 1 - .../rl/basic/cart_pole/train_config_grpo.json | 5 ++- .../rl/basic/cart_pole/train_config_grpo.yaml | 1 - .../agents/rl/push_cube/train_config.json | 5 ++- .../rl/push_cube/train_config_grpo.json | 3 +- .../rl/motion_policy_evaluation/test_cli.py | 36 +++++++++++++++++++ .../test_train_integration.py | 26 +++++++++++++- 14 files changed, 133 insertions(+), 26 deletions(-) diff --git a/docs/source/guides/cli.md b/docs/source/guides/cli.md index 30eb51bd1..6d3c281c1 100644 --- a/docs/source/guides/cli.md +++ b/docs/source/guides/cli.md @@ -417,9 +417,10 @@ Motion Profile Headless run defaults to 100 control steps per episode. |---|---|---| | ``RUN`` | *(optional)* | Training run containing ``run-manifest.json`` | | ``--profile`` | manifest or EmbodiChain task | Registered external Motion Profile ID | +| ``--original-task`` | disabled | Use the EmbodiChain task configs when RUN declares a Motion Profile | | ``--checkpoint`` | ``best`` with RUN | ``best``, ``latest``, a path within RUN, or an explicit `.pt` path | -| ``--config`` | *(optional)* | Training config for an explicit checkpoint | -| ``--gym-config`` | *(optional)* | Task config for an explicit checkpoint | +| ``--config`` | manifest with RUN | Training config for an explicit checkpoint or a RUN override | +| ``--gym-config`` | manifest with RUN | Task config for an explicit checkpoint or a RUN override | | ``--resource-root`` | *(optional)* | Task assets and local reference data passed to the Profile | | ``--episodes`` | EmbodiChain task Headless: ``1`` | Completed EmbodiChain task episodes, or independent Motion Profile runs | | ``--control-steps`` | Viewer continuous | Exact Policy action count | diff --git a/docs/source/guides/motion_policy_evaluation.md b/docs/source/guides/motion_policy_evaluation.md index 877462930..8527f118b 100644 --- a/docs/source/guides/motion_policy_evaluation.md +++ b/docs/source/guides/motion_policy_evaluation.md @@ -19,9 +19,23 @@ outputs/_/ ├── configs/ │ ├── train.yaml │ └── gym.yaml +├── logs/ +│ └── / +├── videos/ +│ ├── train/ +│ └── eval/ +├── evaluations/ +│ └── -motion-policy/evaluation.json └── run-manifest.json ``` +Configurations are captured once when training completes. Checkpoint filenames +carry their training step, recorder filenames carry their episode index, and +each visual evaluation receives its own timestamped directory. Camera recorder +events use `videos/train` or `videos/eval` when `save_path` is omitted; an +explicit `save_path` keeps the configured location. The `evaluations` directory +is created by the first evaluation run. + A simulator RL task manifest has the following structure: ```json @@ -58,6 +72,10 @@ When the manifest declares a `motion_profile`, the CLI uses that Profile by default. Add `--original-task` to recreate the EmbodiChain Environment recorded by the training configurations instead. +`--config` and `--gym-config` replace the corresponding manifest entries for +one command. This is useful when checking a checkpoint with an updated local +configuration while retaining the run's checkpoint selection and report path. + Open the Viewer for a training run: ```bash @@ -254,6 +272,10 @@ Each run writes an `evaluation.json` report containing: - termination reason, integer simulation steps, and control steps; - episode reward, episode length, success, and task metrics. +During native task evaluation, every `info["metrics"]` scalar is forwarded to +the Evaluator on that control step. The latest values are included in the final +report with the completed-episode reward, length, and success aggregates. + A training run writes reports under `/evaluations/`. An explicit checkpoint writes them under the `evaluations/` directory next to that checkpoint. `--output` selects another output directory. diff --git a/docs/source/tutorial/rl.rst b/docs/source/tutorial/rl.rst index f91216172..556eac8d9 100644 --- a/docs/source/tutorial/rl.rst +++ b/docs/source/tutorial/rl.rst @@ -265,6 +265,8 @@ All outputs are written to ``./outputs/_/``: - **logs/**: TensorBoard logs - **checkpoints/**: Model checkpoints - **configs/**: Training config and referenced gym config snapshots +- **videos/**: Training and evaluation recordings grouped by phase +- **evaluations/**: Timestamped visual evaluation reports - **run-manifest.json**: training configs and best/latest checkpoint index used by ``eval-motion-policy`` A simulator training run can be opened directly in its original task Environment: @@ -466,9 +468,9 @@ Best Practices - **Configuration**: Use JSON for all hyperparameters. This makes experiments reproducible and easy to track. -- **Logging**: Metrics are automatically logged to TensorBoard and Weights & Biases. Check ``outputs//logs/`` for TensorBoard logs. +- **Logging**: Metrics are automatically logged to TensorBoard and Weights & Biases. Check ``outputs/_/logs/`` for TensorBoard logs. -- **Checkpoints**: Regular checkpoints are saved to ``outputs//checkpoints/``. Use these to resume training or evaluate policies. +- **Checkpoints**: Regular checkpoints are saved to ``outputs/_/checkpoints/``. Use these to resume training or evaluate policies. See Also -------- diff --git a/embodichain/learning/rl/motion_policy_evaluation/cli.py b/embodichain/learning/rl/motion_policy_evaluation/cli.py index 91bd1360f..01038b195 100644 --- a/embodichain/learning/rl/motion_policy_evaluation/cli.py +++ b/embodichain/learning/rl/motion_policy_evaluation/cli.py @@ -185,13 +185,21 @@ def _resolve_input(args: argparse.Namespace) -> MotionInput: if candidate.is_absolute() else (manifest.root / candidate).resolve() ) - profile = ( - None if args.original_task else args.profile or manifest.motion_profile - ) + configs = dict(manifest.configs) + if args.config is not None: + configs["train"] = Path(args.config).expanduser().resolve() + if args.gym_config is not None: + configs["gym"] = Path(args.gym_config).expanduser().resolve() + if args.original_task: + profile = None + elif args.profile is not None: + profile = args.profile + else: + profile = manifest.motion_profile return MotionInput( checkpoint=checkpoint, profile=profile, - configs=manifest.configs, + configs=configs, run=manifest.root, requested_checkpoint=requested, selected_checkpoint=selected, diff --git a/embodichain/learning/rl/motion_policy_evaluation/native_task.py b/embodichain/learning/rl/motion_policy_evaluation/native_task.py index 30fd66a6e..87151e765 100644 --- a/embodichain/learning/rl/motion_policy_evaluation/native_task.py +++ b/embodichain/learning/rl/motion_policy_evaluation/native_task.py @@ -141,13 +141,13 @@ def __init__( self._episodes: list[dict[str, float | int | bool | str]] = [] self._reported_metrics: dict[str, float] = {} self._closed = False + self._policy_context = _policy_context_from_env(self._base_env) self._previous_no_auto_reset = getattr( self._base_env, "_demo_no_auto_reset", _MISSING, ) self._base_env._demo_no_auto_reset = True - self._policy_context = _policy_context_from_env(self._base_env) @property def policy_context(self) -> PolicyContext: @@ -282,7 +282,11 @@ def metrics(self) -> dict[str, float]: return result def wait_for_reset_or_close(self) -> str: - """Keep a paused Viewer responsive until it is closed.""" + """Keep a paused Viewer responsive until it is closed. + + ``MotionPolicyEvaluator`` calls this method after a task termination + when the selected behavior is ``pause``. + """ while self.viewer_is_open: event = self.poll() if event is not None: diff --git a/embodichain/learning/rl/train.py b/embodichain/learning/rl/train.py index a1b81620e..cb93eeef0 100644 --- a/embodichain/learning/rl/train.py +++ b/embodichain/learning/rl/train.py @@ -58,6 +58,8 @@ from embodichain.utils.module_utils import find_function_from_modules from embodichain.lab.gym.envs.managers.cfg import EventCfg +_CAMERA_RECORDERS = {"record_camera_data", "record_camera_data_async"} + def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: """Parse command-line arguments. @@ -115,6 +117,20 @@ def _resolve_profile_output( return str(output.with_name(f"{output.stem}_rank{rank}{output.suffix}")) +def _event_params( + event_info: dict, + *, + run_base: str | Path, + phase: str, +) -> dict: + """Resolve event parameters that belong to one training run.""" + params = dict(event_info.get("params", {})) + function_name = str(event_info.get("func", "")).rsplit(".", 1)[-1] + if function_name in _CAMERA_RECORDERS: + params.setdefault("save_path", str(Path(run_base) / "videos" / phase)) + return params + + def _train_learning_env( cfg_data: dict, *, @@ -506,7 +522,7 @@ def train_from_config( for event_name, event_info in events_dict.get("train", {}).items(): event_func_str = event_info.get("func") mode = event_info.get("mode", "interval") - params = event_info.get("params", {}) + params = _event_params(event_info, run_base=run_base, phase="train") interval_step = event_info.get("interval_step", 1) event_func = find_function_from_modules( event_func_str, event_modules, raise_if_not_found=True @@ -522,7 +538,7 @@ def train_from_config( for event_name, event_info in events_dict.get("eval", {}).items(): event_func_str = event_info.get("func") mode = event_info.get("mode", "interval") - params = event_info.get("params", {}) + params = _event_params(event_info, run_base=run_base, phase="eval") interval_step = event_info.get("interval_step", 1) event_func = find_function_from_modules( event_func_str, event_modules, raise_if_not_found=True diff --git a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.json b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.json index 7a2deedb2..a0e4e7a18 100644 --- a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.json +++ b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.json @@ -45,8 +45,7 @@ 600, 320, 240 - ], - "save_path": "./outputs/videos/eval" + ] } } } @@ -93,4 +92,4 @@ "max_grad_norm": 0.5 } } -} \ No newline at end of file +} diff --git a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml index c160305e1..94304d23b 100644 --- a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml +++ b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config.yaml @@ -39,7 +39,6 @@ trainer: - 600 - 320 - 240 - save_path: ./outputs/videos/eval renderer: fast-rt policy: name: actor_critic diff --git a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.json b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.json index 80bedf8a5..506049fcf 100644 --- a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.json +++ b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.json @@ -46,8 +46,7 @@ 600, 320, 240 - ], - "save_path": "./outputs/videos/eval" + ] } } } @@ -87,4 +86,4 @@ "truncate_at_first_done": true } } -} \ No newline at end of file +} diff --git a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml index 2895a7b52..f7b77b93d 100644 --- a/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml +++ b/embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml @@ -40,7 +40,6 @@ trainer: - 600 - 320 - 240 - save_path: ./outputs/videos/eval renderer: hybrid policy: name: actor_only diff --git a/embodichain_tasks/configs/agents/rl/push_cube/train_config.json b/embodichain_tasks/configs/agents/rl/push_cube/train_config.json index 263498078..c0a195672 100644 --- a/embodichain_tasks/configs/agents/rl/push_cube/train_config.json +++ b/embodichain_tasks/configs/agents/rl/push_cube/train_config.json @@ -28,8 +28,7 @@ "eye": [-1.4, 1.4, 2.0], "target": [0, 0, 0], "up": [0, 0, 1], - "intrinsics": [600, 600, 320, 240], - "save_path": "./outputs/videos_ppo1/eval" + "intrinsics": [600, 600, 320, 240] } } } @@ -76,4 +75,4 @@ "max_grad_norm": 0.5 } } -} \ No newline at end of file +} diff --git a/embodichain_tasks/configs/agents/rl/push_cube/train_config_grpo.json b/embodichain_tasks/configs/agents/rl/push_cube/train_config_grpo.json index 4dea67c20..96d52e457 100644 --- a/embodichain_tasks/configs/agents/rl/push_cube/train_config_grpo.json +++ b/embodichain_tasks/configs/agents/rl/push_cube/train_config_grpo.json @@ -28,8 +28,7 @@ "eye": [-1.4, 1.4, 2.0], "target": [0, 0, 0], "up": [0, 0, 1], - "intrinsics": [600, 600, 320, 240], - "save_path": "./outputs/videos/eval" + "intrinsics": [600, 600, 320, 240] } } } diff --git a/tests/learning/rl/motion_policy_evaluation/test_cli.py b/tests/learning/rl/motion_policy_evaluation/test_cli.py index b18c9d35f..676da67ff 100644 --- a/tests/learning/rl/motion_policy_evaluation/test_cli.py +++ b/tests/learning/rl/motion_policy_evaluation/test_cli.py @@ -72,6 +72,42 @@ def test_original_task_overrides_manifest_profile(tmp_path): assert resolved.configs["train"] == run / "configs" / "train.yaml" +def test_explicit_configs_override_manifest_configs(tmp_path): + run = tmp_path / "run" + checkpoint = run / "checkpoints" / "policy.pt" + checkpoint.parent.mkdir(parents=True) + checkpoint.write_bytes(b"checkpoint") + train = tmp_path / "train.yaml" + train.write_text("trainer: {}\n", encoding="utf-8") + gym = tmp_path / "gym.yaml" + gym.write_text("id: Original\n", encoding="utf-8") + write_run_manifest( + run, + train_config=train, + gym_config=gym, + latest_checkpoint=checkpoint, + ) + replacement_train = tmp_path / "replacement-train.yaml" + replacement_train.write_text("trainer: {}\n", encoding="utf-8") + replacement_gym = tmp_path / "replacement-gym.yaml" + replacement_gym.write_text("id: Replacement\n", encoding="utf-8") + + resolved = _resolve_input( + parse_args( + ( + str(run), + "--config", + str(replacement_train), + "--gym-config", + str(replacement_gym), + ) + ) + ) + + assert resolved.configs["train"] == replacement_train + assert resolved.configs["gym"] == replacement_gym + + def test_viewer_defaults_to_hybrid_without_a_time_limit(): args = parse_args( ( diff --git a/tests/learning/rl/motion_policy_evaluation/test_train_integration.py b/tests/learning/rl/motion_policy_evaluation/test_train_integration.py index eaaf34da3..167f7a031 100644 --- a/tests/learning/rl/motion_policy_evaluation/test_train_integration.py +++ b/tests/learning/rl/motion_policy_evaluation/test_train_integration.py @@ -17,7 +17,7 @@ from __future__ import annotations from embodichain.learning.rl.motion_policy_evaluation import RunManifest -from embodichain.learning.rl.train import _write_motion_run_manifest +from embodichain.learning.rl.train import _event_params, _write_motion_run_manifest def test_training_summary_writes_minimal_motion_manifest(tmp_path): @@ -41,3 +41,27 @@ def test_training_summary_writes_minimal_motion_manifest(tmp_path): manifest = RunManifest.load(run) assert manifest.motion_profile == "example-motion" assert manifest.checkpoints["latest"] == checkpoint + + +def test_camera_recorder_defaults_to_the_run_video_directory(tmp_path): + params = _event_params( + {"func": "record_camera_data_async", "params": {"name": "main"}}, + run_base=tmp_path / "run", + phase="eval", + ) + + assert params["save_path"] == str(tmp_path / "run" / "videos" / "eval") + + +def test_camera_recorder_keeps_an_explicit_output_directory(tmp_path): + custom = tmp_path / "custom-videos" + params = _event_params( + { + "func": "record_camera_data", + "params": {"save_path": str(custom)}, + }, + run_base=tmp_path / "run", + phase="train", + ) + + assert params["save_path"] == str(custom) From 6dd4f5868c02ad24ad71477f5d405637b6a957ce Mon Sep 17 00:00:00 2001 From: acrlw <13927622+acrlw@users.noreply.github.com> Date: Mon, 17 Aug 2026 12:16:20 +0800 Subject: [PATCH 5/7] docs(rl): document motion policy viewer controls --- docs/source/guides/motion_policy_evaluation.md | 12 +++++++++--- examples/learning/motion_policy_evaluation/README.md | 9 ++++++++- .../test_anymal_c_example.py | 3 ++- 3 files changed, 19 insertions(+), 5 deletions(-) diff --git a/docs/source/guides/motion_policy_evaluation.md b/docs/source/guides/motion_policy_evaluation.md index 8527f118b..904ef1fea 100644 --- a/docs/source/guides/motion_policy_evaluation.md +++ b/docs/source/guides/motion_policy_evaluation.md @@ -87,8 +87,8 @@ embodichain eval-motion-policy outputs/_ \ ``` The Viewer runs until its window closes when neither `--control-steps` nor -`--duration` is set. A Headless run uses an exact integer number of control -steps: +`--duration` is set. Press `R` to start or stop Viewer recording. A Headless +run uses an exact integer number of control steps: ```bash embodichain eval-motion-policy outputs/_ \ @@ -336,7 +336,13 @@ such as `--device`, `--sim-device`, `--control-steps`, and `--scene-config` are forwarded to EmbodiChain. Use W/S to adjust `vx`, A/D to adjust `vy`, Q/E to adjust `yaw`, M to zero all -three commands, and R to reset the task. +three commands, Backspace to reset the task and camera framing, and T to switch +between tracking and free view. The Viewer follows the robot root in the ground +plane. Hold the left mouse button to adjust the orbit angle and use the mouse +wheel to adjust the viewing distance. Right-button panning is locked while +tracking is active. The tracking target continues following the robot during an +orbit drag. Switching back to tracking centers the camera on the current robot +position. The external Policy data path is: diff --git a/examples/learning/motion_policy_evaluation/README.md b/examples/learning/motion_policy_evaluation/README.md index 5dc227685..3e2fe1a97 100644 --- a/examples/learning/motion_policy_evaluation/README.md +++ b/examples/learning/motion_policy_evaluation/README.md @@ -61,7 +61,14 @@ Viewer controls: | A / D | Increase / decrease `vy` | | Q / E | Increase / decrease `yaw` | | M | Set all three commands to zero | -| R | Reset the robot and Policy history | +| Backspace | Reset the robot, Policy history, and camera framing | +| T | Switch between tracking and free view | + +The camera follows the robot root in the ground plane. Hold the left mouse +button to change the orbit angle and use the mouse wheel to change the viewing +distance. Right-button panning is locked while tracking is active. Tracking +continues while the orbit angle is being adjusted. Switching back to tracking +centers the camera on the current robot position. The terminal prints the path to `evaluation.json` when the Viewer closes. Run a Headless smoke test with: diff --git a/tests/learning/rl/motion_policy_evaluation/test_anymal_c_example.py b/tests/learning/rl/motion_policy_evaluation/test_anymal_c_example.py index 68965f64a..155ec9aed 100644 --- a/tests/learning/rl/motion_policy_evaluation/test_anymal_c_example.py +++ b/tests/learning/rl/motion_policy_evaluation/test_anymal_c_example.py @@ -168,7 +168,8 @@ def test_anymal_c_profile_builds_and_runs_torchscript(tmp_path, monkeypatch): resource_root=tmp_path, ) ) - parse_policy_spec(profile.policy_spec) + spec = parse_policy_spec(profile.policy_spec) + assert spec.environment.entrypoint is None config = profile.policy_spec["policy"]["adapter"]["config"] adapter = AnymalCVelocityAdapter( AdapterRequest( From 4899e7f0a9f15bd89e98d95deebadfb1645d32d1 Mon Sep 17 00:00:00 2001 From: acrlw <13927622+acrlw@users.noreply.github.com> Date: Tue, 18 Aug 2026 17:58:35 +0800 Subject: [PATCH 6/7] refactor(rl): unify post-training policy evaluation Add one eval-policy command for Headless, Viewer, and external Profile workflows. Reuse the training runtime, record run manifests, and consolidate the focused documentation and tests. --- ...n.learning.rl.motion_policy_evaluation.rst | 15 - ...odichain.learning.rl.policy_evaluation.rst | 7 + .../embodichain/embodichain.learning.rl.rst | 3 +- .../embodichain.learning.rl.runtime.rst | 29 -- docs/source/guides/cli.md | 73 +-- docs/source/guides/index.rst | 2 +- .../source/guides/motion_policy_evaluation.md | 392 -------------- docs/source/guides/policy_evaluation.md | 196 +++++++ docs/source/tutorial/rl.rst | 25 +- embodichain/__main__.py | 6 +- .../rl/motion_policy_evaluation/__init__.py | 50 -- .../rl/motion_policy_evaluation/checkpoint.py | 60 --- .../rl/policy_evaluation}/__init__.py | 16 +- .../bridge.py | 64 +-- .../cli.py | 481 +++++++++++------- .../manifest.py | 22 +- .../profile.py | 28 +- .../report.py | 4 +- .../viewer.py} | 136 ++--- embodichain/learning/rl/runtime.py | 168 +++--- embodichain/learning/rl/train.py | 60 +-- .../README.md | 14 +- .../anymal_c/__init__.py | 4 +- .../anymal_c/profile.py | 27 +- .../eval_policy.py | 6 +- .../prepare_resources.py | 92 ++-- .../motion_policy_evaluation/test_bridge.py | 159 ------ .../test_checkpoint.py | 30 -- .../rl/motion_policy_evaluation/test_cli.py | 199 -------- .../motion_policy_evaluation/test_manifest.py | 49 -- .../test_native_task.py | 380 -------------- .../motion_policy_evaluation/test_profile.py | 54 -- .../test_train_integration.py | 67 --- .../test_anymal_c_example.py | 109 +--- .../rl/policy_evaluation/test_bridge.py | 84 +++ .../learning/rl/policy_evaluation/test_cli.py | 123 +++++ .../rl/policy_evaluation/test_viewer.py | 221 ++++++++ tests/learning/test_point_mass.py | 8 + tests/learning/test_runtime.py | 51 +- tests/learning/test_train_profile.py | 11 + tests/test_main.py | 2 +- 41 files changed, 1224 insertions(+), 2303 deletions(-) delete mode 100644 docs/source/api_reference/embodichain/embodichain.learning.rl.motion_policy_evaluation.rst create mode 100644 docs/source/api_reference/embodichain/embodichain.learning.rl.policy_evaluation.rst delete mode 100644 docs/source/api_reference/embodichain/embodichain.learning.rl.runtime.rst delete mode 100644 docs/source/guides/motion_policy_evaluation.md create mode 100644 docs/source/guides/policy_evaluation.md delete mode 100644 embodichain/learning/rl/motion_policy_evaluation/__init__.py delete mode 100644 embodichain/learning/rl/motion_policy_evaluation/checkpoint.py rename {tests/learning/rl/motion_policy_evaluation => embodichain/learning/rl/policy_evaluation}/__init__.py (71%) rename embodichain/learning/rl/{motion_policy_evaluation => policy_evaluation}/bridge.py (71%) rename embodichain/learning/rl/{motion_policy_evaluation => policy_evaluation}/cli.py (54%) rename embodichain/learning/rl/{motion_policy_evaluation => policy_evaluation}/manifest.py (88%) rename embodichain/learning/rl/{motion_policy_evaluation => policy_evaluation}/profile.py (81%) rename embodichain/learning/rl/{motion_policy_evaluation => policy_evaluation}/report.py (95%) rename embodichain/learning/rl/{motion_policy_evaluation/native_task.py => policy_evaluation/viewer.py} (83%) rename examples/learning/{motion_policy_evaluation => policy_evaluation}/README.md (92%) rename examples/learning/{motion_policy_evaluation => policy_evaluation}/anymal_c/__init__.py (87%) rename examples/learning/{motion_policy_evaluation => policy_evaluation}/anymal_c/profile.py (88%) rename examples/learning/{motion_policy_evaluation => policy_evaluation}/eval_policy.py (90%) rename examples/learning/{motion_policy_evaluation => policy_evaluation}/prepare_resources.py (69%) delete mode 100644 tests/learning/rl/motion_policy_evaluation/test_bridge.py delete mode 100644 tests/learning/rl/motion_policy_evaluation/test_checkpoint.py delete mode 100644 tests/learning/rl/motion_policy_evaluation/test_cli.py delete mode 100644 tests/learning/rl/motion_policy_evaluation/test_manifest.py delete mode 100644 tests/learning/rl/motion_policy_evaluation/test_native_task.py delete mode 100644 tests/learning/rl/motion_policy_evaluation/test_profile.py delete mode 100644 tests/learning/rl/motion_policy_evaluation/test_train_integration.py rename tests/learning/rl/{motion_policy_evaluation => policy_evaluation}/test_anymal_c_example.py (63%) create mode 100644 tests/learning/rl/policy_evaluation/test_bridge.py create mode 100644 tests/learning/rl/policy_evaluation/test_cli.py create mode 100644 tests/learning/rl/policy_evaluation/test_viewer.py diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.motion_policy_evaluation.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.motion_policy_evaluation.rst deleted file mode 100644 index 03d6beaf2..000000000 --- a/docs/source/api_reference/embodichain/embodichain.learning.rl.motion_policy_evaluation.rst +++ /dev/null @@ -1,15 +0,0 @@ -embodichain.learning.rl.motion_policy_evaluation -================================================ - -.. automodule:: embodichain.learning.rl.motion_policy_evaluation - :members: - :undoc-members: - :show-inheritance: - -Native Task ------------ - -.. automodule:: embodichain.learning.rl.motion_policy_evaluation.native_task - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.policy_evaluation.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.policy_evaluation.rst new file mode 100644 index 000000000..b5aab6c0a --- /dev/null +++ b/docs/source/api_reference/embodichain/embodichain.learning.rl.policy_evaluation.rst @@ -0,0 +1,7 @@ +embodichain.learning.rl.policy_evaluation +========================================= + +.. automodule:: embodichain.learning.rl.policy_evaluation + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.rst index 1393f7505..e2e51b69d 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rl.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rl.rst @@ -18,8 +18,7 @@ collection logic, policy/model builders, and training entry points. buffer collector models - motion_policy_evaluation - runtime + policy_evaluation train utils diff --git a/docs/source/api_reference/embodichain/embodichain.learning.rl.runtime.rst b/docs/source/api_reference/embodichain/embodichain.learning.rl.runtime.rst deleted file mode 100644 index 10d68dbef..000000000 --- a/docs/source/api_reference/embodichain/embodichain.learning.rl.runtime.rst +++ /dev/null @@ -1,29 +0,0 @@ -embodichain.learning.rl.runtime -=============================== - -.. automodule:: embodichain.learning.rl.runtime - - .. rubric:: Module Attributes - - .. autosummary:: - - GymPolicyFactory - - .. rubric:: Functions - - .. autosummary:: - - build_gym_environment - build_gym_policy_runtime - build_learning_environment - build_learning_policy_runtime - resolve_config_reference - resolve_torch_device - seed_policy_runtime - - .. rubric:: Classes - - .. autosummary:: - - GymEnvironmentRuntime - PolicyRuntime diff --git a/docs/source/guides/cli.md b/docs/source/guides/cli.md index 6d3c281c1..427b8c17e 100644 --- a/docs/source/guides/cli.md +++ b/docs/source/guides/cli.md @@ -372,74 +372,53 @@ See the Profiling section under Run Env for report format. Outputs are written t --- -## Motion Policy Evaluation +## Policy Evaluation -Run visual evaluation for a simulator training run. EmbodiChain restores the -training Environment, observation, Policy, action manager, and task reset: +Evaluate the latest checkpoint from an EmbodiChain training run: ```bash -embodichain eval-motion-policy outputs/my_policy_ \ - --checkpoint best \ - --device cuda \ - --sim-device gpu \ - --viewer +embodichain eval-policy outputs/my_policy_ ``` -Evaluate an older EmbodiChain `.pt` checkpoint: +Open a simulator task in the Viewer: ```bash -embodichain eval-motion-policy \ - --checkpoint /path/to/policy.pt \ - --config /path/to/train.yaml \ - --gym-config /path/to/gym.yaml \ - --viewer +embodichain eval-policy outputs/my_policy_ \ + --checkpoint best \ + --viewer \ + --renderer hybrid ``` -Evaluate a checkpoint from an external Policy project with a registered Motion -Profile: +Evaluate an explicit EmbodiChain checkpoint: ```bash -embodichain eval-motion-policy \ - --profile my-joint-policy \ +embodichain eval-policy \ --checkpoint /path/to/policy.pt \ --config /path/to/train.yaml \ - --resource-root /path/to/task-resources \ - --viewer + --gym-config /path/to/gym.yaml ``` -Viewer runs continuously when `--episodes`, `--control-steps`, and `--duration` -are omitted. An EmbodiChain task Headless run completes one episode. An external -Motion Profile Headless run defaults to 100 control steps per episode. - -### Arguments +### Main arguments | Argument | Default | Description | |---|---|---| | ``RUN`` | *(optional)* | Training run containing ``run-manifest.json`` | -| ``--profile`` | manifest or EmbodiChain task | Registered external Motion Profile ID | -| ``--original-task`` | disabled | Use the EmbodiChain task configs when RUN declares a Motion Profile | -| ``--checkpoint`` | ``best`` with RUN | ``best``, ``latest``, a path within RUN, or an explicit `.pt` path | -| ``--config`` | manifest with RUN | Training config for an explicit checkpoint or a RUN override | -| ``--gym-config`` | manifest with RUN | Task config for an explicit checkpoint or a RUN override | -| ``--resource-root`` | *(optional)* | Task assets and local reference data passed to the Profile | -| ``--episodes`` | EmbodiChain task Headless: ``1`` | Completed EmbodiChain task episodes, or independent Motion Profile runs | -| ``--control-steps`` | Viewer continuous | Exact Policy action count | -| ``--duration`` | *(optional)* | Seconds converted upward to integer control steps | -| ``--command`` | Profile | External Motion Profile command values | -| ``--device`` | Training config or ``cpu`` | PyTorch inference device | -| ``--sim-device`` | Inference device or ``cpu`` | Simulation device | -| ``--physics-backend`` | Policy Spec | External Motion Profile backend override | -| ``--renderer`` | ``hybrid`` | ``raster``, ``hybrid``, ``fastrt``, or ``offlinert`` | -| ``--gpu-id`` | ``0`` | GPU index | -| ``--scene-config`` | Profile: ``standard`` | External Motion Profile scene style or YAML path | -| ``--termination-behavior`` | EmbodiChain task: ``auto_reset`` | EmbodiChain task: ``pause`` or ``auto_reset``; Profile also supports ``continue`` | -| ``--viewer`` | Headless | Open the DexSim Viewer | -| ``--cache-dir`` | DexSim cache | External Motion Profile resource cache | -| ``--offline`` | ``False`` | Resolve external Motion Profile resources from cache | +| ``--checkpoint`` | ``latest`` with RUN | ``latest``, ``best``, or a checkpoint path | +| ``--config`` | RUN manifest | Training configuration override | +| ``--gym-config`` | RUN manifest | Simulator task configuration override | +| ``--episodes`` | Training configuration | Number of completed task episodes | +| ``--num-envs`` | Training configuration | Number of parallel Headless environments | +| ``--viewer`` | Headless | Open the original simulator task in the DexSim Viewer | +| ``--control-steps`` | Viewer runs continuously | Exact number of Policy actions | +| ``--duration`` | *(optional)* | Duration converted to integer control steps | +| ``--renderer`` | Training configuration or ``hybrid`` | Viewer renderer | +| ``--device`` | Training configuration | PyTorch inference device | +| ``--sim-device`` | Inference device | Simulation device | | ``--output`` | RUN or checkpoint evaluations | Evaluation output parent directory | -See {doc}`motion_policy_evaluation` for EmbodiChain training runs, external -Motion Profiles, run manifests, reports, and complete examples. +External Motion Profiles use the same command with `--profile`. See +{doc}`policy_evaluation` for training-run layout, execution paths, Viewer +controls, output reports, and the complete ANYmal-C example. --- diff --git a/docs/source/guides/index.rst b/docs/source/guides/index.rst index 3c36b2137..2caee765a 100644 --- a/docs/source/guides/index.rst +++ b/docs/source/guides/index.rst @@ -13,5 +13,5 @@ Practical guides for common tasks in EmbodiChain. add_robot preview_asset run_env - motion_policy_evaluation + policy_evaluation cli diff --git a/docs/source/guides/motion_policy_evaluation.md b/docs/source/guides/motion_policy_evaluation.md deleted file mode 100644 index 904ef1fea..000000000 --- a/docs/source/guides/motion_policy_evaluation.md +++ /dev/null @@ -1,392 +0,0 @@ -# Visual Motion Policy Evaluation - -An EmbodiChain simulator RL task can open its trained `.pt` checkpoint in the -original task Environment and DexSim Viewer. Evaluation restores the Policy and -Environment from the training configuration, loads the checkpoint, and reuses -the task's observation, action processing, reset logic, objects, and termination -conditions. Lightweight learning tasks such as PointMass use the same command -for Headless evaluation. - -## Training runs - -`train-rl` writes a `[new] run-manifest.json` file to each output directory. The -manifest indexes the saved configurations and best/latest checkpoints: - -```text -outputs/_/ -├── checkpoints/ -│ └── policy_*.pt -├── configs/ -│ ├── train.yaml -│ └── gym.yaml -├── logs/ -│ └── / -├── videos/ -│ ├── train/ -│ └── eval/ -├── evaluations/ -│ └── -motion-policy/evaluation.json -└── run-manifest.json -``` - -Configurations are captured once when training completes. Checkpoint filenames -carry their training step, recorder filenames carry their episode index, and -each visual evaluation receives its own timestamped directory. Camera recorder -events use `videos/train` or `videos/eval` when `save_path` is omitted; an -explicit `save_path` keeps the configured location. The `evaluations` directory -is created by the first evaluation run. - -A simulator RL task manifest has the following structure: - -```json -{ - "schema_version": 1, - "motion_profile": null, - "configs": { - "train": "configs/train.yaml", - "gym": "configs/gym.yaml" - }, - "checkpoints": { - "best": "checkpoints/cart_pole_grpo_best.pt", - "latest": "checkpoints/cart_pole_grpo_step_.pt" - } -} -``` - -| Field | Description | -|---|---| -| `schema_version` | Manifest schema version; currently `1` | -| `motion_profile` | `null` for an EmbodiChain training run; an external Profile ID when configured | -| `configs.train` | Snapshot of the training configuration | -| `configs.gym` | Snapshot of the Gym configuration for a simulator task | -| `checkpoints.best` | Relative path to the best checkpoint; `null` when no best checkpoint was produced | -| `checkpoints.latest` | Relative path to the checkpoint saved at the end of training | - -Every path is relative to the run directory. `--checkpoint best` selects the -best checkpoint and falls back to `latest` when the manifest value is `null`. -`--checkpoint latest` selects the latest checkpoint directly. A relative path -to another checkpoint in the run is also accepted. The resolved checkpoint is -recorded in `evaluation.json`. - -When the manifest declares a `motion_profile`, the CLI uses that Profile by -default. Add `--original-task` to recreate the EmbodiChain Environment recorded -by the training configurations instead. - -`--config` and `--gym-config` replace the corresponding manifest entries for -one command. This is useful when checking a checkpoint with an updated local -configuration while retaining the run's checkpoint selection and report path. - -Open the Viewer for a training run: - -```bash -embodichain eval-motion-policy outputs/_ \ - --checkpoint best \ - --device cuda:0 \ - --sim-device gpu \ - --viewer -``` - -The Viewer runs until its window closes when neither `--control-steps` nor -`--duration` is set. Press `R` to start or stop Viewer recording. A Headless -run uses an exact integer number of control steps: - -```bash -embodichain eval-motion-policy outputs/_ \ - --checkpoint latest \ - --device cuda:0 \ - --sim-device gpu \ - --control-steps 500 -``` - -## CartPole example - -Train a Policy with the repository's CartPole GRPO configuration: - -```bash -embodichain train-rl \ - --config embodichain_tasks/configs/agents/rl/basic/cart_pole/train_config_grpo.yaml -``` - -Pass the new run directory to the evaluation command: - -```bash -embodichain eval-motion-policy \ - outputs/cart_pole_grpo_ \ - --checkpoint best \ - --viewer -``` - -The CLI finds the `.pt`, training configuration, and Gym configuration through -`run-manifest.json`. It recreates the CartPole Environment with `num_envs=1`, -rebuilds the `actor_only` Policy, loads its weights, and opens the Viewer. The -report is written under `/evaluations/`. - -## Existing checkpoints - -For a training output created before run manifests were introduced, pass the -checkpoint and its training configuration explicitly: - -```bash -embodichain eval-motion-policy \ - --checkpoint /path/to/policy.pt \ - --config /path/to/train.yaml \ - --gym-config /path/to/gym.yaml \ - --device cuda:0 \ - --sim-device gpu \ - --viewer -``` - -When `trainer.gym_config` already points to the task configuration, -`--gym-config` can be omitted. - -## Execution pipeline - -```mermaid -flowchart LR - Config["Training configuration"] --> Runtime["Shared RL Runtime"] - Checkpoint[".pt checkpoint"] --> Policy["Policy + weights"] - Runtime --> Env["Original EmbodiChain Environment"] - Runtime --> Policy - Env -->|observation| Adapter["EmbodiChain Policy Adapter"] - Policy --> Adapter - Adapter -->|raw Policy action| TaskEnv["EmbodiChain Environment Bridge"] - TaskEnv -->|original action manager| Env - Env -->|env.step| TaskEnv - Evaluator["DexSim MotionPolicyEvaluator"] --> Adapter - Evaluator --> TaskEnv - TaskEnv --> Report["evaluation.json"] -``` - -Each control cycle follows this sequence: - -```mermaid -sequenceDiagram - autonumber - participant Eval as MotionPolicyEvaluator - participant Adapter as EmbodiChainTaskPolicyAdapter - participant TaskEnv as EmbodiChainTaskEnvironment - participant Policy as Policy - participant Env as Original Environment - - Eval->>TaskEnv: reset() - TaskEnv->>Env: reset() - Env-->>TaskEnv: observation and task state - TaskEnv-->>Eval: EvaluationFrame - Eval->>Adapter: reset(frame) - loop Integer control steps - Eval->>Adapter: infer(frame) - Adapter->>Policy: prepare observation and run deterministic inference - Policy-->>Adapter: action tensor - Adapter-->>Eval: PolicyOutput - Eval->>TaskEnv: step(action) - TaskEnv->>Env: apply ActionManager and call env.step() - Env-->>TaskEnv: observation, reward, terminated, truncated, info - TaskEnv-->>Eval: EnvironmentStep - end -``` - -| Component | Responsibility | -|---|---| -| `embodichain.learning.rl.runtime` `[new]` | Builds the Environment and Policy shared by training and evaluation | -| `embodichain.learning.rl.evaluation` `[updated]` | Provides shared observation flattening, deterministic inference, and action conversion | -| `EmbodiChainTaskPolicyAdapter` `[new]` | Reads the original observation from the frame, calls the shared inference path, and returns the raw action | -| `EmbodiChainTaskEnvironment` `[new]` | Calls the original Environment reset, action manager, and step methods, then exposes task results | -| `MotionPolicyEvaluator` | Schedules the Adapter and Environment, counts control steps, applies termination behavior, and updates the Viewer | - -This path loads `.pt` checkpoints through their training-time PyTorch Policy -definition. A standalone DexSim Policy project can use ONNX through its own -Adapter. - -## Viewer and Headless modes - -A simulator RL training run uses `trainer.gym_config` to recreate its original -Environment with `num_envs=1`. An external Policy uses a Motion Profile to -provide the robot assets, Adapter, and Environment. The following Viewer paths -have been validated: - -| Input | Viewer scene | Validated tasks | -|---|---|---| -| EmbodiChain training run | Robot, task objects, and scene restored from the Gym configuration | CartPole GRPO, PushCube PPO | -| External Policy | Robot and scene created by the Motion Profile | ANYmal-C velocity TorchScript `.pt` | - -A lightweight learning task uses `trainer.learning_env` to recreate its tensor -environment for Headless evaluation. PointMass stores its positions, -velocities, targets, and obstacles in PyTorch tensors, so its results are -reported through terminal metrics and `evaluation.json`. - -### Policy loading - -Evaluation rebuilds the Policy from its `policy` configuration, loads the Policy -state dictionary from the checkpoint, and runs deterministic inference. PPO, -GRPO, and APG select the training update algorithm; evaluation follows the -configured Policy type and weights. - -| Training algorithm | Policy type | Viewer validation | Headless validation | -|---|---|---|---| -| PPO | `actor_critic` | PushCube | PushCube, PointMass | -| GRPO | `actor_only` | CartPole | CartPole | -| APG | `actor_only` | — | PointMass | - -A custom EmbodiChain Policy registers its model builder and keeps the matching -network configuration in the training run. An external checkpoint uses a Motion -Profile to describe its model, robot assets, observation, and action processing. -A robot-state Policy can use the Motion Policy Kit default Environment. A Policy -with task objects supplies a complete task Environment. For example, -`UR5CubeResetEnvironment` creates a UR5, table, cube, and target, while -`AllegroReorientCubeEnvironment` creates an Allegro Hand, an object cube, and a -target pose. - -## Termination and duration - -| Option | Behavior | -|---|---| -| `--episodes N` | Completes `N` task episodes | -| `--control-steps N` | Executes exactly `N` Policy actions | -| `--duration SECONDS` | Converts the duration to an integer number of Policy control steps before execution | -| `--termination-behavior auto_reset` | Restores the original Environment after termination; used by the Viewer by default | -| `--termination-behavior pause` | Keeps the terminal task state visible | - -`--control-steps` and `--duration` are mutually exclusive. Simulation time is -calculated from integer step counters using `physics_dt` and -`sim_steps_per_control`. - -## Rendering - -The default renderer is `hybrid`: - -```bash -embodichain eval-motion-policy --viewer --renderer hybrid -``` - -For an EmbodiChain training task, the original Environment provides its robot, -task objects, and scene parameters. An external Motion Profile can use -`--scene-config standard`, `classic`, or a custom YAML file to configure the -ground, lighting, and post-processing. - -## Output - -Each run writes an `evaluation.json` report containing: - -- run, checkpoint, and configuration paths; -- inference and simulation devices, renderer, versions, and commits; -- termination reason, integer simulation steps, and control steps; -- episode reward, episode length, success, and task metrics. - -During native task evaluation, every `info["metrics"]` scalar is forwarded to -the Evaluator on that control step. The latest values are included in the final -report with the completed-episode reward, length, and success aggregates. - -A training run writes reports under `/evaluations/`. An explicit -checkpoint writes them under the `evaluations/` directory next to that -checkpoint. `--output` selects another output directory. - -## External Policy example - -`examples/learning/motion_policy_evaluation/` contains a runnable ANYmal-C -velocity example. Its public TorchScript `.pt` accepts `vx`, `vy`, and `yaw`, -uses a 48-dimensional observation, and outputs 12 joint actions. The example -includes resource preparation, a Motion Profile, a complete Adapter, tests, and -run commands. - -```text -examples/learning/motion_policy_evaluation/ -├── README.md -├── prepare_resources.py -├── eval_policy.py -└── anymal_c/ - ├── __init__.py - └── profile.py -``` - -Prepare the upstream model and robot resources: - -```bash -python examples/learning/motion_policy_evaluation/prepare_resources.py -``` - -The script fetches `mjw_anymal.pt`, its Policy configuration and license, and -the ANYmal-C URDF and meshes from newton-assets commit -`261cd1f429619d8ef4f546bd788ab9dea906b5e1`. It verifies the expected resource -digests and prints each download and Git operation. Re-running the command -continues an existing checkout. The resulting cache has this structure: - -```text -~/.cache/embodichain/examples/anymal_c_velocity/ -└── upstream/ - └── anybotics_anymal_c/ - ├── rl_policies/ - │ ├── mjw_anymal.pt - │ └── anymal.yaml - ├── urdf/anymal.urdf - └── meshes/... -``` - -Open the Viewer: - -```bash -python examples/learning/motion_policy_evaluation/eval_policy.py \ - --viewer \ - --renderer hybrid -``` - -`eval_policy.py` imports the adjacent `anymal_c/profile.py`, registers its -Profile in the current process, and fills the checkpoint and robot resource -paths from the default cache. Run it directly from the repository root. Options -such as `--device`, `--sim-device`, `--control-steps`, and `--scene-config` are -forwarded to EmbodiChain. - -Use W/S to adjust `vx`, A/D to adjust `vy`, Q/E to adjust `yaw`, M to zero all -three commands, Backspace to reset the task and camera framing, and T to switch -between tracking and free view. The Viewer follows the robot root in the ground -plane. Hold the left mouse button to adjust the orbit angle and use the mouse -wheel to adjust the viewing distance. Right-button panning is locked while -tracking is active. The tracking target continues following the robot during an -orbit drag. Switching back to tracking centers the camera on the current robot -position. - -The external Policy data path is: - -```mermaid -flowchart LR - CLI[eval-motion-policy] --> Profile[build_profile] - Profile --> Spec[Policy Spec
robot, control parameters, frequency] - Spec --> Setup[Adapter.setup
load TorchScript and JointMap] - Setup --> State[RobotState] - Command[vx + vy + yaw] --> Obs[48-dimensional observation] - State --> Obs - Obs --> Actor[TorchScript actor] - Actor --> Action[scale + default position] - Action --> Sim[advance the Environment] -``` - -The Adapter follows the upstream Policy definition when concatenating body -linear velocity, body angular velocity, projected gravity, the three-dimensional -command, joint position, joint velocity, and previous action. The model contains -its observation normalizer. The actor output is converted to joint position -targets with `default_position + 0.5 * action`. Simulation runs at 200 Hz, and -the Policy runs every four simulation steps at 50 Hz. - -To integrate another external Policy, copy this example and replace: - -1. the model and robot resource sources, fixed revisions, paths, and digests; -2. the initial pose, joint control parameters, and control frequency in `build_profile()`; -3. model loading and joint order in `Adapter.setup()`; -4. observation construction, normalization, network forward pass, and action processing in `Adapter.infer()`; -5. `PROFILE_ID` and the Profile name used by `eval_policy.py`. - -Task data can be resolved to local paths through Policy Spec `resources`, or an -Adapter can call an existing project reader from `__init__()` or `setup()`. -Python code can provide changing runtime data through `EvaluationFrame.inputs`. - -The DexSim Motion Policy Kit documentation describes the default Environment, -task Environment interface, Adapter API, task resources, and Python API. - -## Command summary - -| Scenario | Command | -|---|---| -| New training run | `embodichain eval-motion-policy --checkpoint best --viewer` | -| Original task when a Profile is configured | `embodichain eval-motion-policy --original-task --viewer` | -| Existing `.pt` | `embodichain eval-motion-policy --checkpoint --config --viewer` | -| Headless smoke test | `embodichain eval-motion-policy --control-steps 20` | -| External ANYmal-C `.pt` | `python examples/learning/motion_policy_evaluation/eval_policy.py --viewer` | diff --git a/docs/source/guides/policy_evaluation.md b/docs/source/guides/policy_evaluation.md new file mode 100644 index 000000000..f83edd9b0 --- /dev/null +++ b/docs/source/guides/policy_evaluation.md @@ -0,0 +1,196 @@ +# Policy Evaluation + +`embodichain eval-policy` evaluates a saved EmbodiChain policy after training. +It reconstructs the policy and environment from the training configuration, +loads the selected checkpoint, and writes a standalone evaluation report. + +The command runs Headless by default. Add `--viewer` to open the original +simulator task in the DexSim Viewer. + +## Training output + +`train-rl` records the files required by a later evaluation: + +```text +outputs/_/ +├── checkpoints/ +│ └── policy_*.pt +├── configs/ +│ ├── train.yaml +│ └── gym.yaml +├── logs/ +├── videos/ +│ ├── train/ +│ └── eval/ +└── run-manifest.json +``` + +`configs/gym.yaml` is present for simulator tasks. The first evaluation adds: + +```text +evaluations/ +└── -policy/ + └── evaluation.json +``` + +`run-manifest.json` connects the run directory to its configuration snapshots +and checkpoints: + +```json +{ + "schema_version": 1, + "configs": { + "train": "configs/train.yaml", + "gym": "configs/gym.yaml" + }, + "checkpoints": { + "best": "checkpoints/cart_pole_grpo_best.pt", + "latest": "checkpoints/cart_pole_grpo_step_4096.pt" + } +} +``` + +All paths in the manifest are relative to the run directory. `best` is `null` +when training did not select a best checkpoint. + +## Evaluate a training run + +The shortest command selects `latest` and runs the configured number of +Headless evaluation episodes: + +```bash +embodichain eval-policy outputs/_ +``` + +Select the best checkpoint and override the episode count: + +```bash +embodichain eval-policy outputs/_ \ + --checkpoint best \ + --episodes 10 +``` + +Open the original simulator task in the Viewer: + +```bash +embodichain eval-policy outputs/_ \ + --checkpoint best \ + --viewer \ + --renderer hybrid \ + --device cuda:0 \ + --sim-device gpu +``` + +The Viewer uses one environment and keeps running until the window closes. Use +`--episodes`, `--control-steps`, or `--duration` to select another stopping +condition. + +For a checkpoint created before `run-manifest.json` was introduced, provide +its training configuration directly: + +```bash +embodichain eval-policy \ + --checkpoint /path/to/policy.pt \ + --config /path/to/train.yaml \ + --gym-config /path/to/gym.yaml \ + --viewer +``` + +`--gym-config` can be omitted when the training configuration already refers to +the task configuration. + +## Execution paths + +```mermaid +flowchart LR + Run[Training run] --> Manifest[run-manifest.json] + Manifest --> Config[Training config] + Manifest --> Checkpoint[Checkpoint] + Config --> Runtime[EmbodiChain RL runtime] + Checkpoint --> Runtime + Runtime --> Headless[Headless episode evaluation] + Runtime --> Viewer[DexSim MotionPolicyEvaluator] + Headless --> Report[evaluation.json] + Viewer --> Report + Profile[External Motion Profile] --> Viewer +``` + +Headless evaluation calls the existing `evaluate_episodes()` path. Viewer +evaluation keeps the task's original observation, action processing, reset, +reward, termination, objects, and sensors: + +```mermaid +sequenceDiagram + participant Evaluator as MotionPolicyEvaluator + participant Adapter as EmbodiChainTaskPolicyAdapter + participant Task as EmbodiChainTaskEnvironment + participant Policy as EmbodiChain Policy + participant Env as Original task Environment + + Evaluator->>Task: reset() + Task->>Env: reset() + Env-->>Task: observation and task state + Task-->>Evaluator: EvaluationFrame + loop Each control step + Evaluator->>Adapter: infer(frame) + Adapter->>Policy: deterministic inference + Policy-->>Adapter: action + Adapter-->>Evaluator: PolicyOutput + Evaluator->>Task: step(action) + Task->>Env: action processing and env.step() + Env-->>Task: observation, reward, termination and info + Task-->>Evaluator: EnvironmentStep + end +``` + +| Input | Headless | Viewer | +|---|---:|---:| +| EmbodiChain lightweight RL environment | Yes | — | +| EmbodiChain simulator RL environment | Yes | Yes | +| Registered external Motion Profile | Yes | Yes | + +Policy reconstruction follows the model definition stored in the training +configuration. The Viewer path has been validated with CartPole GRPO and +PushCube PPO checkpoints. + +## Viewer controls + +| Key | Action | +|---|---| +| `Backspace` | Reset the task and camera framing | +| `T` | Switch between tracking and free camera modes when the Environment provides a tracking target | +| `R` | Start or stop recording | +| `Esc` | Close the Viewer | + +While tracking is active, drag with the left mouse button to orbit and use the +mouse wheel to zoom. + +## External policy example + +The repository includes a concrete ANYmal-C velocity example under +`examples/learning/policy_evaluation/`. It prepares a public TorchScript +checkpoint and robot assets, registers an adjacent Motion Profile, and forwards +the remaining arguments to `eval-policy`. + +```bash +python examples/learning/policy_evaluation/prepare_resources.py +python examples/learning/policy_evaluation/eval_policy.py \ + --viewer \ + --renderer hybrid \ + --sim-device gpu +``` + +Use W/S for `vx`, A/D for `vy`, Q/E for `yaw`, and M to zero the command. See +the [example README](https://github.com/DexForce/EmbodiChain/tree/main/examples/learning/policy_evaluation) +for the resource layout, observation construction, action conversion, and +Profile implementation. + +This example tracks the robot root in the ground plane. Press `T` to switch +between tracking and free view. + +## Evaluation report + +`evaluation.json` records the selected checkpoint and configs, task and device +information, episode results, and aggregated metrics. Reports are written to +`/evaluations/` for a training run and next to an explicit checkpoint by +default. Use `--output` to select another parent directory. diff --git a/docs/source/tutorial/rl.rst b/docs/source/tutorial/rl.rst index 556eac8d9..64bf26cc4 100644 --- a/docs/source/tutorial/rl.rst +++ b/docs/source/tutorial/rl.rst @@ -82,7 +82,7 @@ APG and PPO. Launch either config with the same CLI: embodichain train-rl --config embodichain_tasks/configs/agents/rl/basic/point_mass/train_ppo.yaml Configuration Sections ----------------------- +--------------------- Runtime Settings ^^^^^^^^^^^^^^^^ @@ -134,7 +134,7 @@ Example: } Policy Configuration -^^^^^^^^^^^^^^^^^^^^ +^^^^^^^^^^^^^^^^^^^ The ``policy`` section defines the neural network policy: @@ -265,24 +265,21 @@ All outputs are written to ``./outputs/_/``: - **logs/**: TensorBoard logs - **checkpoints/**: Model checkpoints - **configs/**: Training config and referenced gym config snapshots -- **videos/**: Training and evaluation recordings grouped by phase -- **evaluations/**: Timestamped visual evaluation reports -- **run-manifest.json**: training configs and best/latest checkpoint index used by ``eval-motion-policy`` +- **evaluations/**: Timestamped policy evaluation reports +- **run-manifest.json**: Training configs and best/latest checkpoint index used by ``eval-policy`` -A simulator training run can be opened directly in its original task Environment: +A training run can be evaluated Headless or opened in its simulator task: .. code-block:: bash - embodichain eval-motion-policy outputs/_ --viewer + embodichain eval-policy outputs/_ + embodichain eval-policy outputs/_ --viewer -An external Policy project can record its Motion Profile ID in -``trainer.motion_profile``. - -See :doc:`../guides/motion_policy_evaluation` for EmbodiChain ``.pt`` training +See :doc:`../guides/policy_evaluation` for EmbodiChain ``.pt`` training runs and the external Motion Profile example. Training Process -~~~~~~~~~~~~~~~~ +~~~~~~~~~~~~~~~ The training process follows this sequence: @@ -342,7 +339,7 @@ Available Algorithms - **GRPO**: Group Relative Policy Optimization (no Critic, step-wise returns, masked group normalization). Use ``actor_only`` policy. Set ``kl_coef=0`` for from-scratch training (CartPole, dense reward); ``kl_coef=0.02`` for VLA/LLM fine-tuning. Adding a New Algorithm ----------------------- +--------------------- To add a new algorithm: @@ -480,4 +477,4 @@ See Also - :doc:`basic_env` — Creating basic Gymnasium environments - :doc:`modular_env` — Advanced modular environments with managers - :doc:`/resources/task/index` — List of available RL task environments -- :doc:`/guides/motion_policy_evaluation` — Visual evaluation of EmbodiChain ``.pt`` checkpoints and external Motion Profiles +- :doc:`/guides/policy_evaluation` — Headless and Viewer evaluation of EmbodiChain ``.pt`` checkpoints and external Motion Profiles diff --git a/embodichain/__main__.py b/embodichain/__main__.py index 584dd05e8..643f2ce48 100644 --- a/embodichain/__main__.py +++ b/embodichain/__main__.py @@ -82,9 +82,9 @@ class Command: help="Train an RL agent from a JSON or YAML config.", ), Command( - name="eval-motion-policy", - target="embodichain.learning.rl.motion_policy_evaluation.cli:cli", - help="Visualize a policy through DexSim Motion Policy Kit.", + name="eval-policy", + target="embodichain.learning.rl.policy_evaluation.cli:cli", + help="Evaluate a trained policy in Headless or Viewer mode.", ), Command( name="annotate-grasp", diff --git a/embodichain/learning/rl/motion_policy_evaluation/__init__.py b/embodichain/learning/rl/motion_policy_evaluation/__init__.py deleted file mode 100644 index d68087a27..000000000 --- a/embodichain/learning/rl/motion_policy_evaluation/__init__.py +++ /dev/null @@ -1,50 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Visual motion-policy evaluation through DexSim Motion Policy Kit.""" - -from __future__ import annotations - -from .bridge import ( - MotionEvaluationResult, - create_motion_profile_evaluator, - evaluate_motion_profile, -) -from .checkpoint import load_policy_state_dict -from .manifest import RunManifest, write_run_manifest -from .profile import ( - MotionProfile, - MotionProfileRequest, - build_motion_profile, - get_motion_profile_names, - register_motion_profile, -) -from .report import write_evaluation_report - -__all__ = [ - "MotionEvaluationResult", - "MotionProfile", - "MotionProfileRequest", - "RunManifest", - "build_motion_profile", - "create_motion_profile_evaluator", - "evaluate_motion_profile", - "get_motion_profile_names", - "load_policy_state_dict", - "register_motion_profile", - "write_run_manifest", - "write_evaluation_report", -] diff --git a/embodichain/learning/rl/motion_policy_evaluation/checkpoint.py b/embodichain/learning/rl/motion_policy_evaluation/checkpoint.py deleted file mode 100644 index 817eddb4c..000000000 --- a/embodichain/learning/rl/motion_policy_evaluation/checkpoint.py +++ /dev/null @@ -1,60 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -"""Load the policy weights stored in an EmbodiChain checkpoint.""" - -from __future__ import annotations - -from collections.abc import Mapping -from pathlib import Path -from typing import Any - -import torch - -__all__ = ["load_policy_state_dict"] - - -def load_policy_state_dict( - checkpoint: str | Path, - *, - map_location: str | torch.device = "cpu", -) -> Mapping[str, Any]: - """Load the ``policy`` state mapping from an EmbodiChain ``.pt`` file. - - Args: - checkpoint: EmbodiChain training checkpoint. - map_location: Device passed to :func:`torch.load`. - - Returns: - Policy state mapping ready for ``module.load_state_dict()``. - - Raises: - FileNotFoundError: If the checkpoint does not exist. - TypeError: If the checkpoint or policy payload is not a mapping. - ValueError: If the checkpoint has no policy weights. - """ - path = Path(checkpoint).expanduser().resolve() - if not path.is_file(): - raise FileNotFoundError(f"Policy checkpoint does not exist: {path}") - payload = torch.load(path, map_location=map_location, weights_only=True) - if not isinstance(payload, Mapping): - raise TypeError("Policy checkpoint root must be a mapping") - state = payload.get("policy") - if not isinstance(state, Mapping): - raise TypeError("Policy checkpoint field 'policy' must be a mapping") - if not state: - raise ValueError("Policy checkpoint field 'policy' is empty") - return state diff --git a/tests/learning/rl/motion_policy_evaluation/__init__.py b/embodichain/learning/rl/policy_evaluation/__init__.py similarity index 71% rename from tests/learning/rl/motion_policy_evaluation/__init__.py rename to embodichain/learning/rl/policy_evaluation/__init__.py index 82062a260..4dbc590b4 100644 --- a/tests/learning/rl/motion_policy_evaluation/__init__.py +++ b/embodichain/learning/rl/policy_evaluation/__init__.py @@ -14,6 +14,20 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Tests for EmbodiChain motion-policy evaluation.""" +"""External Policy Profiles for ``embodichain eval-policy``.""" from __future__ import annotations + +from .profile import ( + MotionProfile, + MotionProfileRequest, + build_motion_profile, + register_motion_profile, +) + +__all__ = [ + "MotionProfile", + "MotionProfileRequest", + "build_motion_profile", + "register_motion_profile", +] diff --git a/embodichain/learning/rl/motion_policy_evaluation/bridge.py b/embodichain/learning/rl/policy_evaluation/bridge.py similarity index 71% rename from embodichain/learning/rl/motion_policy_evaluation/bridge.py rename to embodichain/learning/rl/policy_evaluation/bridge.py index 1299212a4..a336b9ccc 100644 --- a/embodichain/learning/rl/motion_policy_evaluation/bridge.py +++ b/embodichain/learning/rl/policy_evaluation/bridge.py @@ -14,24 +14,20 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Run an EmbodiChain Motion Profile through DexSim Motion Policy Kit.""" +"""Run an external Policy Profile through DexSim Motion Policy Kit.""" from __future__ import annotations from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path -from types import MappingProxyType from typing import Any from dexsim.kit.motion_policy import ( - MotionPolicyEvaluator, - PolicyAdapter, PolicySpec, ResolvedPolicy, ResourceResolver, RunOptions, - create_motion_policy_evaluator, load_scene_config, parse_policy_spec, policy_spec_to_dict, @@ -39,14 +35,11 @@ run_motion_policy, scene_config_to_dict, ) -from dexsim.kit.motion_policy.environment import PolicyEnvironment -from dexsim.kit.motion_policy.evaluator import InputProvider from .profile import MotionProfile __all__ = [ "MotionEvaluationResult", - "create_motion_profile_evaluator", "evaluate_motion_profile", ] @@ -62,22 +55,6 @@ class MotionEvaluationResult: summary: Mapping[str, Any] viewer: bool - def __post_init__(self) -> None: - object.__setattr__( - self, "policy_spec", MappingProxyType(dict(self.policy_spec)) - ) - object.__setattr__( - self, - "scene_config", - MappingProxyType(dict(self.scene_config)), - ) - object.__setattr__( - self, - "episodes", - tuple(MappingProxyType(dict(value)) for value in self.episodes), - ) - object.__setattr__(self, "summary", MappingProxyType(dict(self.summary))) - def evaluate_motion_profile( profile: MotionProfile, @@ -95,8 +72,6 @@ def evaluate_motion_profile( termination_behavior: str | None = None, cache_dir: str | Path | None = None, offline: bool = False, - input_provider: InputProvider | None = None, - environment: PolicyEnvironment | None = None, ) -> MotionEvaluationResult: """Resolve one Motion Profile and run its visual evaluation. @@ -115,8 +90,6 @@ def evaluate_motion_profile( termination_behavior: Policy termination handling override. cache_dir: Motion Policy Kit resource cache. offline: Use resources already available in the cache. - input_provider: Build per-frame task inputs for the Adapter. - environment: Prebuilt task environment owned by the Evaluator. Returns: Normalized inputs, episode results, and aggregate metrics. @@ -125,8 +98,6 @@ def evaluate_motion_profile( raise ValueError("episodes must be positive") if viewer and episodes != 1: raise ValueError("Viewer evaluation supports one episode") - if environment is not None and episodes != 1: - raise ValueError("A prebuilt environment supports one episode") parsed, resolved = _resolve_profile(profile, cache_dir, offline) resolved_scene = load_scene_config(scene_config) options = RunOptions( @@ -147,8 +118,6 @@ def evaluate_motion_profile( run_motion_policy( resolved, options, - environment=environment, - input_provider=input_provider, ), ) for index in range(episodes) @@ -163,37 +132,6 @@ def evaluate_motion_profile( ) -def create_motion_profile_evaluator( - profile: MotionProfile, - options: RunOptions | None = None, - *, - cache_dir: str | Path | None = None, - offline: bool = False, - adapter: PolicyAdapter | None = None, - environment: PolicyEnvironment | None = None, -) -> MotionPolicyEvaluator: - """Create a DexSim Evaluator for an EmbodiChain Motion Profile. - - Args: - profile: Provider-built profile containing the DexSim Policy Spec. - options: DexSim evaluation options. - cache_dir: Motion Policy Kit resource cache. - offline: Use resources already available in the cache. - adapter: Prebuilt policy adapter owned by the returned Evaluator. - environment: Prebuilt task environment owned by the returned Evaluator. - - Returns: - Configured evaluator ready for ``reset()``, ``step()``, or ``run()``. - """ - _parsed, resolved = _resolve_profile(profile, cache_dir, offline) - return create_motion_policy_evaluator( - resolved, - options, - adapter=adapter, - environment=environment, - ) - - def _resolve_profile( profile: MotionProfile, cache_dir: str | Path | None, diff --git a/embodichain/learning/rl/motion_policy_evaluation/cli.py b/embodichain/learning/rl/policy_evaluation/cli.py similarity index 54% rename from embodichain/learning/rl/motion_policy_evaluation/cli.py rename to embodichain/learning/rl/policy_evaluation/cli.py index 01038b195..853bd3daa 100644 --- a/embodichain/learning/rl/motion_policy_evaluation/cli.py +++ b/embodichain/learning/rl/policy_evaluation/cli.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Command line for visual evaluation of EmbodiChain policy checkpoints.""" +"""Unified policy evaluation for EmbodiChain training runs.""" from __future__ import annotations @@ -25,6 +25,7 @@ from pathlib import Path from typing import Any +import numpy as np import torch from embodichain import __version__ @@ -32,27 +33,23 @@ discover_task_packages, execute_init_hooks, ) +from embodichain.learning.rl.evaluation import evaluate_episodes from embodichain.learning.rl.runtime import ( + PolicyRuntime, build_gym_policy_runtime, build_learning_policy_runtime, - resolve_torch_device, - seed_policy_runtime, ) from embodichain.utils.utility import load_config -from .bridge import MotionEvaluationResult, evaluate_motion_profile -from .checkpoint import load_policy_state_dict from .manifest import RunManifest -from .native_task import NativeTaskEvaluationResult, evaluate_native_task -from .profile import MotionProfileRequest, build_motion_profile from .report import write_evaluation_report __all__ = ["cli", "parse_args", "run"] @dataclass(frozen=True) -class MotionInput: - """Resolved checkpoint, configs, and Profile selection.""" +class EvaluationInput: + """Checkpoint and configuration selected for one evaluation.""" checkpoint: Path profile: str | None @@ -62,28 +59,36 @@ class MotionInput: selected_checkpoint: str +@dataclass(frozen=True) +class NativeRuntime: + """Reconstructed EmbodiChain task and its runtime choices.""" + + runtime: PolicyRuntime + device: torch.device + simulation_device: torch.device + seed: int + renderer: str + uses_simulator: bool + trainer: Mapping[str, Any] + + def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: - """Parse ``embodichain eval-motion-policy`` arguments.""" + """Parse ``embodichain eval-policy`` arguments.""" parser = argparse.ArgumentParser( - prog="embodichain eval-motion-policy", - description="Open a DexSim Viewer for an EmbodiChain policy checkpoint.", + prog="embodichain eval-policy", + description="Evaluate an EmbodiChain or external policy checkpoint.", ) parser.add_argument("run", nargs="?", help="EmbodiChain training run directory.") - mode = parser.add_mutually_exclusive_group() - mode.add_argument("--profile", help="Registered Motion Profile name.") - mode.add_argument( - "--original-task", - action="store_true", - help="Use the original EmbodiChain task instead of the manifest Profile.", - ) + parser.add_argument("--profile", help="Registered external Policy Profile.") parser.add_argument( "--checkpoint", - help="Checkpoint path, or best/latest when RUN is supplied.", + help="latest, best, or a checkpoint path; defaults to latest with RUN.", ) parser.add_argument("--config", help="Training config for an explicit checkpoint.") - parser.add_argument("--gym-config", help="Task config for an explicit checkpoint.") - parser.add_argument("--resource-root", help="Profile-specific local resource root.") + parser.add_argument("--gym-config", help="Task config override.") + parser.add_argument("--resource-root", help="External Profile resource root.") parser.add_argument("--episodes", type=int) + parser.add_argument("--num-envs", type=int) count = parser.add_mutually_exclusive_group() count.add_argument("--control-steps", type=int) count.add_argument("--duration", type=float) @@ -91,11 +96,10 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.add_argument("--device", help="PyTorch inference device.") parser.add_argument("--sim-device", choices=("cpu", "gpu")) parser.add_argument("--seed", type=int) - parser.add_argument("--physics-backend", choices=("default",)) + parser.add_argument("--physics-backend") parser.add_argument( "--renderer", choices=("raster", "hybrid", "fastrt", "offlinert"), - default="hybrid", ) parser.add_argument("--gpu-id", type=int, default=0) parser.add_argument("--scene-config") @@ -111,52 +115,20 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: def run(args: argparse.Namespace) -> Path: - """Resolve inputs, run the evaluation, and write ``evaluation.json``.""" + """Run Headless or Viewer evaluation and write ``evaluation.json``.""" resolved = _resolve_input(args) + if resolved.profile is not None: + return _run_profile(args, resolved) discover_task_packages() execute_init_hooks() - if resolved.profile is None: - return _run_native_task(args, resolved) - - device = torch.device(args.device or "cpu") - profile = build_motion_profile( - resolved.profile, - MotionProfileRequest( - checkpoint=resolved.checkpoint, - device=device, - configs=resolved.configs, - resource_root=( - None if args.resource_root is None else Path(args.resource_root) - ), - renderer=args.renderer, - ), - ) - for warning in profile.warnings: - print(f"Warning: {warning}", file=sys.stderr) - result = evaluate_motion_profile( - profile, - episodes=args.episodes or 1, - viewer=args.viewer, - control_steps=args.control_steps, - duration=args.duration, - command=None if args.command is None else tuple(args.command), - scene_config=args.scene_config or "standard", - physics_backend=args.physics_backend, - simulation_device=args.sim_device or "cpu", - renderer=args.renderer, - gpu_id=args.gpu_id, - termination_behavior=args.termination_behavior, - cache_dir=args.cache_dir, - offline=args.offline, - ) - return write_evaluation_report( - _output_parent(args.output, resolved), - _profile_report(result, resolved, device), - ) + _validate_native_options(args) + if args.viewer: + return _run_native_viewer(args, resolved) + return _run_native_headless(args, resolved) def cli(argv: Sequence[str] | None = None) -> None: - """Run motion-policy evaluation from the unified EmbodiChain CLI.""" + """Run policy evaluation from the unified EmbodiChain CLI.""" try: report = run(parse_args(argv)) except ( @@ -167,14 +139,14 @@ def cli(argv: Sequence[str] | None = None) -> None: TypeError, ValueError, ) as error: - raise SystemExit(f"eval-motion-policy: {error}") from error + raise SystemExit(f"eval-policy: {error}") from error print(f"Evaluation report: {report}") -def _resolve_input(args: argparse.Namespace) -> MotionInput: +def _resolve_input(args: argparse.Namespace) -> EvaluationInput: if args.run is not None: manifest = RunManifest.load(args.run) - requested = args.checkpoint or "best" + requested = args.checkpoint or "latest" if requested in {"best", "latest"}: selected, checkpoint = manifest.select_checkpoint(requested) else: @@ -190,15 +162,9 @@ def _resolve_input(args: argparse.Namespace) -> MotionInput: configs["train"] = Path(args.config).expanduser().resolve() if args.gym_config is not None: configs["gym"] = Path(args.gym_config).expanduser().resolve() - if args.original_task: - profile = None - elif args.profile is not None: - profile = args.profile - else: - profile = manifest.motion_profile - return MotionInput( + return EvaluationInput( checkpoint=checkpoint, - profile=profile, + profile=args.profile, configs=configs, run=manifest.root, requested_checkpoint=requested, @@ -206,16 +172,15 @@ def _resolve_input(args: argparse.Namespace) -> MotionInput: ) if args.checkpoint is None: raise ValueError("--checkpoint is required without RUN") - checkpoint = Path(args.checkpoint).expanduser().resolve() configs = {} if args.config is not None: configs["train"] = Path(args.config).expanduser().resolve() if args.gym_config is not None: configs["gym"] = Path(args.gym_config).expanduser().resolve() if args.profile is None and "train" not in configs: - raise ValueError("--config is required for a native EmbodiChain checkpoint") - return MotionInput( - checkpoint=checkpoint, + raise ValueError("--config is required for an EmbodiChain checkpoint") + return EvaluationInput( + checkpoint=Path(args.checkpoint).expanduser().resolve(), profile=args.profile, configs=configs, run=None, @@ -224,39 +189,146 @@ def _resolve_input(args: argparse.Namespace) -> MotionInput: ) -def _run_native_task(args: argparse.Namespace, resolved: MotionInput) -> Path: - _validate_native_options(args) +def _run_native_headless( + args: argparse.Namespace, + resolved: EvaluationInput, +) -> Path: + native = _build_native_runtime(args, resolved, viewer=False) + episodes = ( + args.episodes + if args.episodes is not None + else int(native.trainer.get("num_eval_episodes", 5)) + ) + try: + metrics = evaluate_episodes( + policy=native.runtime.policy, + env=native.runtime.env, + num_episodes=episodes, + device=native.device, + seed=native.seed, + ) + finally: + native.runtime.close() + _flush_simulator(native.uses_simulator) + return write_evaluation_report( + _output_parent(args.output, resolved), + _headless_report(native, resolved, episodes, metrics), + ) + + +def _run_native_viewer( + args: argparse.Namespace, + resolved: EvaluationInput, +) -> Path: + from .viewer import evaluate_native_viewer + + native = _build_native_runtime(args, resolved, viewer=True) + try: + result = evaluate_native_viewer( + native.runtime, + seed=native.seed, + episodes=args.episodes, + control_steps=args.control_steps, + duration=args.duration, + termination_behavior=args.termination_behavior or "auto_reset", + ) + finally: + _flush_simulator(True) + return write_evaluation_report( + _output_parent(args.output, resolved), + _viewer_report(result, native, resolved), + ) + + +def _run_profile(args: argparse.Namespace, resolved: EvaluationInput) -> Path: + from .bridge import evaluate_motion_profile + from .profile import MotionProfileRequest, build_motion_profile + + device = _torch_device(args.device or "cpu") + renderer = args.renderer or "hybrid" + profile = build_motion_profile( + resolved.profile, + MotionProfileRequest( + checkpoint=resolved.checkpoint, + device=device, + configs=resolved.configs, + resource_root=( + None if args.resource_root is None else Path(args.resource_root) + ), + renderer=renderer, + ), + ) + for warning in profile.warnings: + print(f"Warning: {warning}", file=sys.stderr) + result = evaluate_motion_profile( + profile, + episodes=args.episodes if args.episodes is not None else 1, + viewer=args.viewer, + control_steps=args.control_steps, + duration=args.duration, + command=None if args.command is None else tuple(args.command), + scene_config=args.scene_config or "standard", + physics_backend=args.physics_backend, + simulation_device=args.sim_device or "cpu", + renderer=renderer, + gpu_id=args.gpu_id, + termination_behavior=args.termination_behavior, + cache_dir=args.cache_dir, + offline=args.offline, + ) + return write_evaluation_report( + _output_parent(args.output, resolved), + _profile_report(result, resolved, device), + ) + + +def _build_native_runtime( + args: argparse.Namespace, + resolved: EvaluationInput, + *, + viewer: bool, +) -> NativeRuntime: train_config = resolved.configs.get("train") if train_config is None: - raise ValueError("Training config is required for native task evaluation") + raise ValueError("Training config is required for an EmbodiChain checkpoint") config = load_config(train_config) config["trainer"] = dict(config["trainer"]) gym_config = resolved.configs.get("gym") if gym_config is not None: config["trainer"]["gym_config"] = str(gym_config) trainer = config["trainer"] - device = resolve_torch_device(args.device or trainer.get("device", "cpu")) - if args.sim_device == "gpu": - simulation_device = resolve_torch_device(f"cuda:{args.gpu_id}") - elif args.sim_device == "cpu": - simulation_device = torch.device("cpu") - else: - simulation_device = device + device = _torch_device(args.device or trainer.get("device", "cpu")) + simulation_device = _simulation_device(args, device) seed = int( args.seed if args.seed is not None else trainer.get("eval_seed", int(trainer.get("seed", 1)) + 10_000) ) - seed_policy_runtime(seed, device) - uses_simulator = "learning_env" not in trainer + np.random.seed(seed) + torch.manual_seed(seed) + if device.type == "cuda": + torch.cuda.manual_seed_all(seed) + uses_simulator = "gym_config" in trainer + if viewer and not uses_simulator: + raise ValueError("--viewer requires a simulator training task") + renderer = args.renderer or str(trainer.get("renderer", "hybrid")) + num_envs = ( + 1 + if viewer + else int( + args.num_envs + if args.num_envs is not None + else trainer.get("num_eval_envs", 4) + ) + ) if uses_simulator: runtime = build_gym_policy_runtime( config, device=device, simulation_device=simulation_device, - num_envs=1, - headless=not args.viewer, - renderer=args.renderer, + num_envs=num_envs, + headless=not viewer, + renderer=renderer, gpu_id=args.gpu_id, config_dir=train_config.parent, ) @@ -264,55 +336,27 @@ def _run_native_task(args: argparse.Namespace, resolved: MotionInput) -> Path: runtime = build_learning_policy_runtime( config, device=device, - num_envs=1, + num_envs=num_envs, ) try: - try: - runtime.policy.load_state_dict( - load_policy_state_dict(resolved.checkpoint, map_location="cpu") - ) - except Exception: - runtime.env.close() - raise - - episodes = args.episodes - if ( - episodes is None - and not args.viewer - and args.control_steps is None - and args.duration is None - ): - episodes = 1 - result = evaluate_native_task( - runtime, - seed=seed, - viewer=args.viewer, - episodes=episodes, - control_steps=args.control_steps, - duration=args.duration, - termination_behavior=args.termination_behavior or "auto_reset", - ) - finally: - if uses_simulator: - from embodichain.lab.sim.sim_manager import SimulationManager - - SimulationManager.flush_cleanup_queue() - return write_evaluation_report( - _output_parent(args.output, resolved), - _native_report( - result, - resolved, - device=device, - simulation_device=simulation_device, - seed=seed, - renderer=args.renderer, - ), + runtime.policy.load_state_dict(_load_policy_state_dict(resolved.checkpoint)) + except Exception: + runtime.close() + _flush_simulator(uses_simulator) + raise + return NativeRuntime( + runtime=runtime, + device=device, + simulation_device=simulation_device, + seed=seed, + renderer=renderer, + uses_simulator=uses_simulator, + trainer=trainer, ) def _validate_native_options(args: argparse.Namespace) -> None: - """Reject options that belong to external Motion Profiles.""" - values = { + profile_options = { "--resource-root": args.resource_root, "--command": args.command, "--physics-backend": args.physics_backend, @@ -320,12 +364,62 @@ def _validate_native_options(args: argparse.Namespace) -> None: "--cache-dir": args.cache_dir, "--offline": args.offline, } - selected = [name for name, value in values.items() if value not in (None, False)] + selected = [ + name for name, value in profile_options.items() if value not in (None, False) + ] if selected: raise ValueError(f"{', '.join(selected)} requires --profile") + if not args.viewer and ( + args.control_steps is not None + or args.duration is not None + or args.termination_behavior is not None + ): + raise ValueError( + "--control-steps, --duration, and --termination-behavior require --viewer" + ) + + +def _load_policy_state_dict(checkpoint: Path) -> Mapping[str, Any]: + payload = torch.load(checkpoint, map_location="cpu", weights_only=True) + if not isinstance(payload, Mapping) or not isinstance( + payload.get("policy"), Mapping + ): + raise TypeError("Checkpoint must contain a 'policy' state mapping") + return payload["policy"] + + +def _torch_device(value: str) -> torch.device: + device = torch.device(value) + if device.type == "cuda": + index = ( + device.index if device.index is not None else torch.cuda.current_device() + ) + torch.cuda.set_device(index) + return torch.device(f"cuda:{index}") + if device.type != "cpu": + raise ValueError(f"Unsupported device type: {device.type}") + return device + + +def _simulation_device( + args: argparse.Namespace, + inference_device: torch.device, +) -> torch.device: + if args.sim_device == "gpu": + return _torch_device(f"cuda:{args.gpu_id}") + if args.sim_device == "cpu": + return torch.device("cpu") + return inference_device + +def _flush_simulator(enabled: bool) -> None: + if enabled: + from embodichain.lab.sim.sim_manager import SimulationManager -def _output_parent(configured: str | None, resolved: MotionInput) -> Path: + SimulationManager.flush_cleanup_queue() + + +def _output_parent(configured: str | None, resolved: EvaluationInput) -> Path: if configured is not None: return Path(configured) if resolved.run is not None: @@ -333,82 +427,97 @@ def _output_parent(configured: str | None, resolved: MotionInput) -> Path: return resolved.checkpoint.parent / "evaluations" -def _profile_report( - result: MotionEvaluationResult, - resolved: MotionInput, - device: torch.device, +def _checkpoint_inputs(resolved: EvaluationInput) -> dict[str, Any]: + return { + "run": resolved.run, + "checkpoint": { + "path": resolved.checkpoint, + "requested": resolved.requested_checkpoint, + "selected": resolved.selected_checkpoint, + }, + "configs": resolved.configs, + } + + +def _headless_report( + native: NativeRuntime, + resolved: EvaluationInput, + episodes: int, + metrics: Mapping[str, float], +) -> dict[str, Any]: + return { + "mode": "headless", + "inputs": { + **_checkpoint_inputs(resolved), + "task_id": native.runtime.env_id, + "seed": native.seed, + "num_envs": int(native.runtime.env.num_envs), + "device": str(native.device), + "embodichain_version": __version__, + }, + "result": {"episodes": episodes, "metrics": metrics}, + } + + +def _viewer_report( + result: Any, + native: NativeRuntime, + resolved: EvaluationInput, ) -> dict[str, Any]: import dexsim return { - "mode": "viewer" if result.viewer else "headless", + "mode": "viewer", "inputs": { - "run": resolved.run, - "checkpoint": { - "path": resolved.checkpoint, - "requested": resolved.requested_checkpoint, - "selected": resolved.selected_checkpoint, - }, - "profile": { - "id": result.profile.profile_id, - "provider_version": result.profile.provider_version, - "provenance": result.profile.provenance, - "warnings": result.profile.warnings, - }, - "configs": resolved.configs, - "policy_spec": result.policy_spec, - "scene_config": result.scene_config, - "inference_device": str(device), + **_checkpoint_inputs(resolved), + "task_id": result.task_id, + "seed": native.seed, + "inference_device": str(native.device), + "simulation_device": str(native.simulation_device), + "renderer": native.renderer, "embodichain_version": __version__, "dexsim_version": getattr(dexsim, "__version__", None), "dexsim_commit": getattr(dexsim, "__commit_id__", None), }, "result": { + "reason": result.reason, + "simulation_time": result.simulation_time, + "simulation_steps": result.simulation_steps, + "control_steps": result.control_steps, + "requested_duration": result.requested_duration, + "effective_duration": result.effective_duration, "episodes": result.episodes, - "summary": result.summary, + "metrics": result.metrics, }, } -def _native_report( - result: NativeTaskEvaluationResult, - resolved: MotionInput, - *, +def _profile_report( + result: Any, + resolved: EvaluationInput, device: torch.device, - simulation_device: torch.device, - seed: int, - renderer: str, ) -> dict[str, Any]: import dexsim return { "mode": "viewer" if result.viewer else "headless", "inputs": { - "run": resolved.run, - "checkpoint": { - "path": resolved.checkpoint, - "requested": resolved.requested_checkpoint, - "selected": resolved.selected_checkpoint, + **_checkpoint_inputs(resolved), + "profile": { + "id": result.profile.profile_id, + "provider_version": result.profile.provider_version, + "provenance": result.profile.provenance, + "warnings": result.profile.warnings, }, - "configs": resolved.configs, - "task_id": result.task_id, - "seed": seed, + "policy_spec": result.policy_spec, + "scene_config": result.scene_config, "inference_device": str(device), - "simulation_device": str(simulation_device), - "renderer": renderer, "embodichain_version": __version__, "dexsim_version": getattr(dexsim, "__version__", None), "dexsim_commit": getattr(dexsim, "__commit_id__", None), }, "result": { - "reason": result.reason, - "simulation_time": result.simulation_time, - "simulation_steps": result.simulation_steps, - "control_steps": result.control_steps, - "physics_backend": result.physics_backend, - "requested_duration": result.requested_duration, - "effective_duration": result.effective_duration, "episodes": result.episodes, - "metrics": result.metrics, + "summary": result.summary, }, } diff --git a/embodichain/learning/rl/motion_policy_evaluation/manifest.py b/embodichain/learning/rl/policy_evaluation/manifest.py similarity index 88% rename from embodichain/learning/rl/motion_policy_evaluation/manifest.py rename to embodichain/learning/rl/policy_evaluation/manifest.py index 19b9c5f7f..cc8baa75d 100644 --- a/embodichain/learning/rl/motion_policy_evaluation/manifest.py +++ b/embodichain/learning/rl/policy_evaluation/manifest.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Minimal index connecting a training run to motion-policy evaluation.""" +"""Index a training run for standalone policy evaluation.""" from __future__ import annotations @@ -23,7 +23,6 @@ from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path -from types import MappingProxyType from typing import Any __all__ = ["RUN_MANIFEST_NAME", "RunManifest", "write_run_manifest"] @@ -36,18 +35,13 @@ class RunManifest: """Resolved paths from one EmbodiChain training run.""" root: Path - motion_profile: str | None configs: Mapping[str, Path] checkpoints: Mapping[str, Path | None] def __post_init__(self) -> None: object.__setattr__(self, "root", Path(self.root).resolve()) - object.__setattr__(self, "configs", MappingProxyType(dict(self.configs))) - object.__setattr__( - self, - "checkpoints", - MappingProxyType(dict(self.checkpoints)), - ) + object.__setattr__(self, "configs", dict(self.configs)) + object.__setattr__(self, "checkpoints", dict(self.checkpoints)) @classmethod def load(cls, run: str | Path) -> RunManifest: @@ -73,12 +67,9 @@ def load(cls, run: str | Path) -> RunManifest: "checkpoints", allow_none=True, ) - profile = value.get("motion_profile") - if profile is not None and not isinstance(profile, str): - raise TypeError("Run manifest motion_profile must be a string or null") - return cls(root, profile, configs, checkpoints) + return cls(root, configs, checkpoints) - def select_checkpoint(self, requested: str = "best") -> tuple[str, Path]: + def select_checkpoint(self, requested: str = "latest") -> tuple[str, Path]: """Select ``best`` or ``latest`` and return its resolved path. Args: @@ -109,7 +100,6 @@ def write_run_manifest( latest_checkpoint: str | Path, best_checkpoint: str | Path | None = None, gym_config: str | Path | None = None, - motion_profile: str | None = None, ) -> Path: """Snapshot training configs and write the minimal run manifest. @@ -119,7 +109,6 @@ def write_run_manifest( latest_checkpoint: Final saved checkpoint. best_checkpoint: Best checkpoint when evaluation selected one. gym_config: Referenced task config when the trainer uses one. - motion_profile: Default Motion Profile for visual evaluation. Returns: Written manifest path. @@ -139,7 +128,6 @@ def write_run_manifest( } value: dict[str, Any] = { "schema_version": 1, - "motion_profile": motion_profile, "configs": configs, "checkpoints": checkpoints, } diff --git a/embodichain/learning/rl/motion_policy_evaluation/profile.py b/embodichain/learning/rl/policy_evaluation/profile.py similarity index 81% rename from embodichain/learning/rl/motion_policy_evaluation/profile.py rename to embodichain/learning/rl/policy_evaluation/profile.py index ed107c885..a215a084b 100644 --- a/embodichain/learning/rl/motion_policy_evaluation/profile.py +++ b/embodichain/learning/rl/policy_evaluation/profile.py @@ -14,14 +14,13 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Motion Profile providers for task-specific model and control semantics.""" +"""External Policy Profile registration and construction.""" from __future__ import annotations from collections.abc import Callable, Mapping from dataclasses import dataclass, field from pathlib import Path -from types import MappingProxyType from typing import Any import torch @@ -30,7 +29,6 @@ "MotionProfile", "MotionProfileRequest", "build_motion_profile", - "get_motion_profile_names", "register_motion_profile", ] @@ -64,7 +62,7 @@ def __post_init__(self) -> None: else Path(self.resource_root).expanduser().resolve() ) object.__setattr__(self, "checkpoint", checkpoint) - object.__setattr__(self, "configs", MappingProxyType(configs)) + object.__setattr__(self, "configs", configs) object.__setattr__(self, "resource_root", root) @@ -73,27 +71,14 @@ class MotionProfile: """DexSim Policy Spec and report metadata built by one provider.""" profile_id: str - checkpoint: Path policy_spec: Mapping[str, Any] provider_version: int = 1 provenance: Mapping[str, Any] = field(default_factory=dict) warnings: tuple[str, ...] = () def __post_init__(self) -> None: - checkpoint = Path(self.checkpoint).expanduser().resolve() - if not checkpoint.is_file(): - raise FileNotFoundError(f"Motion checkpoint does not exist: {checkpoint}") - object.__setattr__(self, "checkpoint", checkpoint) - object.__setattr__( - self, - "policy_spec", - MappingProxyType(dict(self.policy_spec)), - ) - object.__setattr__( - self, - "provenance", - MappingProxyType(dict(self.provenance)), - ) + object.__setattr__(self, "policy_spec", dict(self.policy_spec)) + object.__setattr__(self, "provenance", dict(self.provenance)) object.__setattr__(self, "warnings", tuple(self.warnings)) @@ -115,11 +100,6 @@ def register_motion_profile(name: str, provider: MotionProfileProvider) -> None: _PROVIDERS[name] = provider -def get_motion_profile_names() -> tuple[str, ...]: - """Return the registered profile names.""" - return tuple(sorted(_PROVIDERS)) - - def build_motion_profile( name: str, request: MotionProfileRequest, diff --git a/embodichain/learning/rl/motion_policy_evaluation/report.py b/embodichain/learning/rl/policy_evaluation/report.py similarity index 95% rename from embodichain/learning/rl/motion_policy_evaluation/report.py rename to embodichain/learning/rl/policy_evaluation/report.py index 02060fea0..545b41f2e 100644 --- a/embodichain/learning/rl/motion_policy_evaluation/report.py +++ b/embodichain/learning/rl/policy_evaluation/report.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Write one timestamped motion-policy evaluation report.""" +"""Write timestamped policy evaluation reports.""" from __future__ import annotations @@ -44,7 +44,7 @@ def write_evaluation_report( output = Path(parent).expanduser().resolve() output.mkdir(parents=True, exist_ok=True) stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") - directory = output / f"{stamp}-motion-policy" + directory = output / f"{stamp}-policy" directory.mkdir() report = { "schema_version": 1, diff --git a/embodichain/learning/rl/motion_policy_evaluation/native_task.py b/embodichain/learning/rl/policy_evaluation/viewer.py similarity index 83% rename from embodichain/learning/rl/motion_policy_evaluation/native_task.py rename to embodichain/learning/rl/policy_evaluation/viewer.py index 87151e765..509176684 100644 --- a/embodichain/learning/rl/motion_policy_evaluation/native_task.py +++ b/embodichain/learning/rl/policy_evaluation/viewer.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Connect an EmbodiChain RL task to DexSim Motion Policy Evaluator.""" +"""Connect an EmbodiChain task Viewer to Motion Policy Evaluator.""" from __future__ import annotations @@ -24,7 +24,6 @@ from dataclasses import dataclass from typing import Any -import numpy as np import torch from dexsim.kit.motion_policy import ( @@ -45,37 +44,34 @@ __all__ = [ "EmbodiChainTaskEnvironment", "EmbodiChainTaskPolicyAdapter", - "NativeTaskEvaluationResult", - "evaluate_native_task", + "NativeViewerResult", + "evaluate_native_viewer", ] _MISSING = object() @dataclass(frozen=True) -class NativeTaskEvaluationResult: - """Result of evaluating one Policy in its original EmbodiChain task.""" +class NativeViewerResult: + """Result of visualizing one Policy in its EmbodiChain task.""" task_id: str reason: str simulation_time: float simulation_steps: int control_steps: int - physics_backend: str effective_duration: float requested_duration: float | None episodes: tuple[Mapping[str, float | int | bool | str], ...] metrics: Mapping[str, float] - viewer: bool class EmbodiChainTaskPolicyAdapter: """Run an EmbodiChain Policy from the task observation in each frame.""" - def __init__(self, policy: torch.nn.Module, device: torch.device, num_envs: int): + def __init__(self, policy: torch.nn.Module, device: torch.device): self.policy = policy self.device = device - self.num_envs = num_envs self._previous_training = policy.training def setup(self, context: PolicyContext) -> None: @@ -97,7 +93,7 @@ def infer(self, frame: EvaluationFrame) -> PolicyOutput: self.policy, frame.observation, device=self.device, - num_envs=self.num_envs, + num_envs=1, ) return PolicyOutput(action=action) @@ -118,22 +114,19 @@ def __init__( env: Any, *, seed: int, - viewer: bool, ) -> None: if int(env.num_envs) != 1: raise ValueError("Visual task evaluation requires num_envs=1") self.env = env self._base_env = getattr(env, "unwrapped", env) - self._viewer = viewer - if viewer: - world = self._world() - if world is None or not world.is_window_initialized(): - raise ValueError( - "Viewer evaluation requires an EmbodiChain simulator task " - "with an initialized window" - ) + world = self._world() + if world is None or not world.is_window_initialized(): + raise ValueError( + "Viewer evaluation requires an initialized simulator window" + ) self._seed = seed self._first_reset = True + self._reset_key_down = False self._control_step = 0 self._frame: EvaluationFrame | None = None self._episode_return = 0.0 @@ -162,8 +155,6 @@ def physics_backend(self) -> str: @property def viewer_is_open(self) -> bool: """Return whether the original task Viewer remains open.""" - if not self._viewer: - return False world = self._world() return bool(world is not None and world.is_window_initialized()) @@ -177,15 +168,11 @@ def current_frame(self) -> EvaluationFrame: @property def episodes(self) -> tuple[Mapping[str, float | int | bool | str], ...]: """Return completed episode summaries.""" - return tuple(dict(episode) for episode in self._episodes) + return tuple(self._episodes) def open_viewer(self, title: str) -> None: """Apply the evaluation title to the task Viewer.""" - if not self._viewer: - return - world = self._world() - if world is not None and world.is_window_initialized(): - world.get_windows().set_window_title(title) + self._world().get_windows().set_window_title(title) def reset(self) -> EvaluationFrame: """Run the task's original reset and return its observation.""" @@ -200,8 +187,6 @@ def reset(self) -> EvaluationFrame: def poll(self) -> str | None: """Report when the native Viewer is closed or Escape is pressed.""" - if not self._viewer: - return None world = self._world() if world is None or not world.is_window_initialized(): return "viewer closed" @@ -210,6 +195,11 @@ def poll(self) -> str | None: native = world.get_windows().native() if native.key_state(InputKey.SCANCODE_ESCAPE): return "viewer closed" + reset_down = bool(native.key_state(InputKey.SCANCODE_BACKSPACE)) + reset_pressed = reset_down and not self._reset_key_down + self._reset_key_down = reset_down + if reset_pressed: + return "manual reset" return None def step(self, action: object) -> EnvironmentStep: @@ -246,10 +236,9 @@ def step(self, action: object) -> EnvironmentStep: "success": success, } ) - if self._viewer: - remaining = self._policy_context.policy_dt - (time.perf_counter() - started) - if remaining > 0.0: - time.sleep(remaining) + remaining = self._policy_context.policy_dt - (time.perf_counter() - started) + if remaining > 0.0: + time.sleep(remaining) return EnvironmentStep( frame=self._frame, termination_reason=reason, @@ -260,23 +249,21 @@ def metrics(self) -> dict[str, float]: """Return task metrics and completed episode aggregates.""" result = dict(self._reported_metrics) if self._episodes: + count = len(self._episodes) result.update( { - "eval/avg_reward": float( - np.mean( - [float(episode["reward"]) for episode in self._episodes] - ) - ), - "eval/avg_length": float( - np.mean( - [float(episode["length"]) for episode in self._episodes] - ) - ), - "eval/success_rate": float( - np.mean( - [bool(episode["success"]) for episode in self._episodes] - ) - ), + "eval/avg_reward": sum( + float(episode["reward"]) for episode in self._episodes + ) + / count, + "eval/avg_length": sum( + float(episode["length"]) for episode in self._episodes + ) + / count, + "eval/success_rate": sum( + bool(episode["success"]) for episode in self._episodes + ) + / count, } ) return result @@ -333,17 +320,16 @@ def _world(self) -> Any | None: return None if sim is None else sim.get_world() -def evaluate_native_task( +def evaluate_native_viewer( runtime: PolicyRuntime, *, seed: int, - viewer: bool, episodes: int | None, control_steps: int | None, duration: float | None, termination_behavior: str = "auto_reset", -) -> NativeTaskEvaluationResult: - """Evaluate an EmbodiChain Policy in the task used for training.""" +) -> NativeViewerResult: + """Visualize an EmbodiChain Policy in the task used for training.""" if episodes is not None and episodes <= 0: raise ValueError("episodes must be positive") if control_steps is not None and control_steps <= 0: @@ -355,31 +341,26 @@ def evaluate_native_task( if termination_behavior == "continue": raise ValueError("Native task evaluation supports pause or auto_reset") + environment = None + adapter = None + evaluator = None try: environment = EmbodiChainTaskEnvironment( runtime.env, seed=seed, - viewer=viewer, ) - except Exception: - runtime.env.close() - raise - adapter = None - evaluator = None - try: adapter = EmbodiChainTaskPolicyAdapter( runtime.policy, runtime.device, - num_envs=1, ) if duration is not None: control_steps = math.ceil( duration / environment.policy_context.policy_dt - 1e-12 ) total_steps = 0 - reason = "viewer closed" if viewer else "episode target reached" + reason = "viewer closed" options = RunOptions( - headless=not viewer, + headless=False, termination_behavior=( "continue" if termination_behavior == "auto_reset" else "pause" ), @@ -415,28 +396,27 @@ def evaluate_native_task( episode_results = environment.episodes metrics = environment.metrics() context = environment.policy_context - backend = environment.physics_backend finally: - if evaluator is None: + if evaluator is not None: + evaluator.close() + elif environment is not None: if adapter is not None: adapter.close() environment.close() else: - evaluator.close() + runtime.close() simulation_steps = total_steps * context.sim_steps_per_control - return NativeTaskEvaluationResult( + return NativeViewerResult( task_id=runtime.env_id, reason=reason, simulation_time=simulation_steps * context.physics_dt, simulation_steps=simulation_steps, control_steps=total_steps, - physics_backend=backend, effective_duration=total_steps * context.policy_dt, requested_duration=duration, episodes=episode_results, metrics=metrics, - viewer=viewer, ) @@ -491,20 +471,10 @@ def _step_metrics(info: object, reward: float) -> dict[str, float]: def _policy_context_from_env(env: Any) -> PolicyContext: - """Read timing from a simulator task or lightweight learning task.""" - if hasattr(env, "physics_dt") and hasattr(env, "step_dt"): - return PolicyContext( - robot=None, - physics_dt=float(env.physics_dt), - sim_steps_per_control=int(env.cfg.sim_steps_per_control), - policy_dt=float(env.step_dt), - ) - if not hasattr(env, "dt"): - raise ValueError("Lightweight task evaluation requires an env.dt value") - policy_dt = float(env.dt) + """Read timing from the simulator task.""" return PolicyContext( robot=None, - physics_dt=policy_dt, - sim_steps_per_control=1, - policy_dt=policy_dt, + physics_dt=float(env.physics_dt), + sim_steps_per_control=int(env.cfg.sim_steps_per_control), + policy_dt=float(env.step_dt), ) diff --git a/embodichain/learning/rl/runtime.py b/embodichain/learning/rl/runtime.py index d70fb07f3..ed8f0c7f7 100644 --- a/embodichain/learning/rl/runtime.py +++ b/embodichain/learning/rl/runtime.py @@ -18,12 +18,10 @@ from __future__ import annotations -from collections.abc import Callable from dataclasses import dataclass from pathlib import Path from typing import Any -import numpy as np import torch from embodichain.lab.gym.utils.gym_utils import config_to_cfg, get_manager_modules @@ -37,24 +35,14 @@ from embodichain.utils.utility import load_config __all__ = [ - "GymEnvironmentRuntime", - "GymPolicyFactory", "PolicyRuntime", - "build_gym_environment", "build_gym_policy_runtime", - "build_learning_environment", "build_learning_policy_runtime", - "resolve_config_reference", - "resolve_torch_device", - "seed_policy_runtime", ] -GymPolicyFactory = Callable[[int, int, torch.device], torch.nn.Module] -"""Build an evaluation Policy after the task dimensions are known.""" - @dataclass(frozen=True) -class GymEnvironmentRuntime: +class _GymEnvironmentRuntime: """A simulator task reconstructed from one training configuration.""" env: Any @@ -76,46 +64,12 @@ class PolicyRuntime: gym_config: dict[str, Any] | None = None gym_config_path: Path | None = None + def close(self) -> None: + """Close the Environment without terminating the current process.""" + _close_environment(self.env) -def resolve_torch_device(value: str | torch.device) -> torch.device: - """Validate and activate one CPU or CUDA device.""" - if not isinstance(value, (str, torch.device)): - raise TypeError("device must be a string or torch.device") - try: - device = torch.device(value) - except RuntimeError as error: - raise ValueError(f"Failed to parse device {value!r}: {error}") from error - - if device.type == "cuda": - if not torch.cuda.is_available(): - raise ValueError( - "CUDA was requested but torch.cuda.is_available() is False." - ) - index = device.index - if index is None: - index = torch.cuda.current_device() - if index < 0 or index >= torch.cuda.device_count(): - raise ValueError( - f"CUDA device index {index} is out of range " - f"(available devices: {torch.cuda.device_count()})." - ) - torch.cuda.set_device(index) - return torch.device(f"cuda:{index}") - if device.type != "cpu": - raise ValueError(f"Unsupported device type: {device.type}") - return torch.device("cpu") - - -def seed_policy_runtime(seed: int, device: torch.device) -> None: - """Seed NumPy and PyTorch for one Policy runtime.""" - np.random.seed(seed) - torch.manual_seed(seed) - torch.backends.cudnn.deterministic = True - if device.type == "cuda": - torch.cuda.manual_seed_all(seed) - -def resolve_config_reference( +def _resolve_config_reference( value: str | Path, *, base_dir: str | Path | None = None, @@ -131,7 +85,7 @@ def resolve_config_reference( return path -def build_learning_environment( +def _build_learning_environment( config: dict[str, Any], *, device: torch.device, @@ -141,15 +95,15 @@ def build_learning_environment( env_block = config["trainer"]["learning_env"] if isinstance(env_block, str): env_name = env_block - env_cfg: dict[str, Any] = {} + env_config: dict[str, Any] = {} else: env_name = env_block["name"] - env_cfg = dict(env_block.get("cfg", {})) + env_config = dict(env_block.get("cfg", {})) return str(env_name), build_learning_env( - env_name, + str(env_name), num_envs=num_envs, device=device, - **env_cfg, + **env_config, ) @@ -160,45 +114,20 @@ def build_learning_policy_runtime( num_envs: int, ) -> PolicyRuntime: """Build a lightweight Environment and its configured Policy.""" - env_name, env = build_learning_environment( + env_name, env = _build_learning_environment( config, - num_envs=num_envs, device=device, + num_envs=num_envs, ) try: - observation_dim = int(env.single_observation_space.shape[-1]) - action_dim = int(env.single_action_space.shape[-1]) - policy_block = config["policy"] - actor_cfg = policy_block.get("actor") - critic_cfg = policy_block.get("critic") - actor = ( - build_mlp_from_cfg(actor_cfg, observation_dim, action_dim) - if actor_cfg is not None - else None - ) - critic = ( - build_mlp_from_cfg(critic_cfg, observation_dim, 1) - if critic_cfg is not None - else None - ) - policy = build_policy( - policy_block, - env.single_observation_space, - env.single_action_space, - device, - actor=actor, - critic=critic, - ) - if "initial_log_std" in policy_block and hasattr(policy, "log_std"): - with torch.no_grad(): - policy.log_std.fill_(float(policy_block["initial_log_std"])) + policy = _build_learning_policy(config["policy"], env, device) except Exception: env.close() raise return PolicyRuntime(env, policy, device, env_name) -def build_gym_environment( +def _build_gym_environment( config: dict[str, Any], *, simulation_device: torch.device, @@ -208,10 +137,10 @@ def build_gym_environment( gpu_id: int, config_dir: str | Path | None = None, profiler: EnvProfilerCfg | None = None, -) -> GymEnvironmentRuntime: +) -> _GymEnvironmentRuntime: """Build the simulator Environment declared by a training config.""" trainer_cfg = config["trainer"] - gym_config_path = resolve_config_reference( + gym_config_path = _resolve_config_reference( trainer_cfg["gym_config"], base_dir=config_dir, ) @@ -231,7 +160,7 @@ def build_gym_environment( ) env_cfg.profiler = profiler env = build_env(gym_config["id"], base_env_cfg=env_cfg) - return GymEnvironmentRuntime( + return _GymEnvironmentRuntime( env=env, env_id=str(gym_config["id"]), env_cfg=env_cfg, @@ -250,11 +179,10 @@ def build_gym_policy_runtime( gpu_id: int, config_dir: str | Path | None = None, profiler: EnvProfilerCfg | None = None, - policy_factory: GymPolicyFactory | None = None, simulation_device: torch.device | None = None, ) -> PolicyRuntime: """Build a simulator task and the Policy declared by its training config.""" - task = build_gym_environment( + task = _build_gym_environment( config, simulation_device=simulation_device or device, num_envs=num_envs, @@ -275,20 +203,15 @@ def build_gym_policy_runtime( if action_manager is not None else len(env.get_wrapper_attr("active_joint_ids")) ) - if policy_factory is None: - policy = _build_gym_policy( - config["policy"], - env=env, - device=device, - observation_dim=observation_dim, - action_dim=environment_action_dim, - ) - else: - policy = policy_factory(observation_dim, environment_action_dim, device) - if not isinstance(policy, torch.nn.Module): - raise TypeError("policy_factory must return a torch.nn.Module.") + policy = _build_gym_policy( + config["policy"], + env=env, + device=device, + observation_dim=observation_dim, + action_dim=environment_action_dim, + ) except Exception: - env.close() + _close_environment(env) raise return PolicyRuntime( env=env, @@ -348,3 +271,42 @@ def _build_gym_policy( env.action_space, device, ) + + +def _build_learning_policy( + policy_block: dict[str, Any], + env: Any, + device: torch.device, +) -> torch.nn.Module: + observation_dim = int(env.single_observation_space.shape[-1]) + action_dim = int(env.single_action_space.shape[-1]) + actor_cfg = policy_block.get("actor") + critic_cfg = policy_block.get("critic") + policy = build_policy( + policy_block, + env.single_observation_space, + env.single_action_space, + device, + actor=( + build_mlp_from_cfg(actor_cfg, observation_dim, action_dim) + if actor_cfg is not None + else None + ), + critic=( + build_mlp_from_cfg(critic_cfg, observation_dim, 1) + if critic_cfg is not None + else None + ), + ) + if "initial_log_std" in policy_block and hasattr(policy, "log_std"): + with torch.no_grad(): + policy.log_std.fill_(float(policy_block["initial_log_std"])) + return policy + + +def _close_environment(env: Any) -> None: + target = getattr(env, "unwrapped", env) + if getattr(target, "sim", None) is not None: + target.close(exit_process=False) + else: + env.close() diff --git a/embodichain/learning/rl/train.py b/embodichain/learning/rl/train.py index cb93eeef0..9556efd31 100644 --- a/embodichain/learning/rl/train.py +++ b/embodichain/learning/rl/train.py @@ -20,6 +20,7 @@ import os import time from collections.abc import Sequence +from copy import deepcopy from pathlib import Path import numpy as np @@ -28,7 +29,7 @@ from torch.utils.tensorboard import SummaryWriter from embodichain.learning.rl.models import get_registered_policy_names -from embodichain.learning.rl.motion_policy_evaluation.manifest import ( +from embodichain.learning.rl.policy_evaluation.manifest import ( write_run_manifest, ) from embodichain.learning.rl.algo import ( @@ -40,9 +41,8 @@ DifferentiableTrainer, DifferentiableTrainerCfg, ) -from embodichain.learning.rl.env import build_learning_env from embodichain.learning.rl.runtime import ( - build_gym_environment, + _build_learning_environment, build_gym_policy_runtime, build_learning_policy_runtime, ) @@ -50,6 +50,7 @@ from embodichain.learning.rl.utils.trainer import Trainer from embodichain.utils import logger from embodichain.lab.gym.utils.registration import ( + build_env, discover_task_packages, execute_init_hooks, ) @@ -123,7 +124,7 @@ def _event_params( run_base: str | Path, phase: str, ) -> dict: - """Resolve event parameters that belong to one training run.""" + """Place default camera recordings under the current training run.""" params = dict(event_info.get("params", {})) function_name = str(event_info.get("func", "")).rsplit(".", 1)[-1] if function_name in _CAMERA_RECORDERS: @@ -180,16 +181,10 @@ def _train_learning_env( enable_eval = bool(trainer_cfg.get("enable_eval", False)) eval_env = None if enable_eval: - env_block = trainer_cfg["learning_env"] - if isinstance(env_block, str): - env_cfg = {} - else: - env_cfg = dict(env_block.get("cfg", {})) - eval_env = build_learning_env( - env_name, + _eval_name, eval_env = _build_learning_environment( + cfg_data, num_envs=int(trainer_cfg.get("num_eval_envs", 16)), device=device, - **env_cfg, ) algorithm = build_algo( @@ -276,13 +271,6 @@ def _train_learning_env( trainer.train(total_timesteps) trainer.save_checkpoint() summary = trainer.get_summary() - _write_motion_run_manifest( - run_base, - config_path, - trainer_cfg, - summary, - ) - return summary finally: writer.close() if use_wandb: @@ -290,6 +278,8 @@ def _train_learning_env( env.close() if eval_env is not None: eval_env.close() + _write_policy_run_manifest(run_base, config_path, summary) + return summary def train_from_config( @@ -477,16 +467,11 @@ def train_from_config( eval_env = None num_eval_envs = trainer_cfg.get("num_eval_envs", 4) if enable_eval and rank == 0: - eval_runtime = build_gym_environment( - cfg_data, - simulation_device=device, - num_envs=int(num_eval_envs), - headless=True, - renderer=renderer, - gpu_id=gpu_id, - config_dir=Path(config_path).expanduser().resolve().parent, - ) - eval_env = eval_runtime.env + eval_gym_env_cfg = deepcopy(gym_env_cfg) + eval_gym_env_cfg.num_envs = int(num_eval_envs) + eval_gym_env_cfg.sim_cfg.headless = True + eval_gym_env_cfg.profiler = None + eval_env = build_env(gym_config_data["id"], base_env_cfg=eval_gym_env_cfg) logger.log_info( f"Evaluation environment created (num_envs={num_eval_envs}, headless=True)" ) @@ -589,6 +574,7 @@ def train_from_config( f"Total steps: {total_steps} (iterations≈{iterations}, world_size={world_size})" ) + summary = None try: trainer.train(total_steps) except KeyboardInterrupt: @@ -596,6 +582,8 @@ def train_from_config( logger.log_info("Training interrupted by user") finally: trainer.save_checkpoint() + if rank == 0: + summary = trainer.get_summary() if writer is not None: writer.close() if use_wandb and rank == 0: @@ -622,26 +610,25 @@ def train_from_config( if distributed and torch.distributed.is_initialized(): torch.distributed.destroy_process_group() - if rank == 0: - _write_motion_run_manifest( + if summary is not None: + _write_policy_run_manifest( run_base, config_path, - trainer_cfg, - trainer.get_summary(), + summary, gym_config=gym_config_path, ) + if rank == 0: logger.log_info("Training finished") -def _write_motion_run_manifest( +def _write_policy_run_manifest( run_base: str | Path, config_path: str | Path, - trainer_cfg: dict, summary: dict, *, gym_config: str | Path | None = None, ) -> Path: - """Write the final checkpoint and config index for motion evaluation.""" + """Write the checkpoint and configuration index for policy evaluation.""" latest = summary.get("latest_checkpoint_path") if latest is None: raise RuntimeError("Training finished without a checkpoint") @@ -651,7 +638,6 @@ def _write_motion_run_manifest( gym_config=gym_config, latest_checkpoint=latest, best_checkpoint=summary.get("best_checkpoint_path"), - motion_profile=trainer_cfg.get("motion_profile"), ) diff --git a/examples/learning/motion_policy_evaluation/README.md b/examples/learning/policy_evaluation/README.md similarity index 92% rename from examples/learning/motion_policy_evaluation/README.md rename to examples/learning/policy_evaluation/README.md index 3e2fe1a97..8e3155b00 100644 --- a/examples/learning/motion_policy_evaluation/README.md +++ b/examples/learning/policy_evaluation/README.md @@ -16,7 +16,7 @@ target processing. ## Directory layout ```text -motion_policy_evaluation/ +policy_evaluation/ ├── README.md ├── prepare_resources.py ├── eval_policy.py # Register the local Profile and run the example @@ -47,8 +47,8 @@ prints the model, asset, checkout, and digest verification progress. Re-running the command continues an existing Git checkout after an interrupted download. ```bash -python examples/learning/motion_policy_evaluation/prepare_resources.py -python examples/learning/motion_policy_evaluation/eval_policy.py \ +python examples/learning/policy_evaluation/prepare_resources.py +python examples/learning/policy_evaluation/eval_policy.py \ --viewer \ --renderer hybrid ``` @@ -74,7 +74,7 @@ The terminal prints the path to `evaluation.json` when the Viewer closes. Run a Headless smoke test with: ```bash -python examples/learning/motion_policy_evaluation/eval_policy.py \ +python examples/learning/policy_evaluation/eval_policy.py \ --device cpu \ --sim-device cpu \ --control-steps 20 @@ -86,18 +86,18 @@ current process. Run the script directly from the repository root. To use another cache directory: ```bash -python examples/learning/motion_policy_evaluation/prepare_resources.py \ +python examples/learning/policy_evaluation/prepare_resources.py \ --output /tmp/anymal_c_velocity ANYMAL_C_EXAMPLE_CACHE=/tmp/anymal_c_velocity \ - python examples/learning/motion_policy_evaluation/eval_policy.py --viewer + python examples/learning/policy_evaluation/eval_policy.py --viewer ``` ## Execution pipeline ```mermaid flowchart LR - CLI[eval-motion-policy] --> Profile[build_profile] + CLI[eval-policy] --> Profile[build_profile] Profile --> Spec[Policy Spec
assets, control parameters, frequency] Spec --> Setup[Adapter.setup
load TorchScript and joint mapping] Setup --> State[read RobotState] diff --git a/examples/learning/motion_policy_evaluation/anymal_c/__init__.py b/examples/learning/policy_evaluation/anymal_c/__init__.py similarity index 87% rename from examples/learning/motion_policy_evaluation/anymal_c/__init__.py rename to examples/learning/policy_evaluation/anymal_c/__init__.py index e18fe4904..0303a1211 100644 --- a/examples/learning/motion_policy_evaluation/anymal_c/__init__.py +++ b/examples/learning/policy_evaluation/anymal_c/__init__.py @@ -14,11 +14,11 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Register the concrete ANYmal-C velocity Motion Profile.""" +"""Register the ANYmal-C velocity Motion Profile.""" from __future__ import annotations -from embodichain.learning.rl.motion_policy_evaluation import register_motion_profile +from embodichain.learning.rl.policy_evaluation import register_motion_profile from .profile import PROFILE_ID, build_profile diff --git a/examples/learning/motion_policy_evaluation/anymal_c/profile.py b/examples/learning/policy_evaluation/anymal_c/profile.py similarity index 88% rename from examples/learning/motion_policy_evaluation/anymal_c/profile.py rename to examples/learning/policy_evaluation/anymal_c/profile.py index edea62d0b..2f0e89e73 100644 --- a/examples/learning/motion_policy_evaluation/anymal_c/profile.py +++ b/examples/learning/policy_evaluation/anymal_c/profile.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""ANYmal-C velocity Profile for Newton's public TorchScript policy.""" +"""ANYmal-C velocity Profile for a public TorchScript policy.""" from __future__ import annotations @@ -32,7 +32,7 @@ require_finite, ) -from embodichain.learning.rl.motion_policy_evaluation import ( +from embodichain.learning.rl.policy_evaluation import ( MotionProfile, MotionProfileRequest, ) @@ -73,8 +73,6 @@ def build_profile(request: MotionProfileRequest) -> MotionProfile: Returns: A Motion Profile ready for DexSim Motion Policy Kit. """ - if request.checkpoint.suffix.lower() != ".pt": - raise ValueError("The ANYmal-C example requires a TorchScript .pt file") if request.resource_root is None: raise ValueError( "The ANYmal-C example requires --resource-root from prepare_resources.py" @@ -88,7 +86,6 @@ def build_profile(request: MotionProfileRequest) -> MotionProfile: return MotionProfile( profile_id=PROFILE_ID, - checkpoint=request.checkpoint, policy_spec={ "schema_version": 1, "kind": "policy", @@ -173,8 +170,6 @@ def __init__(self, request: AdapterRequest) -> None: config = dict(request.config) self.device = torch.device(str(config["inference_device"])) self.joint_names = tuple(config["joint_names"]) - if self.joint_names != _JOINT_NAMES: - raise ValueError("ANYmal-C joint order is incompatible") self.checkpoint = request.models["actor"] self.previous_action = np.zeros(12, dtype=np.float32) self.joints: JointMap | None = None @@ -182,10 +177,11 @@ def __init__(self, request: AdapterRequest) -> None: def setup(self, context: PolicyContext) -> None: """Load the model and bind the runtime joint order.""" - if context.robot is None: + robot = context.robot + if robot is None: raise RuntimeError("ANYmal-C robot description is required") self.joints = JointMap.from_joint_names( - context.robot.joint_names, + robot.joint_names, self.joint_names, ) self.model = torch.jit.load( @@ -206,17 +202,10 @@ def reset(self, frame: EvaluationFrame) -> None: def infer(self, frame: EvaluationFrame) -> PolicyOutput: """Build one 48-D observation and return 12 joint targets.""" - if frame.robot_state is None: - raise RuntimeError("ANYmal-C robot state is required") observation = self._build_observation(frame) tensor = torch.from_numpy(observation).to(self.device).unsqueeze(0) with torch.inference_mode(): output = self._model()(tensor) - if not isinstance(output, torch.Tensor) or tuple(output.shape) != (1, 12): - shape = ( - None if not isinstance(output, torch.Tensor) else tuple(output.shape) - ) - raise ValueError(f"ANYmal-C policy output must be (1, 12), got {shape}") action = require_finite( "ANYmal-C action", output[0].detach().cpu().numpy(), @@ -264,10 +253,6 @@ def _build_observation(self, frame: EvaluationFrame) -> np.ndarray: ), dtype=np.float32, ) - if observation.shape != (48,): - raise ValueError( - f"ANYmal-C observation must have shape (48,), got {observation.shape}" - ) return require_finite("ANYmal-C observation", observation) def _joints(self) -> JointMap: @@ -283,8 +268,6 @@ def _model(self) -> torch.jit.ScriptModule: def _fall_reason(root_pose: np.ndarray) -> str | None: pose = np.asarray(root_pose, dtype=np.float64) - if pose.shape != (4, 4) or not np.all(np.isfinite(pose)): - raise ValueError("ANYmal-C root pose must be a finite 4x4 matrix") height = float(pose[2, 3]) tilt = math.acos(float(np.clip(pose[2, 2], -1.0, 1.0))) reasons = [] diff --git a/examples/learning/motion_policy_evaluation/eval_policy.py b/examples/learning/policy_evaluation/eval_policy.py similarity index 90% rename from examples/learning/motion_policy_evaluation/eval_policy.py rename to examples/learning/policy_evaluation/eval_policy.py index 1f5f69929..b01048b30 100644 --- a/examples/learning/motion_policy_evaluation/eval_policy.py +++ b/examples/learning/policy_evaluation/eval_policy.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Run the public ANYmal-C checkpoint directly from the example directory.""" +"""Evaluate the public ANYmal-C checkpoint from the example directory.""" from __future__ import annotations @@ -33,7 +33,7 @@ def example_arguments(argv: list[str]) -> list[str]: """Add the example Profile, checkpoint, and resource paths. Args: - argv: Evaluation options accepted by ``eval-motion-policy``. + argv: Evaluation options accepted by ``eval-policy``. Returns: Arguments ready for the EmbodiChain evaluation CLI. @@ -62,7 +62,7 @@ def main(argv: list[str] | None = None) -> None: sys.path.insert(0, str(REPOSITORY_ROOT)) from anymal_c import register - from embodichain.learning.rl.motion_policy_evaluation.cli import cli + from embodichain.learning.rl.policy_evaluation.cli import cli register() cli(example_arguments(sys.argv[1:] if argv is None else argv)) diff --git a/examples/learning/motion_policy_evaluation/prepare_resources.py b/examples/learning/policy_evaluation/prepare_resources.py similarity index 69% rename from examples/learning/motion_policy_evaluation/prepare_resources.py rename to examples/learning/policy_evaluation/prepare_resources.py index 703b41938..7e000e2a4 100644 --- a/examples/learning/motion_policy_evaluation/prepare_resources.py +++ b/examples/learning/policy_evaluation/prepare_resources.py @@ -14,7 +14,7 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Fetch the pinned public policy and robot assets for the ANYmal-C example.""" +"""Prepare the pinned public policy and assets for the ANYmal-C example.""" from __future__ import annotations @@ -34,18 +34,14 @@ ROBOT_LICENSE_RELATIVE_PATH = Path("anybotics_anymal_c/LICENSE") ROBOT_RELATIVE_PATH = Path("anybotics_anymal_c/urdf/anymal.urdf") MESH_RELATIVE_PATH = Path("anybotics_anymal_c/meshes/base.dae") -MODEL_SHA256 = "00765c1c07e497be3825672b05f9cefff9238f2df72fb0bcb5ac9541155b945f" -POLICY_CONFIG_SHA256 = ( - "b5a463ac418c7f40ebe494c7bcf0d8031f021db70a0625dbcc28a718de8ee817" -) -POLICY_LICENSE_SHA256 = ( - "59899c6091b540582ed617e8eeaac4919dc985ccfc35459ee9752b699be5205b" -) -ROBOT_LICENSE_SHA256 = ( - "cef384faae108293b03b5e16a00bc3db8212d44575f69df6296438a3f901700b" -) -ROBOT_SHA256 = "d6bd20292cdd4873ffdeeb6f8ca3f96c4a0096565d78d8b6204f6edf0d19fb83" -MESH_SHA256 = "785bea9b33831f8c741fc0ca070162e73cbf560ea9b03c53abf8978be877fc48" +SHA256 = { + MODEL_RELATIVE_PATH: "00765c1c07e497be3825672b05f9cefff9238f2df72fb0bcb5ac9541155b945f", + POLICY_CONFIG_RELATIVE_PATH: "b5a463ac418c7f40ebe494c7bcf0d8031f021db70a0625dbcc28a718de8ee817", + POLICY_LICENSE_RELATIVE_PATH: "59899c6091b540582ed617e8eeaac4919dc985ccfc35459ee9752b699be5205b", + ROBOT_LICENSE_RELATIVE_PATH: "cef384faae108293b03b5e16a00bc3db8212d44575f69df6296438a3f901700b", + ROBOT_RELATIVE_PATH: "d6bd20292cdd4873ffdeeb6f8ca3f96c4a0096565d78d8b6204f6edf0d19fb83", + MESH_RELATIVE_PATH: "785bea9b33831f8c741fc0ca070162e73cbf560ea9b03c53abf8978be877fc48", +} def prepare_resources(output: Path) -> tuple[Path, Path]: @@ -75,21 +71,8 @@ def prepare_resources(output: Path) -> tuple[Path, Path]: ) checkpoint = checkout / MODEL_RELATIVE_PATH - policy_config = checkout / POLICY_CONFIG_RELATIVE_PATH - policy_license = checkout / POLICY_LICENSE_RELATIVE_PATH - robot_license = checkout / ROBOT_LICENSE_RELATIVE_PATH - robot = checkout / ROBOT_RELATIVE_PATH - mesh = checkout / MESH_RELATIVE_PATH - if not checkpoint.is_file(): - raise FileNotFoundError(f"Downloaded checkpoint does not exist: {checkpoint}") - if not robot.is_file(): - raise FileNotFoundError(f"Downloaded robot asset does not exist: {robot}") - _verify_sha256(checkpoint, MODEL_SHA256) - _verify_sha256(policy_config, POLICY_CONFIG_SHA256) - _verify_sha256(policy_license, POLICY_LICENSE_SHA256) - _verify_sha256(robot_license, ROBOT_LICENSE_SHA256) - _verify_sha256(robot, ROBOT_SHA256) - _verify_sha256(mesh, MESH_SHA256) + for relative, digest in SHA256.items(): + _verify_sha256(checkout / relative, digest) _status("Resource verification completed") return checkpoint, checkout @@ -151,31 +134,18 @@ def _prepare_checkout( ) if checkout.exists(): - actual_revision = _git_output(checkout, "rev-parse", "HEAD") - if actual_revision == revision: + remote_url = _git_output(checkout, "remote", "get-url", "origin") + if remote_url != url: + raise RuntimeError( + f"Resource checkout uses an unexpected remote: {remote_url}" + ) + if _git_output(checkout, "rev-parse", "HEAD") == revision: tracked_changes = _git_output( - checkout, - "status", - "--porcelain", - "--untracked-files=no", + checkout, "status", "--porcelain", "--untracked-files=no" ) if tracked_changes == "": _status(f"Using cached revision {revision[:8]} from {checkout}") return - _status(f"Repairing interrupted checkout in {checkout}") - elif actual_revision is not None: - raise RuntimeError( - f"Existing checkout uses revision {actual_revision}; " - "choose another --output" - ) - else: - _status(f"Resuming interrupted checkout in {checkout}") - - remote_url = _git_output(checkout, "remote", "get-url", "origin") - if remote_url != url: - raise RuntimeError( - f"Incomplete checkout uses an unexpected remote: {remote_url}" - ) else: checkout.parent.mkdir(parents=True, exist_ok=True) checkout.mkdir() @@ -185,21 +155,17 @@ def _prepare_checkout( _git(checkout, "sparse-checkout", "init", "--no-cone") _git(checkout, "sparse-checkout", "set", *includes) - fetched_revision = _git_output(checkout, "rev-parse", "FETCH_HEAD") - if fetched_revision != revision: - _status(f"Fetching revision {revision[:8]} from {url}") - _git( - checkout, - "fetch", - "--progress", - "--filter=blob:none", - "--depth", - "1", - "origin", - revision, - ) - else: - _status(f"Using previously fetched revision {revision[:8]}") + _status(f"Fetching revision {revision[:8]} from {url}") + _git( + checkout, + "fetch", + "--progress", + "--filter=blob:none", + "--depth", + "1", + "origin", + revision, + ) _status(f"Checking out required files in {checkout}") _git( diff --git a/tests/learning/rl/motion_policy_evaluation/test_bridge.py b/tests/learning/rl/motion_policy_evaluation/test_bridge.py deleted file mode 100644 index 754a8467e..000000000 --- a/tests/learning/rl/motion_policy_evaluation/test_bridge.py +++ /dev/null @@ -1,159 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -from types import SimpleNamespace - -from embodichain.learning.rl.motion_policy_evaluation.bridge import ( - create_motion_profile_evaluator, - evaluate_motion_profile, -) -from embodichain.learning.rl.motion_policy_evaluation.profile import MotionProfile - - -def test_bridge_forwards_physics_backend_and_exact_control_steps( - tmp_path, - monkeypatch, -): - checkpoint = tmp_path / "policy.pt" - checkpoint.write_bytes(b"checkpoint") - profile = MotionProfile( - profile_id="example", - checkpoint=checkpoint, - policy_spec={"schema_version": 1}, - ) - options = [] - forwarded = [] - monkeypatch.setattr( - "embodichain.learning.rl.motion_policy_evaluation.bridge.parse_policy_spec", - lambda value: "parsed", - ) - monkeypatch.setattr( - "embodichain.learning.rl.motion_policy_evaluation.bridge.resolve_policy_spec", - lambda spec, resolver: "resolved", - ) - monkeypatch.setattr( - "embodichain.learning.rl.motion_policy_evaluation.bridge.policy_spec_to_dict", - lambda value: {"policy_id": "example"}, - ) - monkeypatch.setattr( - "embodichain.learning.rl.motion_policy_evaluation.bridge.scene_config_to_dict", - lambda value: {"style": "standard"}, - ) - - def run_policy(resolved, run_options, **kwargs): - options.append(run_options) - forwarded.append(kwargs) - return SimpleNamespace( - reason="control steps reached", - simulation_time=0.2, - simulation_steps=40, - control_steps=10, - physics_backend="default", - requested_duration=None, - effective_duration=0.2, - metrics={"tracking/error": 0.25}, - ) - - monkeypatch.setattr( - "embodichain.learning.rl.motion_policy_evaluation.bridge.run_motion_policy", - run_policy, - ) - - input_provider = object() - environment = object() - result = evaluate_motion_profile( - profile, - control_steps=10, - physics_backend="default", - input_provider=input_provider, - environment=environment, - ) - - assert options[0].control_steps == 10 - assert options[0].physics_backend == "default" - assert forwarded == [ - { - "environment": environment, - "input_provider": input_provider, - } - ] - assert result.episodes[0]["control_steps"] == 10 - assert result.episodes[0]["effective_duration"] == 0.2 - assert result.summary["metrics"]["tracking/error"] == 0.25 - - -def test_create_profile_evaluator_forwards_the_environment( - tmp_path, - monkeypatch, -): - checkpoint = tmp_path / "policy.pt" - checkpoint.write_bytes(b"checkpoint") - profile = MotionProfile( - profile_id="example", - checkpoint=checkpoint, - policy_spec={"schema_version": 1}, - ) - evaluator = object() - adapter = object() - environment = object() - calls = [] - monkeypatch.setattr( - "embodichain.learning.rl.motion_policy_evaluation.bridge.parse_policy_spec", - lambda value: "parsed", - ) - monkeypatch.setattr( - "embodichain.learning.rl.motion_policy_evaluation.bridge.resolve_policy_spec", - lambda spec, resolver: "resolved", - ) - monkeypatch.setattr( - "embodichain.learning.rl.motion_policy_evaluation.bridge.create_motion_policy_evaluator", - lambda resolved, options, **kwargs: calls.append((resolved, options, kwargs)) - or evaluator, - ) - - result = create_motion_profile_evaluator( - profile, - adapter=adapter, - environment=environment, - ) - - assert result is evaluator - assert calls == [ - ( - "resolved", - None, - {"adapter": adapter, "environment": environment}, - ), - ] - - -def test_prebuilt_environment_requires_one_episode(tmp_path): - checkpoint = tmp_path / "policy.pt" - checkpoint.write_bytes(b"checkpoint") - profile = MotionProfile( - profile_id="example", - checkpoint=checkpoint, - policy_spec={"schema_version": 1}, - ) - - try: - evaluate_motion_profile(profile, episodes=2, environment=object()) - except ValueError as error: - assert str(error) == "A prebuilt environment supports one episode" - else: - raise AssertionError("Expected multi-episode validation") diff --git a/tests/learning/rl/motion_policy_evaluation/test_checkpoint.py b/tests/learning/rl/motion_policy_evaluation/test_checkpoint.py deleted file mode 100644 index 19976302a..000000000 --- a/tests/learning/rl/motion_policy_evaluation/test_checkpoint.py +++ /dev/null @@ -1,30 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -import torch - -from embodichain.learning.rl.motion_policy_evaluation import load_policy_state_dict - - -def test_load_policy_state_dict_returns_embodichain_policy_weights(tmp_path): - checkpoint = tmp_path / "policy.pt" - torch.save({"policy": {"actor.weight": torch.ones(2, 3)}}, checkpoint) - - state = load_policy_state_dict(checkpoint) - - assert torch.equal(state["actor.weight"], torch.ones(2, 3)) diff --git a/tests/learning/rl/motion_policy_evaluation/test_cli.py b/tests/learning/rl/motion_policy_evaluation/test_cli.py deleted file mode 100644 index 676da67ff..000000000 --- a/tests/learning/rl/motion_policy_evaluation/test_cli.py +++ /dev/null @@ -1,199 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -import importlib - -import pytest -import torch - -from embodichain.learning.rl.motion_policy_evaluation.cli import ( - _resolve_input, - _validate_native_options, - parse_args, -) -from embodichain.learning.rl.motion_policy_evaluation.manifest import ( - write_run_manifest, -) - - -def test_run_uses_manifest_profile_and_latest_checkpoint(tmp_path): - run = tmp_path / "run" - checkpoint = run / "checkpoints" / "policy.pt" - checkpoint.parent.mkdir(parents=True) - checkpoint.write_bytes(b"checkpoint") - train = tmp_path / "train.yaml" - train.write_text("trainer: {}\n", encoding="utf-8") - write_run_manifest( - run, - train_config=train, - latest_checkpoint=checkpoint, - motion_profile="example-motion", - ) - - resolved = _resolve_input(parse_args((str(run),))) - - assert resolved.profile == "example-motion" - assert resolved.checkpoint == checkpoint - assert resolved.selected_checkpoint == "latest" - - -def test_original_task_overrides_manifest_profile(tmp_path): - run = tmp_path / "run" - checkpoint = run / "checkpoints" / "policy.pt" - checkpoint.parent.mkdir(parents=True) - checkpoint.write_bytes(b"checkpoint") - train = tmp_path / "train.yaml" - train.write_text("trainer: {}\n", encoding="utf-8") - write_run_manifest( - run, - train_config=train, - latest_checkpoint=checkpoint, - motion_profile="example-motion", - ) - - resolved = _resolve_input(parse_args((str(run), "--original-task"))) - - assert resolved.profile is None - assert resolved.configs["train"] == run / "configs" / "train.yaml" - - -def test_explicit_configs_override_manifest_configs(tmp_path): - run = tmp_path / "run" - checkpoint = run / "checkpoints" / "policy.pt" - checkpoint.parent.mkdir(parents=True) - checkpoint.write_bytes(b"checkpoint") - train = tmp_path / "train.yaml" - train.write_text("trainer: {}\n", encoding="utf-8") - gym = tmp_path / "gym.yaml" - gym.write_text("id: Original\n", encoding="utf-8") - write_run_manifest( - run, - train_config=train, - gym_config=gym, - latest_checkpoint=checkpoint, - ) - replacement_train = tmp_path / "replacement-train.yaml" - replacement_train.write_text("trainer: {}\n", encoding="utf-8") - replacement_gym = tmp_path / "replacement-gym.yaml" - replacement_gym.write_text("id: Replacement\n", encoding="utf-8") - - resolved = _resolve_input( - parse_args( - ( - str(run), - "--config", - str(replacement_train), - "--gym-config", - str(replacement_gym), - ) - ) - ) - - assert resolved.configs["train"] == replacement_train - assert resolved.configs["gym"] == replacement_gym - - -def test_viewer_defaults_to_hybrid_without_a_time_limit(): - args = parse_args( - ( - "--profile", - "example-motion", - "--checkpoint", - "policy.pt", - "--viewer", - ) - ) - - assert args.renderer == "hybrid" - assert args.physics_backend is None - assert args.control_steps is None - assert args.duration is None - - -def test_run_without_profile_selects_native_task_evaluation(tmp_path, monkeypatch): - cli_module = importlib.import_module( - "embodichain.learning.rl.motion_policy_evaluation.cli" - ) - run = tmp_path / "run" - checkpoint = run / "checkpoints" / "policy.pt" - checkpoint.parent.mkdir(parents=True) - torch.save({"model_state_dict": {}}, checkpoint) - train = tmp_path / "train.yaml" - train.write_text("trainer: {}\n", encoding="utf-8") - write_run_manifest( - run, - train_config=train, - latest_checkpoint=checkpoint, - ) - expected = tmp_path / "evaluation.json" - received = {} - - def fake_native(args, resolved): - received["args"] = args - received["resolved"] = resolved - return expected - - monkeypatch.setattr(cli_module, "discover_task_packages", lambda: None) - monkeypatch.setattr(cli_module, "execute_init_hooks", lambda: None) - monkeypatch.setattr(cli_module, "_run_native_task", fake_native) - - report = cli_module.run(parse_args((str(run), "--viewer"))) - - assert report == expected - assert received["resolved"].profile is None - assert received["resolved"].checkpoint == checkpoint - assert received["args"].viewer is True - - -def test_explicit_native_checkpoint_requires_training_config(tmp_path): - checkpoint = tmp_path / "policy.pt" - checkpoint.write_bytes(b"checkpoint") - - args = parse_args(("--checkpoint", str(checkpoint))) - - with pytest.raises( - ValueError, - match="--config is required for a native EmbodiChain checkpoint", - ): - _resolve_input(args) - - -@pytest.mark.parametrize( - "arguments, option", - [ - (("--command", "0.5"), "--command"), - (("--physics-backend", "default"), "--physics-backend"), - (("--scene-config", "classic"), "--scene-config"), - (("--cache-dir", "cache"), "--cache-dir"), - (("--offline",), "--offline"), - (("--resource-root", "resources"), "--resource-root"), - ], -) -def test_native_task_rejects_profile_only_options(arguments, option): - args = parse_args( - ( - "--checkpoint", - "policy.pt", - "--config", - "train.yaml", - *arguments, - ) - ) - - with pytest.raises(ValueError, match=option): - _validate_native_options(args) diff --git a/tests/learning/rl/motion_policy_evaluation/test_manifest.py b/tests/learning/rl/motion_policy_evaluation/test_manifest.py deleted file mode 100644 index 812a6ac2c..000000000 --- a/tests/learning/rl/motion_policy_evaluation/test_manifest.py +++ /dev/null @@ -1,49 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -from embodichain.learning.rl.motion_policy_evaluation import ( - RunManifest, - write_run_manifest, -) - - -def test_manifest_snapshots_configs_and_selects_latest_fallback(tmp_path): - run = tmp_path / "run" - checkpoint = run / "checkpoints" / "policy.pt" - checkpoint.parent.mkdir(parents=True) - checkpoint.write_bytes(b"checkpoint") - train = tmp_path / "train.yaml" - train.write_text("trainer: {}\n", encoding="utf-8") - gym = tmp_path / "gym.yaml" - gym.write_text("id: Example\n", encoding="utf-8") - - write_run_manifest( - run, - train_config=train, - gym_config=gym, - latest_checkpoint=checkpoint, - motion_profile="example-motion", - ) - manifest = RunManifest.load(run) - selected, path = manifest.select_checkpoint("best") - - assert manifest.motion_profile == "example-motion" - assert manifest.configs["train"] == run / "configs" / "train.yaml" - assert manifest.configs["gym"] == run / "configs" / "gym.yaml" - assert selected == "latest" - assert path == checkpoint diff --git a/tests/learning/rl/motion_policy_evaluation/test_native_task.py b/tests/learning/rl/motion_policy_evaluation/test_native_task.py deleted file mode 100644 index 82596e504..000000000 --- a/tests/learning/rl/motion_policy_evaluation/test_native_task.py +++ /dev/null @@ -1,380 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -from types import SimpleNamespace - -import pytest -import torch -from tensordict import TensorDict - -from embodichain.learning.rl.evaluation import infer_policy_action -from embodichain.learning.rl.motion_policy_evaluation.native_task import ( - EmbodiChainTaskPolicyAdapter, - _policy_context_from_env, - evaluate_native_task, -) -from embodichain.learning.rl.runtime import PolicyRuntime -from dexsim.kit.motion_policy import EvaluationFrame, PolicyContext - - -class Policy(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.weight = torch.nn.Parameter(torch.tensor([[2.0], [-1.0]])) - - def get_action( - self, - tensordict: TensorDict, - deterministic: bool = False, - ) -> TensorDict: - assert deterministic - tensordict["action"] = tensordict["obs"] @ self.weight - return tensordict - - -class ActionManager: - def __init__(self) -> None: - self.calls = 0 - - def convert_policy_action_to_env_action(self, action: torch.Tensor): - self.calls += 1 - return action + 0.5 - - -class Environment: - num_envs = 1 - physics_dt = 0.005 - step_dt = 0.02 - - def __init__(self) -> None: - self.unwrapped = self - self.cfg = SimpleNamespace(sim_steps_per_control=4) - self.action_manager = ActionManager() - self.actions = [] - self.episode_step = 0 - self.reset_seeds = [] - self.closed = 0 - - def reset(self, seed=None): - self.reset_seeds.append(seed) - self.episode_step = 0 - return self._observation(), {} - - def step(self, action): - self.actions.append(action.clone()) - self.episode_step += 1 - done = self.episode_step == 2 - return ( - self._observation(), - torch.tensor([1.25]), - torch.tensor([done]), - torch.tensor([False]), - { - "success": torch.tensor([done]), - "metrics": {"task_progress": torch.tensor([float(self.episode_step)])}, - }, - ) - - def close(self): - self.closed += 1 - - def _observation(self): - return { - "policy": torch.tensor( - [[float(self.episode_step), 1.0]], - dtype=torch.float32, - ) - } - - -class SimulatorEnvironment(Environment): - def __init__(self) -> None: - super().__init__() - self.sim = SimpleNamespace(get_world=lambda: None) - self.exit_process_values = [] - - def close(self, *, exit_process=None): - self.closed += 1 - self.exit_process_values.append(exit_process) - - -class ViewerWindow: - def __init__(self) -> None: - self.titles = [] - - def set_window_title(self, title): - self.titles.append(title) - - def native(self): - return self - - def key_state(self, _key): - return False - - -class ViewerWorld: - def __init__(self) -> None: - self.window = ViewerWindow() - self.open = True - self.update_dts = [] - - def is_window_initialized(self): - return self.open - - def get_windows(self): - return self.window - - def update(self, dt): - self.update_dts.append(dt) - self.open = False - - -class ViewerSimulatorEnvironment(SimulatorEnvironment): - def __init__(self) -> None: - super().__init__() - self.world = ViewerWorld() - self.sim = SimpleNamespace(get_world=lambda: self.world) - - -def test_native_adapter_uses_shared_deterministic_inference_chain(): - policy = Policy() - observation = {"policy": torch.tensor([[0.25, 0.75]])} - expected = infer_policy_action( - policy, - observation, - device="cpu", - num_envs=1, - ) - adapter = EmbodiChainTaskPolicyAdapter(policy, torch.device("cpu"), 1) - adapter.setup(PolicyContext(None, 0.005, 4, 0.02)) - - output = adapter.infer( - EvaluationFrame( - control_step=0, - policy_time=0.0, - simulation_step=0, - simulation_time=0.0, - observation=observation, - ) - ) - - assert torch.equal(output.action, expected) - adapter.close() - - -def test_lightweight_task_timing_uses_one_environment_step(): - context = _policy_context_from_env(SimpleNamespace(dt=0.125)) - - assert context.physics_dt == pytest.approx(0.125) - assert context.sim_steps_per_control == 1 - assert context.policy_dt == pytest.approx(0.125) - - -def test_native_task_runs_original_action_conversion_once_per_step(): - env = Environment() - runtime = PolicyRuntime( - env=env, - policy=Policy(), - device=torch.device("cpu"), - env_id="ExampleTask", - ) - - result = evaluate_native_task( - runtime, - seed=17, - viewer=False, - episodes=2, - control_steps=None, - duration=None, - ) - - assert result.reason == "episode target reached" - assert result.control_steps == 4 - assert result.simulation_steps == 16 - assert result.effective_duration == pytest.approx(0.08) - assert len(result.episodes) == 2 - assert result.metrics == pytest.approx( - { - "reward": 1.25, - "task_progress": 2.0, - "eval/avg_reward": 2.5, - "eval/avg_length": 2.0, - "eval/success_rate": 1.0, - } - ) - assert env.action_manager.calls == 4 - assert len(env.actions) == 4 - assert env.reset_seeds == [17, None] - assert env.closed == 1 - - -def test_native_task_control_limit_can_stop_before_episode_end(): - env = Environment() - runtime = PolicyRuntime( - env=env, - policy=Policy(), - device=torch.device("cpu"), - env_id="ExampleTask", - ) - - result = evaluate_native_task( - runtime, - seed=1, - viewer=False, - episodes=None, - control_steps=1, - duration=None, - ) - - assert result.reason == "control steps reached" - assert result.control_steps == 1 - assert result.episodes == () - - -def test_native_task_closes_simulator_without_exiting_process(): - env = SimulatorEnvironment() - runtime = PolicyRuntime( - env=env, - policy=Policy(), - device=torch.device("cpu"), - env_id="ExampleSimulatorTask", - ) - - result = evaluate_native_task( - runtime, - seed=1, - viewer=False, - episodes=None, - control_steps=1, - duration=None, - ) - - assert result.reason == "control steps reached" - assert env.closed == 1 - assert env.exit_process_values == [False] - - -def test_native_task_closes_resources_when_evaluator_creation_fails(monkeypatch): - env = Environment() - policy = Policy() - runtime = PolicyRuntime( - env=env, - policy=policy, - device=torch.device("cpu"), - env_id="ExampleTask", - ) - - def fail_evaluator_creation(**_kwargs): - raise RuntimeError("evaluator setup failed") - - monkeypatch.setattr( - "embodichain.learning.rl.motion_policy_evaluation.native_task.create_motion_policy_evaluator", - fail_evaluator_creation, - ) - - with pytest.raises(RuntimeError, match="evaluator setup failed"): - evaluate_native_task( - runtime, - seed=1, - viewer=False, - episodes=1, - control_steps=None, - duration=None, - ) - - assert env.closed == 1 - assert policy.training is True - - -def test_native_task_closes_environment_when_adapter_creation_fails(monkeypatch): - env = Environment() - runtime = PolicyRuntime( - env=env, - policy=Policy(), - device=torch.device("cpu"), - env_id="ExampleTask", - ) - - def fail_adapter_creation(*_args, **_kwargs): - raise RuntimeError("adapter setup failed") - - monkeypatch.setattr( - "embodichain.learning.rl.motion_policy_evaluation.native_task.EmbodiChainTaskPolicyAdapter", - fail_adapter_creation, - ) - - with pytest.raises(RuntimeError, match="adapter setup failed"): - evaluate_native_task( - runtime, - seed=1, - viewer=False, - episodes=1, - control_steps=None, - duration=None, - ) - - assert env.closed == 1 - - -def test_native_task_pause_waits_for_viewer_close_after_termination(): - env = ViewerSimulatorEnvironment() - runtime = PolicyRuntime( - env=env, - policy=Policy(), - device=torch.device("cpu"), - env_id="ExampleViewerTask", - ) - - result = evaluate_native_task( - runtime, - seed=1, - viewer=True, - episodes=None, - control_steps=None, - duration=None, - termination_behavior="pause", - ) - - assert result.reason == "viewer closed" - assert result.control_steps == 2 - assert env.world.update_dts == [0.0] - - -def test_native_task_viewer_requires_a_simulator_environment(): - env = Environment() - runtime = PolicyRuntime( - env=env, - policy=Policy(), - device=torch.device("cpu"), - env_id="ExampleTask", - ) - - with pytest.raises( - ValueError, - match="requires an EmbodiChain simulator task", - ): - evaluate_native_task( - runtime, - seed=1, - viewer=True, - episodes=None, - control_steps=None, - duration=None, - ) - - assert env.closed == 1 diff --git a/tests/learning/rl/motion_policy_evaluation/test_profile.py b/tests/learning/rl/motion_policy_evaluation/test_profile.py deleted file mode 100644 index ba2eac2ce..000000000 --- a/tests/learning/rl/motion_policy_evaluation/test_profile.py +++ /dev/null @@ -1,54 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -import torch - -from embodichain.learning.rl.motion_policy_evaluation import ( - MotionProfile, - MotionProfileRequest, - build_motion_profile, - register_motion_profile, -) - - -def test_profile_provider_receives_checkpoint_and_training_configs(tmp_path): - checkpoint = tmp_path / "policy.pt" - checkpoint.write_bytes(b"checkpoint") - train = tmp_path / "train.yaml" - train.write_text("trainer: {}\n", encoding="utf-8") - requests = [] - - def provider(request): - requests.append(request) - return MotionProfile( - profile_id="test-profile-provider", - checkpoint=request.checkpoint, - policy_spec={"schema_version": 1}, - ) - - register_motion_profile("test-profile-provider", provider) - request = MotionProfileRequest( - checkpoint=checkpoint, - device=torch.device("cpu"), - configs={"train": train}, - ) - - profile = build_motion_profile("test-profile-provider", request) - - assert profile.checkpoint == checkpoint - assert requests[0].configs["train"] == train diff --git a/tests/learning/rl/motion_policy_evaluation/test_train_integration.py b/tests/learning/rl/motion_policy_evaluation/test_train_integration.py deleted file mode 100644 index 167f7a031..000000000 --- a/tests/learning/rl/motion_policy_evaluation/test_train_integration.py +++ /dev/null @@ -1,67 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -from embodichain.learning.rl.motion_policy_evaluation import RunManifest -from embodichain.learning.rl.train import _event_params, _write_motion_run_manifest - - -def test_training_summary_writes_minimal_motion_manifest(tmp_path): - run = tmp_path / "run" - checkpoint = run / "checkpoints" / "policy.pt" - checkpoint.parent.mkdir(parents=True) - checkpoint.write_bytes(b"checkpoint") - config = tmp_path / "train.yaml" - config.write_text("trainer: {}\n", encoding="utf-8") - - _write_motion_run_manifest( - run, - config, - {"motion_profile": "example-motion"}, - { - "latest_checkpoint_path": str(checkpoint), - "best_checkpoint_path": None, - }, - ) - - manifest = RunManifest.load(run) - assert manifest.motion_profile == "example-motion" - assert manifest.checkpoints["latest"] == checkpoint - - -def test_camera_recorder_defaults_to_the_run_video_directory(tmp_path): - params = _event_params( - {"func": "record_camera_data_async", "params": {"name": "main"}}, - run_base=tmp_path / "run", - phase="eval", - ) - - assert params["save_path"] == str(tmp_path / "run" / "videos" / "eval") - - -def test_camera_recorder_keeps_an_explicit_output_directory(tmp_path): - custom = tmp_path / "custom-videos" - params = _event_params( - { - "func": "record_camera_data", - "params": {"save_path": str(custom)}, - }, - run_base=tmp_path / "run", - phase="train", - ) - - assert params["save_path"] == str(custom) diff --git a/tests/learning/rl/motion_policy_evaluation/test_anymal_c_example.py b/tests/learning/rl/policy_evaluation/test_anymal_c_example.py similarity index 63% rename from tests/learning/rl/motion_policy_evaluation/test_anymal_c_example.py rename to tests/learning/rl/policy_evaluation/test_anymal_c_example.py index 155ec9aed..c07d6ef9f 100644 --- a/tests/learning/rl/motion_policy_evaluation/test_anymal_c_example.py +++ b/tests/learning/rl/policy_evaluation/test_anymal_c_example.py @@ -16,12 +16,14 @@ from __future__ import annotations -import importlib.util -import subprocess from pathlib import Path import numpy as np +import pytest import torch + +pytest.importorskip("dexsim.kit.motion_policy.evaluator") + from dexsim.kit.motion_policy import ( AdapterRequest, EvaluationFrame, @@ -31,7 +33,10 @@ parse_policy_spec, ) -from embodichain.learning.rl.motion_policy_evaluation import MotionProfileRequest +from embodichain.learning.rl.policy_evaluation import ( + MotionProfileRequest, + build_motion_profile, +) _JOINT_NAMES = ( "LF_HAA", @@ -64,94 +69,15 @@ def forward(self, observation: torch.Tensor) -> torch.Tensor: return torch.cat((observation[:, 9:12], padding), dim=1) -def _run_git(directory: Path, *args: str) -> str: - result = subprocess.run( - ["git", "-C", str(directory), *args], - check=True, - text=True, - capture_output=True, - ) - return result.stdout.strip() - - -def test_resource_preparation_resumes_interrupted_checkout(tmp_path, capsys): - script = ( - Path(__file__).resolve().parents[4] - / "examples/learning/motion_policy_evaluation/prepare_resources.py" - ) - spec = importlib.util.spec_from_file_location("anymal_c_prepare_resources", script) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - - upstream = tmp_path / "upstream" - upstream.mkdir() - _run_git(upstream, "init", "--quiet") - resource = upstream / "resources/robot.urdf" - resource.parent.mkdir() - resource.write_text("", encoding="utf-8") - _run_git(upstream, "add", ".") - _run_git( - upstream, - "-c", - "user.name=Test User", - "-c", - "user.email=test@example.com", - "commit", - "--quiet", - "-m", - "add resource", - ) - revision = _run_git(upstream, "rev-parse", "HEAD") - - checkout = tmp_path / "cache/upstream" - checkout.mkdir(parents=True) - _run_git(checkout, "init", "--quiet") - _run_git(checkout, "remote", "add", "origin", str(upstream)) - _run_git(checkout, "sparse-checkout", "init", "--no-cone") - _run_git(checkout, "sparse-checkout", "set", "/resources/**") - _run_git(checkout, "fetch", "--quiet", "origin", revision) - - module._prepare_checkout( - checkout, - str(upstream), - revision, - ("/resources/**",), - ) - - assert _run_git(checkout, "rev-parse", "HEAD") == revision - assert (checkout / "resources/robot.urdf").is_file() - assert "Resuming interrupted checkout" in capsys.readouterr().out - - module._prepare_checkout( - checkout, - str(upstream), - revision, - ("/resources/**",), - ) - assert "Using cached revision" in capsys.readouterr().out - - (checkout / "resources/robot.urdf").unlink() - module._prepare_checkout( - checkout, - str(upstream), - revision, - ("/resources/**",), - ) - assert (checkout / "resources/robot.urdf").is_file() - assert "Repairing interrupted checkout" in capsys.readouterr().out - - def test_anymal_c_profile_builds_and_runs_torchscript(tmp_path, monkeypatch): example_root = ( - Path(__file__).resolve().parents[4] - / "examples/learning/motion_policy_evaluation" + Path(__file__).resolve().parents[4] / "examples/learning/policy_evaluation" ) monkeypatch.syspath_prepend(str(example_root)) from anymal_c.profile import ( AnymalCVelocityAdapter, - build_profile, ) + from anymal_c import register checkpoint = tmp_path / "mjw_anymal.pt" traced = torch.jit.trace(_CommandPolicy().eval(), torch.zeros((1, 48))) @@ -161,13 +87,13 @@ def test_anymal_c_profile_builds_and_runs_torchscript(tmp_path, monkeypatch): robot_asset.parent.mkdir(parents=True) robot_asset.write_text("\n", encoding="utf-8") - profile = build_profile( - MotionProfileRequest( - checkpoint=checkpoint, - device=torch.device("cpu"), - resource_root=tmp_path, - ) + request = MotionProfileRequest( + checkpoint=checkpoint, + device=torch.device("cpu"), + resource_root=tmp_path, ) + register() + profile = build_motion_profile("newton-anymal-c-velocity", request) spec = parse_policy_spec(profile.policy_spec) assert spec.environment.entrypoint is None config = profile.policy_spec["policy"]["adapter"]["config"] @@ -235,8 +161,7 @@ def test_anymal_c_profile_builds_and_runs_torchscript(tmp_path, monkeypatch): def test_example_script_supplies_default_resource_paths(tmp_path, monkeypatch): example_root = ( - Path(__file__).resolve().parents[4] - / "examples/learning/motion_policy_evaluation" + Path(__file__).resolve().parents[4] / "examples/learning/policy_evaluation" ) monkeypatch.syspath_prepend(str(example_root)) monkeypatch.setenv("ANYMAL_C_EXAMPLE_CACHE", str(tmp_path)) diff --git a/tests/learning/rl/policy_evaluation/test_bridge.py b/tests/learning/rl/policy_evaluation/test_bridge.py new file mode 100644 index 000000000..eb3f120f5 --- /dev/null +++ b/tests/learning/rl/policy_evaluation/test_bridge.py @@ -0,0 +1,84 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("dexsim.kit.motion_policy.evaluator") + +from embodichain.learning.rl.policy_evaluation.bridge import ( + evaluate_motion_profile, +) +from embodichain.learning.rl.policy_evaluation.profile import MotionProfile + + +def test_bridge_forwards_physics_backend_and_exact_control_steps( + monkeypatch, +): + profile = MotionProfile( + profile_id="example", + policy_spec={"schema_version": 1}, + ) + options = [] + monkeypatch.setattr( + "embodichain.learning.rl.policy_evaluation.bridge.parse_policy_spec", + lambda value: "parsed", + ) + monkeypatch.setattr( + "embodichain.learning.rl.policy_evaluation.bridge.resolve_policy_spec", + lambda spec, resolver: "resolved", + ) + monkeypatch.setattr( + "embodichain.learning.rl.policy_evaluation.bridge.policy_spec_to_dict", + lambda value: {"policy_id": "example"}, + ) + monkeypatch.setattr( + "embodichain.learning.rl.policy_evaluation.bridge.scene_config_to_dict", + lambda value: {"style": "standard"}, + ) + + def run_policy(resolved, run_options): + options.append(run_options) + return SimpleNamespace( + reason="control steps reached", + simulation_time=0.2, + simulation_steps=40, + control_steps=10, + physics_backend="default", + requested_duration=None, + effective_duration=0.2, + metrics={"tracking/error": 0.25}, + ) + + monkeypatch.setattr( + "embodichain.learning.rl.policy_evaluation.bridge.run_motion_policy", + run_policy, + ) + + result = evaluate_motion_profile( + profile, + control_steps=10, + physics_backend="default", + ) + + assert options[0].control_steps == 10 + assert options[0].physics_backend == "default" + assert result.episodes[0]["control_steps"] == 10 + assert result.episodes[0]["effective_duration"] == 0.2 + assert result.summary["metrics"]["tracking/error"] == 0.25 diff --git a/tests/learning/rl/policy_evaluation/test_cli.py b/tests/learning/rl/policy_evaluation/test_cli.py new file mode 100644 index 000000000..2d8ad1611 --- /dev/null +++ b/tests/learning/rl/policy_evaluation/test_cli.py @@ -0,0 +1,123 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import importlib + +import pytest + +from embodichain.learning.rl.policy_evaluation.cli import ( + _resolve_input, + _validate_native_options, + parse_args, +) +from embodichain.learning.rl.policy_evaluation.manifest import write_run_manifest + + +def _run(tmp_path): + run = tmp_path / "run" + checkpoint = run / "checkpoints" / "policy.pt" + checkpoint.parent.mkdir(parents=True) + checkpoint.write_bytes(b"checkpoint") + train = tmp_path / "train.yaml" + train.write_text("trainer: {}\n", encoding="utf-8") + write_run_manifest( + run, + train_config=train, + latest_checkpoint=checkpoint, + ) + return run, checkpoint + + +def test_run_defaults_to_latest_checkpoint(tmp_path): + run, checkpoint = _run(tmp_path) + + resolved = _resolve_input(parse_args((str(run),))) + + assert resolved.checkpoint == checkpoint + assert resolved.requested_checkpoint == "latest" + assert resolved.selected_checkpoint == "latest" + + +@pytest.mark.parametrize( + "arguments, handler", + [ + ((), "_run_native_headless"), + (("--viewer",), "_run_native_viewer"), + (("--profile", "example"), "_run_profile"), + ], +) +def test_cli_routes_one_command_to_the_selected_evaluation( + tmp_path, + monkeypatch, + arguments, + handler, +): + module = importlib.import_module("embodichain.learning.rl.policy_evaluation.cli") + run, _checkpoint = _run(tmp_path) + expected = tmp_path / "evaluation.json" + calls = [] + + monkeypatch.setattr(module, "discover_task_packages", lambda: None) + monkeypatch.setattr(module, "execute_init_hooks", lambda: None) + for name in ("_run_native_headless", "_run_native_viewer", "_run_profile"): + monkeypatch.setattr( + module, + name, + lambda args, resolved, name=name: calls.append(name) or expected, + ) + + report = module.run(parse_args((str(run), *arguments))) + + assert report == expected + assert calls == [handler] + + +def test_explicit_checkpoint_requires_training_config(tmp_path): + checkpoint = tmp_path / "policy.pt" + checkpoint.write_bytes(b"checkpoint") + + with pytest.raises(ValueError, match="--config is required"): + _resolve_input(parse_args(("--checkpoint", str(checkpoint)))) + + +def test_native_options_keep_profile_and_viewer_inputs_explicit(): + profile_args = parse_args( + ( + "--checkpoint", + "policy.pt", + "--config", + "train.yaml", + "--command", + "0.5", + ) + ) + viewer_args = parse_args( + ( + "--checkpoint", + "policy.pt", + "--config", + "train.yaml", + "--control-steps", + "10", + ) + ) + + with pytest.raises(ValueError, match="--command requires --profile"): + _validate_native_options(profile_args) + with pytest.raises(ValueError, match="--control-steps.*require --viewer"): + _validate_native_options(viewer_args) diff --git a/tests/learning/rl/policy_evaluation/test_viewer.py b/tests/learning/rl/policy_evaluation/test_viewer.py new file mode 100644 index 000000000..bde501fe0 --- /dev/null +++ b/tests/learning/rl/policy_evaluation/test_viewer.py @@ -0,0 +1,221 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch +from tensordict import TensorDict + +pytest.importorskip("dexsim.kit.motion_policy.evaluator") + +from dexsim.kit.motion_policy import EvaluationFrame, PolicyContext + +from embodichain.learning.rl.evaluation import infer_policy_action +from embodichain.learning.rl.policy_evaluation.viewer import ( + EmbodiChainTaskEnvironment, + EmbodiChainTaskPolicyAdapter, + evaluate_native_viewer, +) +from embodichain.learning.rl.runtime import PolicyRuntime + + +class Policy(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor([[2.0], [-1.0]])) + + def get_action( + self, + tensordict: TensorDict, + deterministic: bool = False, + ) -> TensorDict: + assert deterministic + tensordict["action"] = tensordict["obs"] @ self.weight + return tensordict + + +class Window: + def __init__(self) -> None: + self.titles = [] + self.keys = set() + + def set_window_title(self, title): + self.titles.append(title) + + def native(self): + return self + + def key_state(self, key): + return key in self.keys + + +class World: + def __init__(self) -> None: + self.window = Window() + self.open = True + + def is_window_initialized(self): + return self.open + + def get_windows(self): + return self.window + + +class ActionManager: + def __init__(self) -> None: + self.calls = 0 + + def convert_policy_action_to_env_action(self, action: torch.Tensor): + self.calls += 1 + return action + 0.5 + + +class Environment: + num_envs = 1 + physics_dt = 0.005 + step_dt = 0.02 + + def __init__(self) -> None: + self.unwrapped = self + self.cfg = SimpleNamespace(sim_steps_per_control=4) + self.action_manager = ActionManager() + self.world = World() + self.sim = SimpleNamespace(get_world=lambda: self.world) + self.actions = [] + self.episode_step = 0 + self.reset_seeds = [] + self.exit_process_values = [] + + def reset(self, seed=None): + self.reset_seeds.append(seed) + self.episode_step = 0 + return self._observation(), {} + + def step(self, action): + self.actions.append(action.clone()) + self.episode_step += 1 + done = self.episode_step == 2 + return ( + self._observation(), + torch.tensor([1.25]), + torch.tensor([done]), + torch.tensor([False]), + { + "success": torch.tensor([done]), + "metrics": {"task_progress": torch.tensor([self.episode_step])}, + }, + ) + + def close(self, *, exit_process=None): + self.exit_process_values.append(exit_process) + + def _observation(self): + return { + "policy": torch.tensor( + [[float(self.episode_step), 1.0]], + dtype=torch.float32, + ) + } + + +def _runtime(env: Environment, policy: Policy | None = None) -> PolicyRuntime: + return PolicyRuntime( + env=env, + policy=policy or Policy(), + device=torch.device("cpu"), + env_id="ExampleTask", + ) + + +def test_viewer_adapter_uses_the_shared_deterministic_inference_chain(): + policy = Policy() + observation = {"policy": torch.tensor([[0.25, 0.75]])} + expected = infer_policy_action(policy, observation, device="cpu", num_envs=1) + adapter = EmbodiChainTaskPolicyAdapter(policy, torch.device("cpu")) + adapter.setup(PolicyContext(None, 0.005, 4, 0.02)) + + output = adapter.infer(EvaluationFrame(0, 0.0, 0.0, 0, observation=observation)) + + assert torch.equal(output.action, expected) + adapter.close() + + +def test_viewer_reuses_task_actions_resets_and_metrics(): + env = Environment() + + result = evaluate_native_viewer( + _runtime(env), + seed=17, + episodes=2, + control_steps=None, + duration=None, + ) + + assert result.control_steps == 4 + assert result.simulation_steps == 16 + assert len(result.episodes) == 2 + assert result.metrics == pytest.approx( + { + "reward": 1.25, + "task_progress": 2.0, + "eval/avg_reward": 2.5, + "eval/avg_length": 2.0, + "eval/success_rate": 1.0, + } + ) + assert env.action_manager.calls == 4 + assert env.reset_seeds == [17, None] + assert env.exit_process_values == [False] + + +def test_backspace_requests_one_reset_per_key_press(): + from dexsim.types import InputKey + + env = Environment() + task = EmbodiChainTaskEnvironment(env, seed=1) + env.world.window.keys.add(InputKey.SCANCODE_BACKSPACE) + + assert task.poll() == "manual reset" + assert task.poll() is None + env.world.window.keys.clear() + assert task.poll() is None + env.world.window.keys.add(InputKey.SCANCODE_BACKSPACE) + assert task.poll() == "manual reset" + task.close() + + +def test_viewer_closes_resources_when_evaluator_creation_fails(monkeypatch): + env = Environment() + policy = Policy() + monkeypatch.setattr( + "embodichain.learning.rl.policy_evaluation.viewer.create_motion_policy_evaluator", + lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("setup failed")), + ) + + with pytest.raises(RuntimeError, match="setup failed"): + evaluate_native_viewer( + _runtime(env, policy), + seed=1, + episodes=1, + control_steps=None, + duration=None, + ) + + assert env.exit_process_values == [False] + assert policy.training is True diff --git a/tests/learning/test_point_mass.py b/tests/learning/test_point_mass.py index fd3fdab48..669125b93 100644 --- a/tests/learning/test_point_mass.py +++ b/tests/learning/test_point_mass.py @@ -31,6 +31,7 @@ build_learning_env, ) from embodichain.learning.rl.models import ActorCritic +from embodichain.learning.rl.policy_evaluation.manifest import RunManifest from embodichain.learning.rl.train import train_from_config from embodichain.learning.rl.utils import OptimizerCfg from embodichain.learning.rl.utils.trainer import Trainer @@ -183,6 +184,13 @@ def test_unified_train_entry_runs_apg_and_ppo( assert summary["global_step"] == 8 assert summary["latest_checkpoint_path"] is not None + checkpoint = Path(summary["latest_checkpoint_path"]).resolve() + run = checkpoint.parents[1] + manifest = RunManifest.load(run) + assert ( + manifest.configs["train"] == run / "configs" / f"train.{config_path.suffix[1:]}" + ) + assert manifest.select_checkpoint("latest")[1] == checkpoint def test_sync_collector_accepts_tensor_point_mass_observations() -> None: diff --git a/tests/learning/test_runtime.py b/tests/learning/test_runtime.py index 3dd20cf9b..ca35bab74 100644 --- a/tests/learning/test_runtime.py +++ b/tests/learning/test_runtime.py @@ -22,11 +22,9 @@ import torch from embodichain.learning.rl.runtime import ( - GymEnvironmentRuntime, - build_gym_environment, + _GymEnvironmentRuntime, + _build_gym_environment, build_gym_policy_runtime, - build_learning_policy_runtime, - resolve_config_reference, ) @@ -42,17 +40,6 @@ def _policy_config() -> dict: } -class LearningEnvironment: - single_observation_space = gym.spaces.Box(-1.0, 1.0, shape=(3,)) - single_action_space = gym.spaces.Box(-1.0, 1.0, shape=(2,)) - - def __init__(self) -> None: - self.closed = 0 - - def close(self) -> None: - self.closed += 1 - - class GymEnvironment: flattened_observation_space = gym.spaces.Box(-1.0, 1.0, shape=(3,)) observation_space = gym.spaces.Dict( @@ -74,34 +61,6 @@ def close(self) -> None: self.closed += 1 -def test_config_reference_uses_the_training_config_directory(tmp_path): - config_dir = tmp_path / "configs" - config_dir.mkdir() - gym_config = config_dir / "gym.yaml" - gym_config.write_text("id: Example\n", encoding="utf-8") - - assert resolve_config_reference("gym.yaml", base_dir=config_dir) == gym_config - - -def test_learning_runtime_builds_the_configured_policy(monkeypatch): - env = LearningEnvironment() - monkeypatch.setattr( - "embodichain.learning.rl.runtime.build_learning_environment", - lambda config, device, num_envs: ("PointMass", env), - ) - - runtime = build_learning_policy_runtime( - {"trainer": {"learning_env": "PointMass"}, "policy": _policy_config()}, - device=torch.device("cpu"), - num_envs=4, - ) - - assert runtime.env is env - assert runtime.env_id == "PointMass" - assert runtime.policy.actor[0].in_features == 3 - assert runtime.policy.actor[-1].out_features == 2 - - def test_gym_environment_applies_runtime_overrides(tmp_path, monkeypatch): gym_config = tmp_path / "gym.yaml" gym_config.write_text("id: Example\n", encoding="utf-8") @@ -120,7 +79,7 @@ def test_gym_environment_applies_runtime_overrides(tmp_path, monkeypatch): lambda env_id, base_env_cfg: built, ) - runtime = build_gym_environment( + runtime = _build_gym_environment( {"trainer": {"gym_config": "gym.yaml"}}, simulation_device=torch.device("cpu"), num_envs=1, @@ -140,7 +99,7 @@ def test_gym_environment_applies_runtime_overrides(tmp_path, monkeypatch): def test_gym_runtime_uses_the_same_task_spaces_for_policy_build(monkeypatch): env = GymEnvironment() - task = GymEnvironmentRuntime( + task = _GymEnvironmentRuntime( env=env, env_id="Example", env_cfg=SimpleNamespace(), @@ -148,7 +107,7 @@ def test_gym_runtime_uses_the_same_task_spaces_for_policy_build(monkeypatch): gym_config_path=SimpleNamespace(resolve=lambda: None), ) monkeypatch.setattr( - "embodichain.learning.rl.runtime.build_gym_environment", + "embodichain.learning.rl.runtime._build_gym_environment", lambda *args, **kwargs: task, ) diff --git a/tests/learning/test_train_profile.py b/tests/learning/test_train_profile.py index 8b30b74f2..7934fb7e2 100644 --- a/tests/learning/test_train_profile.py +++ b/tests/learning/test_train_profile.py @@ -21,6 +21,7 @@ import pytest from embodichain.learning.rl.train import ( + _event_params, _resolve_profile_output, parse_args, train_from_config, @@ -65,3 +66,13 @@ def test_learning_env_rejects_profile(tmp_path): with pytest.raises(ValueError, match="--profile_output requires --profile"): train_from_config(str(config_path), profile_output="prof.json") + + +def test_camera_recording_defaults_to_the_run_directory(tmp_path): + params = _event_params( + {"func": "record_camera_data_async", "params": {"name": "main"}}, + run_base=tmp_path / "run", + phase="eval", + ) + + assert params["save_path"] == str(tmp_path / "run" / "videos" / "eval") diff --git a/tests/test_main.py b/tests/test_main.py index 906e00173..b35b23520 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -28,7 +28,7 @@ "benchmark", "data", "decompose-urdf", - "eval-motion-policy", + "eval-policy", "preview-asset", "preview_lerobot_data", "run-env", From 796dded301b11ce667349416035d5e2ed9473612 Mon Sep 17 00:00:00 2001 From: acrlw <13927622+acrlw@users.noreply.github.com> Date: Wed, 19 Aug 2026 11:06:08 +0800 Subject: [PATCH 7/7] refactor(sim): standardize renderer names Use hybrid, fast-rt, and offline-rt consistently across simulation configuration, CLIs, policy evaluation, documentation, and focused tests. --- docs/source/features/toolkits/grasp_generator.rst | 2 +- docs/source/guides/cli.md | 2 +- docs/source/guides/policy_evaluation.md | 2 +- docs/source/guides/preview_asset.md | 2 +- docs/source/overview/sim/sim_manager.md | 2 +- embodichain/lab/gym/utils/gym_utils.py | 4 ++-- embodichain/lab/scripts/analyze_workspace.py | 2 +- embodichain/lab/scripts/preview_asset.py | 2 +- embodichain/lab/sim/cfg.py | 14 +++++++------- embodichain/lab/sim/sim_manager.py | 4 ++-- embodichain/lab/sim/utility/render_utils.py | 3 ++- embodichain/learning/rl/policy_evaluation/cli.py | 2 +- scripts/benchmark/atomic_action/common.py | 2 +- scripts/benchmark/atomic_action/press_benchmark.py | 2 +- scripts/benchmark/atomic_action/run_benchmark.py | 2 +- tests/gym/utils/test_gym_utils.py | 10 +++++----- tests/learning/rl/policy_evaluation/test_cli.py | 7 +++++++ tests/sim/test_cfg.py | 3 ++- 18 files changed, 38 insertions(+), 29 deletions(-) diff --git a/docs/source/features/toolkits/grasp_generator.rst b/docs/source/features/toolkits/grasp_generator.rst index 64ae1a27d..c80c28106 100644 --- a/docs/source/features/toolkits/grasp_generator.rst +++ b/docs/source/features/toolkits/grasp_generator.rst @@ -224,7 +224,7 @@ You can customize the run with additional arguments: .. code-block:: bash - python scripts/tutorials/grasp/grasp_generator.py --num_envs --device --renderer --headless + python scripts/tutorials/grasp/grasp_generator.py --num_envs --device --renderer --headless After confirming the grasp region in the browser, the script will compute a grasp pose, print the elapsed time, and then wait for you to press **Enter** before executing the full grasp trajectory in the simulation. Press **Enter** again to exit once the motion is complete. diff --git a/docs/source/guides/cli.md b/docs/source/guides/cli.md index 427b8c17e..4e9250b9c 100644 --- a/docs/source/guides/cli.md +++ b/docs/source/guides/cli.md @@ -145,7 +145,7 @@ embodichain run-env --gym_config config.yaml \ | ``--num_envs`` | ``1`` | Number of parallel environments | | ``--device`` | ``cpu`` | Device (``cpu`` or ``cuda``) | | ``--headless`` | ``False`` | Run in headless mode | -| ``--renderer`` | ``auto`` | Renderer backend: ``auto``, ``hybrid``, ``fast-rt`` or ``rt`` | +| ``--renderer`` | ``auto`` | Renderer backend: ``auto``, ``hybrid``, ``fast-rt`` or ``offline-rt`` | | ``--arena_space`` | ``5.0`` | Arena space size | | ``--gpu_id`` | ``0`` | GPU ID to use | | ``--preview`` | ``False`` | Enter interactive preview mode | diff --git a/docs/source/guides/policy_evaluation.md b/docs/source/guides/policy_evaluation.md index f83edd9b0..392da7cc2 100644 --- a/docs/source/guides/policy_evaluation.md +++ b/docs/source/guides/policy_evaluation.md @@ -83,7 +83,7 @@ embodichain eval-policy outputs/_ \ The Viewer uses one environment and keeps running until the window closes. Use `--episodes`, `--control-steps`, or `--duration` to select another stopping -condition. +condition. `--renderer` accepts `hybrid`, `fast-rt`, and `offline-rt`. For a checkpoint created before `run-manifest.json` was introduced, provide its training configuration directly: diff --git a/docs/source/guides/preview_asset.md b/docs/source/guides/preview_asset.md index a83fa49bc..7049382fe 100644 --- a/docs/source/guides/preview_asset.md +++ b/docs/source/guides/preview_asset.md @@ -149,7 +149,7 @@ asset.set_local_pose(pose) | `--use_usd_properties` | disabled | Use physical properties stored in the USD file. | | `--fix_base` / `--no-fix_base` | fixed | Fix or unfix articulation bases. | | `--sim_device` | `cpu` | Simulation device. | -| `--renderer` | `hybrid` | Renderer: `hybrid`, `fast-rt`, or `rt`. | +| `--renderer` | `hybrid` | Renderer: `hybrid`, `fast-rt`, or `offline-rt`. | | `--env_map` | none | Built-in IBL resource name or absolute `.hdr`, `.png`, or `.exr` path. | | `--headless` | disabled | Run without the native window. | | `--preview` | disabled | Enter the interactive terminal after loading. | diff --git a/docs/source/overview/sim/sim_manager.md b/docs/source/overview/sim/sim_manager.md index c98a16615..522412ff4 100644 --- a/docs/source/overview/sim/sim_manager.md +++ b/docs/source/overview/sim/sim_manager.md @@ -68,7 +68,7 @@ The {class}`~cfg.RenderCfg` class controls the rendering backend and quality set | Parameter | Type | Default | Description | | :--- | :--- | :--- | :--- | -| `renderer` | `str` | `"auto"` | Renderer backend to use. Options are `'auto'` (pick a default based on the detected GPU), `'hybrid'` (ray tracing for shadows/reflections + rasterization), `'fast-rt'` (full ray tracing), and `'rt'` (offline ray-traced renderer for maximum visual fidelity). | +| `renderer` | `str` | `"auto"` | Renderer backend to use. Options are `'auto'` (pick a default based on the detected GPU), `'hybrid'` (ray tracing for shadows/reflections + rasterization), `'fast-rt'` (full ray tracing), and `'offline-rt'` (offline ray-traced renderer for maximum visual fidelity). | | `spp` | `int` | `1` | Samples per pixel for ray-traced rendering. Must be at least 1. | | `tone_mapping_enabled` | `bool` | `False` | Whether to map HDR RGB output with the modified Reinhard curve. | | `tone_mapping_exposure` | `float` | `1.0` | Non-negative fixed linear exposure multiplier applied before tone mapping. | diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index cf3c1086b..7e4924c77 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -823,7 +823,7 @@ def add_env_launcher_args_to_parser( --num_envs: Number of environments to run in parallel (default: 1) --device: Device to run the environment on (default: 'cpu') --headless: Whether to perform the simulation in headless mode (default: False) - --renderer: Renderer backend to use for the simulation. Options are 'hybrid', 'fast-rt', and 'rt'. (default: 'hybrid') + --renderer: Renderer backend to use for the simulation. Options are 'auto', 'hybrid', 'fast-rt', and 'offline-rt'. --gpu_id: The GPU ID to use for the simulation (default: 0) --gym_config: Path to gym config file (default: '') --action_config: Path to action config file (default: None) @@ -864,7 +864,7 @@ def add_env_launcher_args_to_parser( parser.add_argument( "--renderer", type=str, - choices=["auto", "hybrid", "fast-rt", "rt"], + choices=["auto", "hybrid", "fast-rt", "offline-rt"], default=None if require_gym_config else "auto", help="Renderer backend to use for the simulation. When loading a gym " "config, the configured render_cfg.renderer is used unless this option " diff --git a/embodichain/lab/scripts/analyze_workspace.py b/embodichain/lab/scripts/analyze_workspace.py index 7f1649c14..e26a17746 100644 --- a/embodichain/lab/scripts/analyze_workspace.py +++ b/embodichain/lab/scripts/analyze_workspace.py @@ -1047,7 +1047,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: sim.add_argument( "--renderer", type=str, - choices=["hybrid", "fast-rt", "rt"], + choices=["hybrid", "fast-rt", "offline-rt"], default="hybrid", help="Renderer backend (default: hybrid).", ) diff --git a/embodichain/lab/scripts/preview_asset.py b/embodichain/lab/scripts/preview_asset.py index 73d17a248..bd3915b01 100644 --- a/embodichain/lab/scripts/preview_asset.py +++ b/embodichain/lab/scripts/preview_asset.py @@ -431,7 +431,7 @@ def _create_parser() -> argparse.ArgumentParser: parser.add_argument( "--renderer", type=str, - choices=["hybrid", "fast-rt", "rt"], + choices=["hybrid", "fast-rt", "offline-rt"], default="hybrid", help="Renderer backend (default: hybrid).", ) diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index c36fabbcd..4fef9a2b0 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -56,13 +56,13 @@ # :func:`embodichain.lab.sim.utility.render_utils.select_default_renderer`). Assigning a # concrete renderer here (e.g. in test fixtures) forces that renderer and takes # precedence over auto-selection. -DEFAULT_RENDERER: Literal["auto", "hybrid", "fast-rt", "rt"] = "auto" +DEFAULT_RENDERER: Literal["auto", "hybrid", "fast-rt", "offline-rt"] = "auto" @configclass class RenderCfg: - renderer: Literal["auto", "hybrid", "fast-rt", "rt"] = "auto" - """Renderer backend to use for the simulation. Options are 'auto', 'hybrid', 'fast-rt', and 'rt'. + renderer: Literal["auto", "hybrid", "fast-rt", "offline-rt"] = "auto" + """Renderer backend to use for the simulation. Options are 'auto', 'hybrid', 'fast-rt', and 'offline-rt'. Note: - 'auto' selects a default renderer based on the detected GPU: RTX-series cards use @@ -71,11 +71,11 @@ class RenderCfg: - 'hybrid' uses ray tracing for shadows and reflections while keeping rasterization for primary rendering, providing a balance between performance and visual quality. - 'fast-rt' is a fully ray-traced renderer for maximum visual fidelity, but may have higher computational cost. - - 'rt' is an offline ray-traced renderer for maximum visual fidelity, suitable for high-quality rendering tasks. + - 'offline-rt' is an offline ray-traced renderer for maximum visual fidelity, suitable for high-quality rendering tasks. """ spp: int = 1 - """Samples per pixel for ray tracing rendering. This parameter is only valid when renderer is 'hybrid', 'fast-rt' or 'rt'.""" + """Samples per pixel for ray tracing rendering. This parameter is only valid when renderer is 'hybrid', 'fast-rt' or 'offline-rt'.""" tone_mapping_enabled: bool = False """Whether to map HDR RGB output with the modified Reinhard curve.""" @@ -98,7 +98,7 @@ def to_dexsim_flags(self) -> Renderer: return Renderer.HYBRID elif self.renderer == "fast-rt": return Renderer.FASTRT - elif self.renderer == "rt": + elif self.renderer == "offline-rt": return Renderer.OFFLINERT elif self.renderer == "auto": # 'auto' is normally resolved by the SimulationManager before this is @@ -110,7 +110,7 @@ def to_dexsim_flags(self) -> Renderer: return Renderer.HYBRID else: logger.log_error( - f"Invalid renderer type '{self.renderer}' specified. Must be one of 'auto', 'hybrid', 'fast-rt', or 'rt'." + f"Invalid renderer type '{self.renderer}' specified. Must be one of 'auto', 'hybrid', 'fast-rt', or 'offline-rt'." ) def apply_to_dexsim_config(self, world_config: dexsim.WorldConfig) -> None: diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index c8d1cf875..c9d7081fd 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -437,7 +437,7 @@ def set_default_renderer(cls, renderer: str = "auto", gpu_id: int = 0) -> str: Args: renderer: The renderer to set. One of ``"auto"``, ``"hybrid"``, - ``"fast-rt"``, or ``"rt"``. When ``"auto"``, the renderer is + ``"fast-rt"``, or ``"offline-rt"``. When ``"auto"``, the renderer is resolved immediately from the detected GPU via :func:`embodichain.lab.sim.utility.render_utils.select_default_renderer`. gpu_id: The CUDA device index to query when ``renderer="auto"``. @@ -448,7 +448,7 @@ def set_default_renderer(cls, renderer: str = "auto", gpu_id: int = 0) -> str: from embodichain.lab.sim import cfg from embodichain.lab.sim.utility.render_utils import select_default_renderer - valid = {"auto", "hybrid", "fast-rt", "rt"} + valid = {"auto", "hybrid", "fast-rt", "offline-rt"} if renderer not in valid: logger.log_error( f"Invalid renderer '{renderer}'. Must be one of {sorted(valid)}." diff --git a/embodichain/lab/sim/utility/render_utils.py b/embodichain/lab/sim/utility/render_utils.py index d82bb2644..8469ad767 100644 --- a/embodichain/lab/sim/utility/render_utils.py +++ b/embodichain/lab/sim/utility/render_utils.py @@ -47,7 +47,8 @@ def select_default_renderer(gpu_id: int = 0) -> str: gpu_id: The CUDA device index to query for selecting the renderer. Returns: - The resolved renderer name, one of ``"hybrid"``, ``"fast-rt"``, or ``"rt"``. + The resolved renderer name, one of ``"hybrid"``, ``"fast-rt"``, or + ``"offline-rt"``. """ from embodichain.lab.sim import cfg diff --git a/embodichain/learning/rl/policy_evaluation/cli.py b/embodichain/learning/rl/policy_evaluation/cli.py index 853bd3daa..8c9b9c2d2 100644 --- a/embodichain/learning/rl/policy_evaluation/cli.py +++ b/embodichain/learning/rl/policy_evaluation/cli.py @@ -99,7 +99,7 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.add_argument("--physics-backend") parser.add_argument( "--renderer", - choices=("raster", "hybrid", "fastrt", "offlinert"), + choices=("hybrid", "fast-rt", "offline-rt"), ) parser.add_argument("--gpu-id", type=int, default=0) parser.add_argument("--scene-config") diff --git a/scripts/benchmark/atomic_action/common.py b/scripts/benchmark/atomic_action/common.py index 7220ab625..7766ff8ae 100644 --- a/scripts/benchmark/atomic_action/common.py +++ b/scripts/benchmark/atomic_action/common.py @@ -240,7 +240,7 @@ def add_common_benchmark_args(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--renderer", type=str, - choices=("auto", "hybrid", "fast-rt", "rt"), + choices=("auto", "hybrid", "fast-rt", "offline-rt"), default="auto", help="Renderer backend used by SimulationManager.", ) diff --git a/scripts/benchmark/atomic_action/press_benchmark.py b/scripts/benchmark/atomic_action/press_benchmark.py index a5687af7f..1d4e1268e 100644 --- a/scripts/benchmark/atomic_action/press_benchmark.py +++ b/scripts/benchmark/atomic_action/press_benchmark.py @@ -280,7 +280,7 @@ def add_benchmark_args(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--renderer", type=str, - choices=("auto", "hybrid", "fast-rt", "rt"), + choices=("auto", "hybrid", "fast-rt", "offline-rt"), default="auto", help="Renderer backend used by SimulationManager.", ) diff --git a/scripts/benchmark/atomic_action/run_benchmark.py b/scripts/benchmark/atomic_action/run_benchmark.py index ab20494e6..1a8cdf356 100644 --- a/scripts/benchmark/atomic_action/run_benchmark.py +++ b/scripts/benchmark/atomic_action/run_benchmark.py @@ -104,7 +104,7 @@ def add_benchmark_args(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--renderer", type=str, - choices=("auto", "hybrid", "fast-rt", "rt"), + choices=("auto", "hybrid", "fast-rt", "offline-rt"), default="auto", help="Renderer backend forwarded to each selected benchmark.", ) diff --git a/tests/gym/utils/test_gym_utils.py b/tests/gym/utils/test_gym_utils.py index 12d46874b..2f034a495 100644 --- a/tests/gym/utils/test_gym_utils.py +++ b/tests/gym/utils/test_gym_utils.py @@ -364,12 +364,12 @@ def test_launcher_preserves_gym_renderer_when_cli_omits_override(): add_env_launcher_args_to_parser(parser, require_gym_config=True) args = parser.parse_args(["--gym_config", "gym_config.yaml"]) - gym_config = {"id": "Dummy-v0", "render_cfg": {"renderer": "rt"}} + gym_config = {"id": "Dummy-v0", "render_cfg": {"renderer": "offline-rt"}} merged_config = merge_args_with_gym_config(args, gym_config) assert args.renderer is None assert "renderer" not in merged_config - assert merged_config["render_cfg"]["renderer"] == "rt" + assert merged_config["render_cfg"]["renderer"] == "offline-rt" def test_env_launcher_includes_viser_arguments(): @@ -544,7 +544,7 @@ def test_yaml_gym_config_parses_to_cfg(self, tmp_path): "speed_tolerance": 0.1, }, "render_cfg": { - "renderer": "rt", + "renderer": "offline-rt", "spp": 4, "tone_mapping_enabled": True, "tone_mapping_exposure": 1.25, @@ -597,7 +597,7 @@ def test_yaml_gym_config_parses_to_cfg(self, tmp_path): assert cfg.sim_cfg.physics_config.enable_ccd is True assert cfg.sim_cfg.physics_config.length_tolerance == 0.02 assert cfg.sim_cfg.physics_config.speed_tolerance == 0.1 - assert cfg.sim_cfg.render_cfg.renderer == "rt" + assert cfg.sim_cfg.render_cfg.renderer == "offline-rt" assert cfg.sim_cfg.render_cfg.spp == 4 assert cfg.sim_cfg.render_cfg.tone_mapping_enabled is True assert cfg.sim_cfg.render_cfg.tone_mapping_exposure == 1.25 @@ -656,7 +656,7 @@ def test_build_env_cfg_applies_modifier_before_parsing(self, tmp_path): "enable_ccd": True, }, "render_cfg": { - "renderer": "rt", + "renderer": "offline-rt", "spp": 8, "tone_mapping_enabled": True, }, diff --git a/tests/learning/rl/policy_evaluation/test_cli.py b/tests/learning/rl/policy_evaluation/test_cli.py index 2d8ad1611..d36059fee 100644 --- a/tests/learning/rl/policy_evaluation/test_cli.py +++ b/tests/learning/rl/policy_evaluation/test_cli.py @@ -95,6 +95,13 @@ def test_explicit_checkpoint_requires_training_config(tmp_path): _resolve_input(parse_args(("--checkpoint", str(checkpoint)))) +@pytest.mark.parametrize("renderer", ("hybrid", "fast-rt", "offline-rt")) +def test_cli_accepts_dexsim_renderer_names(renderer): + args = parse_args(("--renderer", renderer)) + + assert args.renderer == renderer + + def test_native_options_keep_profile_and_viewer_inputs_explicit(): profile_args = parse_args( ( diff --git a/tests/sim/test_cfg.py b/tests/sim/test_cfg.py index c9cfc28fe..aef03c050 100644 --- a/tests/sim/test_cfg.py +++ b/tests/sim/test_cfg.py @@ -94,13 +94,14 @@ def test_render_cfg_applies_tone_mapping_and_fixed_exposure() -> None: expected_exposure = 1.25 world_config = dexsim.WorldConfig() render_cfg = RenderCfg( - renderer="rt", + renderer="offline-rt", tone_mapping_enabled=True, tone_mapping_exposure=expected_exposure, ) render_cfg.apply_to_dexsim_config(world_config) + assert world_config.renderer == Renderer.OFFLINERT assert world_config.postprocess_config.tone_mapping_enabled is True assert ( world_config.postprocess_config.tone_mapping_type