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.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 bf1b5b01e..e2e51b69d 100644 --- a/docs/source/api_reference/embodichain/embodichain.learning.rl.rst +++ b/docs/source/api_reference/embodichain/embodichain.learning.rl.rst @@ -18,6 +18,7 @@ collection logic, policy/model builders, and training entry points. buffer collector models + policy_evaluation train utils 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/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 f009d64ab..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 | @@ -372,6 +372,56 @@ See the Profiling section under Run Env for report format. Outputs are written t --- +## Policy Evaluation + +Evaluate the latest checkpoint from an EmbodiChain training run: + +```bash +embodichain eval-policy outputs/my_policy_ +``` + +Open a simulator task in the Viewer: + +```bash +embodichain eval-policy outputs/my_policy_ \ + --checkpoint best \ + --viewer \ + --renderer hybrid +``` + +Evaluate an explicit EmbodiChain checkpoint: + +```bash +embodichain eval-policy \ + --checkpoint /path/to/policy.pt \ + --config /path/to/train.yaml \ + --gym-config /path/to/gym.yaml +``` + +### Main arguments + +| Argument | Default | Description | +|---|---|---| +| ``RUN`` | *(optional)* | Training run containing ``run-manifest.json`` | +| ``--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 | + +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. + +--- + ## 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..2caee765a 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 + policy_evaluation cli diff --git a/docs/source/guides/policy_evaluation.md b/docs/source/guides/policy_evaluation.md new file mode 100644 index 000000000..392da7cc2 --- /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. `--renderer` accepts `hybrid`, `fast-rt`, and `offline-rt`. + +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/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/docs/source/tutorial/rl.rst b/docs/source/tutorial/rl.rst index 54ce4590f..64bf26cc4 100644 --- a/docs/source/tutorial/rl.rst +++ b/docs/source/tutorial/rl.rst @@ -264,6 +264,19 @@ All outputs are written to ``./outputs/_/``: - **logs/**: TensorBoard logs - **checkpoints/**: Model checkpoints +- **configs/**: Training config and referenced gym config snapshots +- **evaluations/**: Timestamped policy evaluation reports +- **run-manifest.json**: Training configs and best/latest checkpoint index used by ``eval-policy`` + +A training run can be evaluated Headless or opened in its simulator task: + +.. code-block:: bash + + embodichain eval-policy outputs/_ + embodichain eval-policy outputs/_ --viewer + +See :doc:`../guides/policy_evaluation` for EmbodiChain ``.pt`` training +runs and the external Motion Profile example. Training Process ~~~~~~~~~~~~~~~ @@ -452,9 +465,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 -------- @@ -464,3 +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/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 3b897a3fa..643f2ce48 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-policy", + target="embodichain.learning.rl.policy_evaluation.cli:cli", + help="Evaluate a trained policy in Headless or Viewer mode.", + ), Command( name="annotate-grasp", target="embodichain.toolkits.graspkit.scripts.annotate_grasp:cli", 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/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/policy_evaluation/__init__.py b/embodichain/learning/rl/policy_evaluation/__init__.py new file mode 100644 index 000000000..4dbc590b4 --- /dev/null +++ b/embodichain/learning/rl/policy_evaluation/__init__.py @@ -0,0 +1,33 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""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/policy_evaluation/bridge.py b/embodichain/learning/rl/policy_evaluation/bridge.py new file mode 100644 index 000000000..a336b9ccc --- /dev/null +++ b/embodichain/learning/rl/policy_evaluation/bridge.py @@ -0,0 +1,189 @@ +# ---------------------------------------------------------------------------- +# 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 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 typing import Any + +from dexsim.kit.motion_policy import ( + PolicySpec, + ResolvedPolicy, + ResourceResolver, + RunOptions, + load_scene_config, + parse_policy_spec, + policy_spec_to_dict, + resolve_policy_spec, + run_motion_policy, + scene_config_to_dict, +) + +from .profile import MotionProfile + +__all__ = [ + "MotionEvaluationResult", + "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 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, +) -> 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. + + 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") + 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, + ), + ) + 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 _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/policy_evaluation/cli.py b/embodichain/learning/rl/policy_evaluation/cli.py new file mode 100644 index 000000000..8c9b9c2d2 --- /dev/null +++ b/embodichain/learning/rl/policy_evaluation/cli.py @@ -0,0 +1,523 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Unified policy evaluation for EmbodiChain training runs.""" + +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 numpy as np +import torch + +from embodichain import __version__ +from embodichain.lab.gym.utils.registration import ( + 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, +) +from embodichain.utils.utility import load_config + +from .manifest import RunManifest +from .report import write_evaluation_report + +__all__ = ["cli", "parse_args", "run"] + + +@dataclass(frozen=True) +class EvaluationInput: + """Checkpoint and configuration selected for one evaluation.""" + + checkpoint: Path + profile: str | None + configs: Mapping[str, Path] + run: Path | None + requested_checkpoint: str + 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-policy`` arguments.""" + parser = argparse.ArgumentParser( + prog="embodichain eval-policy", + description="Evaluate an EmbodiChain or external policy checkpoint.", + ) + parser.add_argument("run", nargs="?", help="EmbodiChain training run directory.") + parser.add_argument("--profile", help="Registered external Policy Profile.") + parser.add_argument( + "--checkpoint", + 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 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) + 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") + parser.add_argument( + "--renderer", + choices=("hybrid", "fast-rt", "offline-rt"), + ) + 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: + """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() + _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 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-policy: {error}") from error + print(f"Evaluation report: {report}") + + +def _resolve_input(args: argparse.Namespace) -> EvaluationInput: + if args.run is not None: + manifest = RunManifest.load(args.run) + requested = args.checkpoint or "latest" + 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() + ) + 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() + return EvaluationInput( + checkpoint=checkpoint, + profile=args.profile, + configs=configs, + run=manifest.root, + requested_checkpoint=requested, + selected_checkpoint=selected, + ) + if args.checkpoint is None: + raise ValueError("--checkpoint is required without RUN") + 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 an EmbodiChain checkpoint") + return EvaluationInput( + checkpoint=Path(args.checkpoint).expanduser().resolve(), + profile=args.profile, + configs=configs, + run=None, + requested_checkpoint=args.checkpoint, + selected_checkpoint="explicit", + ) + + +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 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 = _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) + ) + 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=num_envs, + headless=not viewer, + renderer=renderer, + gpu_id=args.gpu_id, + config_dir=train_config.parent, + ) + else: + runtime = build_learning_policy_runtime( + config, + device=device, + num_envs=num_envs, + ) + try: + 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: + profile_options = { + "--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 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 + + 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: + return resolved.run / "evaluations" + return resolved.checkpoint.parent / "evaluations" + + +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", + "inputs": { + **_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, + "metrics": result.metrics, + }, + } + + +def _profile_report( + result: Any, + resolved: EvaluationInput, + device: torch.device, +) -> dict[str, Any]: + import dexsim + + return { + "mode": "viewer" if result.viewer else "headless", + "inputs": { + **_checkpoint_inputs(resolved), + "profile": { + "id": result.profile.profile_id, + "provider_version": result.profile.provider_version, + "provenance": result.profile.provenance, + "warnings": result.profile.warnings, + }, + "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, + }, + } diff --git a/embodichain/learning/rl/policy_evaluation/manifest.py b/embodichain/learning/rl/policy_evaluation/manifest.py new file mode 100644 index 000000000..cc8baa75d --- /dev/null +++ b/embodichain/learning/rl/policy_evaluation/manifest.py @@ -0,0 +1,193 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Index a training run for standalone policy evaluation.""" + +from __future__ import annotations + +import json +import shutil +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +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 + 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", dict(self.configs)) + object.__setattr__(self, "checkpoints", 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, + ) + return cls(root, configs, checkpoints) + + def select_checkpoint(self, requested: str = "latest") -> 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, +) -> 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. + + 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, + "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/policy_evaluation/profile.py b/embodichain/learning/rl/policy_evaluation/profile.py new file mode 100644 index 000000000..a215a084b --- /dev/null +++ b/embodichain/learning/rl/policy_evaluation/profile.py @@ -0,0 +1,128 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""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 typing import Any + +import torch + +__all__ = [ + "MotionProfile", + "MotionProfileRequest", + "build_motion_profile", + "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", 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 + 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: + object.__setattr__(self, "policy_spec", dict(self.policy_spec)) + object.__setattr__(self, "provenance", 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 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/policy_evaluation/report.py b/embodichain/learning/rl/policy_evaluation/report.py new file mode 100644 index 000000000..545b41f2e --- /dev/null +++ b/embodichain/learning/rl/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 timestamped policy evaluation reports.""" + +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}-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/policy_evaluation/viewer.py b/embodichain/learning/rl/policy_evaluation/viewer.py new file mode 100644 index 000000000..509176684 --- /dev/null +++ b/embodichain/learning/rl/policy_evaluation/viewer.py @@ -0,0 +1,480 @@ +# ---------------------------------------------------------------------------- +# 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 task Viewer to 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 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", + "NativeViewerResult", + "evaluate_native_viewer", +] + +_MISSING = object() + + +@dataclass(frozen=True) +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 + effective_duration: float + requested_duration: float | None + episodes: tuple[Mapping[str, float | int | bool | str], ...] + metrics: Mapping[str, float] + + +class EmbodiChainTaskPolicyAdapter: + """Run an EmbodiChain Policy from the task observation in each frame.""" + + def __init__(self, policy: torch.nn.Module, device: torch.device): + self.policy = policy + self.device = device + 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=1, + ) + 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, + ) -> 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) + 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 + self._episode_length = 0 + 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 + + @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.""" + 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(self._episodes) + + def open_viewer(self, title: str) -> None: + """Apply the evaluation title to the task Viewer.""" + self._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.""" + 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" + 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: + """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) + self._reported_metrics.update(metrics) + 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, + } + ) + 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]: + """Return task metrics and completed episode aggregates.""" + result = dict(self._reported_metrics) + if self._episodes: + count = len(self._episodes) + result.update( + { + "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 + + def wait_for_reset_or_close(self) -> str: + """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: + 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_viewer( + runtime: PolicyRuntime, + *, + seed: int, + episodes: int | None, + control_steps: int | None, + duration: float | None, + termination_behavior: str = "auto_reset", +) -> 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: + 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") + + environment = None + adapter = None + evaluator = None + try: + environment = EmbodiChainTaskEnvironment( + runtime.env, + seed=seed, + ) + adapter = EmbodiChainTaskPolicyAdapter( + runtime.policy, + runtime.device, + ) + if duration is not None: + control_steps = math.ceil( + duration / environment.policy_context.policy_dt - 1e-12 + ) + total_steps = 0 + reason = "viewer closed" + options = RunOptions( + headless=False, + 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", + ) + 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 + finally: + if evaluator is not None: + evaluator.close() + elif environment is not None: + if adapter is not None: + adapter.close() + environment.close() + else: + runtime.close() + + simulation_steps = total_steps * context.sim_steps_per_control + return NativeViewerResult( + task_id=runtime.env_id, + reason=reason, + simulation_time=simulation_steps * context.physics_dt, + simulation_steps=simulation_steps, + control_steps=total_steps, + effective_duration=total_steps * context.policy_dt, + requested_duration=duration, + episodes=episode_results, + metrics=metrics, + ) + + +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 the simulator task.""" + 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), + ) diff --git a/embodichain/learning/rl/runtime.py b/embodichain/learning/rl/runtime.py new file mode 100644 index 000000000..ed8f0c7f7 --- /dev/null +++ b/embodichain/learning/rl/runtime.py @@ -0,0 +1,312 @@ +# ---------------------------------------------------------------------------- +# 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 dataclasses import dataclass +from pathlib import Path +from typing import Any + +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__ = [ + "PolicyRuntime", + "build_gym_policy_runtime", + "build_learning_policy_runtime", +] + + +@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 close(self) -> None: + """Close the Environment without terminating the current process.""" + _close_environment(self.env) + + +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_config: dict[str, Any] = {} + else: + env_name = env_block["name"] + env_config = dict(env_block.get("cfg", {})) + return str(env_name), build_learning_env( + str(env_name), + num_envs=num_envs, + device=device, + **env_config, + ) + + +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, + device=device, + num_envs=num_envs, + ) + try: + policy = _build_learning_policy(config["policy"], env, device) + 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, + 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")) + ) + policy = _build_gym_policy( + config["policy"], + env=env, + device=device, + observation_dim=observation_dim, + action_dim=environment_action_dim, + ) + except Exception: + _close_environment(env) + 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, + ) + + +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 41c44a3f1..9556efd31 100644 --- a/embodichain/learning/rl/train.py +++ b/embodichain/learning/rl/train.py @@ -20,16 +20,18 @@ import os import time from collections.abc import Sequence +from copy import deepcopy from pathlib import Path import numpy as np 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.policy_evaluation.manifest import ( + write_run_manifest, +) from embodichain.learning.rl.algo import ( RolloutKind, build_algo, @@ -39,9 +41,12 @@ DifferentiableTrainer, DifferentiableTrainerCfg, ) -from embodichain.learning.rl.env import build_learning_env +from embodichain.learning.rl.runtime import ( + _build_learning_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 ( @@ -49,14 +54,13 @@ 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 +_CAMERA_RECORDERS = {"record_camera_data", "record_camera_data_async"} + def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: """Parse command-line arguments. @@ -114,51 +118,33 @@ 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 _event_params( + event_info: dict, + *, + run_base: str | Path, + phase: str, +) -> dict: + """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: + params.setdefault("save_path", str(Path(run_base) / "videos" / phase)) + return params 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,32 +168,25 @@ 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: - 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, ) - policy = _build_learning_policy(policy_block, env, device) algorithm = build_algo( algorithm_block["name"], dict(algorithm_block.get("cfg", {})), @@ -291,7 +270,7 @@ 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() finally: writer.close() if use_wandb: @@ -299,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( @@ -307,7 +288,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 +297,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 +310,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 +426,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,23 +438,37 @@ 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.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) @@ -498,59 +476,7 @@ def train_from_config( 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() @@ -581,7 +507,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 @@ -597,7 +523,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 @@ -648,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: @@ -655,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: @@ -681,8 +610,35 @@ def train_from_config( if distributed and torch.distributed.is_initialized(): torch.distributed.destroy_process_group() - if rank == 0: - logger.log_info("Training finished") + if summary is not None: + _write_policy_run_manifest( + run_base, + config_path, + summary, + gym_config=gym_config_path, + ) + if rank == 0: + logger.log_info("Training finished") + + +def _write_policy_run_manifest( + run_base: str | Path, + config_path: str | Path, + summary: dict, + *, + gym_config: str | Path | None = None, +) -> Path: + """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") + return write_run_manifest( + run_base, + train_config=config_path, + gym_config=gym_config, + latest_checkpoint=latest, + best_checkpoint=summary.get("best_checkpoint_path"), + ) def cli(argv: Sequence[str] | None = None) -> None: 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/examples/learning/policy_evaluation/README.md b/examples/learning/policy_evaluation/README.md new file mode 100644 index 000000000..8e3155b00 --- /dev/null +++ b/examples/learning/policy_evaluation/README.md @@ -0,0 +1,137 @@ +# 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 +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/policy_evaluation/prepare_resources.py +python examples/learning/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 | +| 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: + +```bash +python examples/learning/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/policy_evaluation/prepare_resources.py \ + --output /tmp/anymal_c_velocity + +ANYMAL_C_EXAMPLE_CACHE=/tmp/anymal_c_velocity \ + python examples/learning/policy_evaluation/eval_policy.py --viewer +``` + +## Execution pipeline + +```mermaid +flowchart LR + 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] + 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/policy_evaluation/anymal_c/__init__.py b/examples/learning/policy_evaluation/anymal_c/__init__.py new file mode 100644 index 000000000..0303a1211 --- /dev/null +++ b/examples/learning/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 ANYmal-C velocity Motion Profile.""" + +from __future__ import annotations + +from embodichain.learning.rl.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/policy_evaluation/anymal_c/profile.py b/examples/learning/policy_evaluation/anymal_c/profile.py new file mode 100644 index 000000000..2f0e89e73 --- /dev/null +++ b/examples/learning/policy_evaluation/anymal_c/profile.py @@ -0,0 +1,278 @@ +# ---------------------------------------------------------------------------- +# 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 a 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.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.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, + 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"]) + 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.""" + robot = context.robot + if robot is None: + raise RuntimeError("ANYmal-C robot description is required") + self.joints = JointMap.from_joint_names( + 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.""" + observation = self._build_observation(frame) + tensor = torch.from_numpy(observation).to(self.device).unsqueeze(0) + with torch.inference_mode(): + output = self._model()(tensor) + 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, + ) + 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) + 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/policy_evaluation/eval_policy.py b/examples/learning/policy_evaluation/eval_policy.py new file mode 100644 index 000000000..b01048b30 --- /dev/null +++ b/examples/learning/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. +# ---------------------------------------------------------------------------- + +"""Evaluate the public ANYmal-C checkpoint 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-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.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/policy_evaluation/prepare_resources.py b/examples/learning/policy_evaluation/prepare_resources.py new file mode 100644 index 000000000..7e000e2a4 --- /dev/null +++ b/examples/learning/policy_evaluation/prepare_resources.py @@ -0,0 +1,209 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Prepare the pinned public policy and 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") +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]: + """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 + for relative, digest in SHA256.items(): + _verify_sha256(checkout / relative, digest) + _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(): + 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" + ) + if tracked_changes == "": + _status(f"Using cached revision {revision[:8]} from {checkout}") + return + 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) + + _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( + 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/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_anymal_c_example.py b/tests/learning/rl/policy_evaluation/test_anymal_c_example.py new file mode 100644 index 000000000..c07d6ef9f --- /dev/null +++ b/tests/learning/rl/policy_evaluation/test_anymal_c_example.py @@ -0,0 +1,181 @@ +# ---------------------------------------------------------------------------- +# 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 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, + PolicyContext, + RobotDescription, + RobotState, + parse_policy_spec, +) + +from embodichain.learning.rl.policy_evaluation import ( + MotionProfileRequest, + build_motion_profile, +) + +_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 test_anymal_c_profile_builds_and_runs_torchscript(tmp_path, monkeypatch): + example_root = ( + Path(__file__).resolve().parents[4] / "examples/learning/policy_evaluation" + ) + monkeypatch.syspath_prepend(str(example_root)) + from anymal_c.profile import ( + AnymalCVelocityAdapter, + ) + from anymal_c import register + + 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") + + 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"] + 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/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/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..d36059fee --- /dev/null +++ b/tests/learning/rl/policy_evaluation/test_cli.py @@ -0,0 +1,130 @@ +# ---------------------------------------------------------------------------- +# 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)))) + + +@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( + ( + "--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 new file mode 100644 index 000000000..ca35bab74 --- /dev/null +++ b/tests/learning/test_runtime.py @@ -0,0 +1,125 @@ +# ---------------------------------------------------------------------------- +# 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, +) + + +def _policy_config() -> dict: + network = { + "type": "mlp", + "network_cfg": {"hidden_sizes": [8], "activation": "relu"}, + } + return { + "name": "actor_critic", + "actor": network, + "critic": network, + } + + +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_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/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/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 diff --git a/tests/test_main.py b/tests/test_main.py index 5466e7c11..b35b23520 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -28,6 +28,7 @@ "benchmark", "data", "decompose-urdf", + "eval-policy", "preview-asset", "preview_lerobot_data", "run-env",