diff --git a/.github/workflows/public-ci.yml b/.github/workflows/public-ci.yml index 6218f80dd..938fb3af0 100644 --- a/.github/workflows/public-ci.yml +++ b/.github/workflows/public-ci.yml @@ -40,6 +40,17 @@ env: GITLEAKS_SHA256: "5bc41815076e6ed6ef8fbecc9d9b75bcae31f39029ceb55da08086315316e3ba" jobs: + model-licenses: + name: Model release licenses + if: github.repository == 'NVlabs/ProtoMotions' + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: { lfs: false } + - uses: actions/setup-python@v5 + with: { python-version: "3.10" } + - run: python scripts/check_model_licenses.py + # Catches a credential in a contributed pull request before it is merged -- # the one class of mistake that rewriting history afterwards does not undo, # because the commit stays reachable from the pull request ref forever. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 891084c37..ba1ac1c05 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -2,6 +2,44 @@ We started ProtoMotions3 to promote open research and help all members of the community. All participation in this project is open and not bound to corporate affiliation. +## Automated Checks + +Linting, formatting, license headers and spell checking all run through +[pre-commit](https://pre-commit.com/). Install the hook once and every commit is checked +before it is created: + +```bash +pip install pre-commit +pre-commit install +``` + +**Run these tools through pre-commit rather than invoking them yourself.** +`.pre-commit-config.yaml` pins an exact version of each tool, and pre-commit installs +those versions into an isolated environment of its own. A `ruff` that happens to be on +your `PATH` is almost certainly a different version, and the formatter's output changes +between releases — so a file your local `ruff` wants to reformat may be one the pinned +version already considers correct. Going through pre-commit is what makes everyone's +result the same, and the same as CI's. + +To check specific files without committing: + +```bash +pre-commit run --files ... +``` + +Prefer that over `pre-commit run --all-files`: parts of the tree predate these checks, so +running over everything produces a large diff unrelated to your change. + +## Releasing models + +Code remains Apache-2.0. NVIDIA-provided model releases default to the NVIDIA +Open Model License unless explicitly approved otherwise. Follow the +[model-release checklist](legal/MODEL-LICENSING.md): include a License section +in the model card, the complete license and an attribution notice with every +model package, including exported models and adapters. Document any exception +in the card and the reviewed exception manifest. Run +`python scripts/check_model_licenses.py` before submitting a model release. + ## Code Reviews All submissions, including submissions by project members, require review. We use GitHub pull requests for this purpose. Consult @@ -59,4 +97,4 @@ All submissions, including submissions by project members, require review. We us personal information I submit with it, including my sign-off) is maintained indefinitely and may be redistributed consistent with this project or the open source license(s) involved. - ``` \ No newline at end of file + ``` diff --git a/README.md b/README.md index 3eb18647f..5118eef4f 100644 --- a/README.md +++ b/README.md @@ -234,7 +234,12 @@ We welcome contributions! Please read our [**Contributing Guide**](CONTRIBUTING. ## License -ProtoMotions3 is released under the [**Apache-2.0 License**](LICENSE.md). +ProtoMotions3 source code is released under the [**Apache-2.0 License**](LICENSE.md). + +NVIDIA-provided pretrained models are released under the **NVIDIA Open Model +License**, unless their model card explicitly states otherwise. Each model +directory includes its model card, license and attribution notice. See +[model licensing and redistribution](legal/MODEL-LICENSING.md). Third-party software and bundled asset notices are listed in [legal/](legal/), including Unitree, BeyondMimic, Isaac Lab, and SMPL/SMPL-H attribution and license notices. diff --git a/data/pretrained_models/AGENTS.md b/data/pretrained_models/AGENTS.md new file mode 100644 index 000000000..fb40aef11 --- /dev/null +++ b/data/pretrained_models/AGENTS.md @@ -0,0 +1,11 @@ +# Model release requirements + +- Follow [the model licensing policy](../../legal/MODEL-LICENSING.md). +- NVIDIA-provided model releases default to the NVIDIA Open Model License; + source code remains Apache-2.0. +- Every model needs a License section in its model card, a complete license + copy and an attribution notice. Include these in standalone exports too. +- Alternative licenses require a reviewed exception recorded in both the card + and the exception manifest. Preserve applicable third-party terms. +- Run `python scripts/check_model_licenses.py` from the repository root. +- Do not modify weights or serialized configs just to change licensing notices. diff --git a/data/pretrained_models/gpc_prior/soma_bones/LICENSE.pdf b/data/pretrained_models/gpc_prior/soma_bones/LICENSE.pdf new file mode 100644 index 000000000..d803837f7 Binary files /dev/null and b/data/pretrained_models/gpc_prior/soma_bones/LICENSE.pdf differ diff --git a/data/pretrained_models/gpc_prior/soma_bones/MODEL_CARD.md b/data/pretrained_models/gpc_prior/soma_bones/MODEL_CARD.md index 8078a2c1b..a4df71cd8 100644 --- a/data/pretrained_models/gpc_prior/soma_bones/MODEL_CARD.md +++ b/data/pretrained_models/gpc_prior/soma_bones/MODEL_CARD.md @@ -1,5 +1,17 @@ # SOMA BONES-SEED GPC Prior +## License + +Model license: NVIDIA Open Model License + +The model weights, checkpoints, serialized model configurations and exported +models are licensed under the [NVIDIA Open Model License](LICENSE.pdf). +See [NOTICE.txt](NOTICE.txt) for attribution. Include both files when +redistributing these artifacts. + +ProtoMotions source code remains Apache-2.0. Training datasets and robot assets +have separate terms; this model license does not grant rights to those materials. + ## Overview This checkpoint is an autoregressive Generative Pre-trained Controller (GPC) diff --git a/data/pretrained_models/gpc_prior/soma_bones/NOTICE.txt b/data/pretrained_models/gpc_prior/soma_bones/NOTICE.txt new file mode 100644 index 000000000..559a7c65e --- /dev/null +++ b/data/pretrained_models/gpc_prior/soma_bones/NOTICE.txt @@ -0,0 +1 @@ +Licensed by NVIDIA Corporation under the NVIDIA Open Model License diff --git a/data/pretrained_models/masked_mimic/smpl/LICENSE.pdf b/data/pretrained_models/masked_mimic/smpl/LICENSE.pdf new file mode 100644 index 000000000..d803837f7 Binary files /dev/null and b/data/pretrained_models/masked_mimic/smpl/LICENSE.pdf differ diff --git a/data/pretrained_models/masked_mimic/smpl/MODEL_CARD.md b/data/pretrained_models/masked_mimic/smpl/MODEL_CARD.md index 35b073e7c..6947182c6 100644 --- a/data/pretrained_models/masked_mimic/smpl/MODEL_CARD.md +++ b/data/pretrained_models/masked_mimic/smpl/MODEL_CARD.md @@ -1,5 +1,17 @@ # SMPL MaskedMimic +## License + +Model license: NVIDIA Open Model License + +The model weights, checkpoints, serialized model configurations and exported +models are licensed under the [NVIDIA Open Model License](LICENSE.pdf). +See [NOTICE.txt](NOTICE.txt) for attribution. Include both files when +redistributing these artifacts. + +ProtoMotions source code remains Apache-2.0. Training datasets and robot assets +have separate terms; this model license does not grant rights to those materials. + ## Overview This is a MaskedMimic controller for the SMPL humanoid. The policy produces diff --git a/data/pretrained_models/masked_mimic/smpl/NOTICE.txt b/data/pretrained_models/masked_mimic/smpl/NOTICE.txt new file mode 100644 index 000000000..559a7c65e --- /dev/null +++ b/data/pretrained_models/masked_mimic/smpl/NOTICE.txt @@ -0,0 +1 @@ +Licensed by NVIDIA Corporation under the NVIDIA Open Model License diff --git a/data/pretrained_models/motion_tracker/g1-bones-deploy/LICENSE.pdf b/data/pretrained_models/motion_tracker/g1-bones-deploy/LICENSE.pdf new file mode 100644 index 000000000..d803837f7 Binary files /dev/null and b/data/pretrained_models/motion_tracker/g1-bones-deploy/LICENSE.pdf differ diff --git a/data/pretrained_models/motion_tracker/g1-bones-deploy/MODEL_CARD.md b/data/pretrained_models/motion_tracker/g1-bones-deploy/MODEL_CARD.md index 0ea4c5610..fc03f2e03 100644 --- a/data/pretrained_models/motion_tracker/g1-bones-deploy/MODEL_CARD.md +++ b/data/pretrained_models/motion_tracker/g1-bones-deploy/MODEL_CARD.md @@ -1,5 +1,17 @@ # G1 BONES-SEED Deployment Tracker +## License + +Model license: NVIDIA Open Model License + +The model weights, checkpoints, serialized model configurations and exported +models are licensed under the [NVIDIA Open Model License](LICENSE.pdf). +See [NOTICE.txt](NOTICE.txt) for attribution. Include both files when +redistributing these artifacts. + +ProtoMotions source code remains Apache-2.0. Training datasets and robot assets +have separate terms; this model license does not grant rights to those materials. + ## Overview This is a general motion-tracking policy for the Unitree G1 humanoid. It is the diff --git a/data/pretrained_models/motion_tracker/g1-bones-deploy/NOTICE.txt b/data/pretrained_models/motion_tracker/g1-bones-deploy/NOTICE.txt new file mode 100644 index 000000000..559a7c65e --- /dev/null +++ b/data/pretrained_models/motion_tracker/g1-bones-deploy/NOTICE.txt @@ -0,0 +1 @@ +Licensed by NVIDIA Corporation under the NVIDIA Open Model License diff --git a/data/pretrained_models/motion_tracker/smpl-terrains/LICENSE.pdf b/data/pretrained_models/motion_tracker/smpl-terrains/LICENSE.pdf new file mode 100644 index 000000000..d803837f7 Binary files /dev/null and b/data/pretrained_models/motion_tracker/smpl-terrains/LICENSE.pdf differ diff --git a/data/pretrained_models/motion_tracker/smpl-terrains/MODEL_CARD.md b/data/pretrained_models/motion_tracker/smpl-terrains/MODEL_CARD.md index 0ba15f010..80ad6c5c3 100644 --- a/data/pretrained_models/motion_tracker/smpl-terrains/MODEL_CARD.md +++ b/data/pretrained_models/motion_tracker/smpl-terrains/MODEL_CARD.md @@ -1,5 +1,17 @@ # SMPL AMASS Terrain Tracker +## License + +Model license: NVIDIA Open Model License + +The model weights, checkpoints, serialized model configurations and exported +models are licensed under the [NVIDIA Open Model License](LICENSE.pdf). +See [NOTICE.txt](NOTICE.txt) for attribution. Include both files when +redistributing these artifacts. + +ProtoMotions source code remains Apache-2.0. Training datasets and robot assets +have separate terms; this model license does not grant rights to those materials. + ## Overview This is a general SMPL motion tracker trained on a mixture of procedural terrain diff --git a/data/pretrained_models/motion_tracker/smpl-terrains/NOTICE.txt b/data/pretrained_models/motion_tracker/smpl-terrains/NOTICE.txt new file mode 100644 index 000000000..559a7c65e --- /dev/null +++ b/data/pretrained_models/motion_tracker/smpl-terrains/NOTICE.txt @@ -0,0 +1 @@ +Licensed by NVIDIA Corporation under the NVIDIA Open Model License diff --git a/data/pretrained_models/motion_tracker/smpl/LICENSE.pdf b/data/pretrained_models/motion_tracker/smpl/LICENSE.pdf new file mode 100644 index 000000000..d803837f7 Binary files /dev/null and b/data/pretrained_models/motion_tracker/smpl/LICENSE.pdf differ diff --git a/data/pretrained_models/motion_tracker/smpl/MODEL_CARD.md b/data/pretrained_models/motion_tracker/smpl/MODEL_CARD.md index 7589e0812..61ed6719f 100644 --- a/data/pretrained_models/motion_tracker/smpl/MODEL_CARD.md +++ b/data/pretrained_models/motion_tracker/smpl/MODEL_CARD.md @@ -1,5 +1,17 @@ # SMPL AMASS Flat-Terrain Tracker +## License + +Model license: NVIDIA Open Model License + +The model weights, checkpoints, serialized model configurations and exported +models are licensed under the [NVIDIA Open Model License](LICENSE.pdf). +See [NOTICE.txt](NOTICE.txt) for attribution. Include both files when +redistributing these artifacts. + +ProtoMotions source code remains Apache-2.0. Training datasets and robot assets +have separate terms; this model license does not grant rights to those materials. + ## Overview This is a general motion-tracking policy for the SMPL humanoid on flat terrain. diff --git a/data/pretrained_models/motion_tracker/smpl/NOTICE.txt b/data/pretrained_models/motion_tracker/smpl/NOTICE.txt new file mode 100644 index 000000000..559a7c65e --- /dev/null +++ b/data/pretrained_models/motion_tracker/smpl/NOTICE.txt @@ -0,0 +1 @@ +Licensed by NVIDIA Corporation under the NVIDIA Open Model License diff --git a/data/pretrained_models/motion_tracker/soma-bones/LICENSE.pdf b/data/pretrained_models/motion_tracker/soma-bones/LICENSE.pdf new file mode 100644 index 000000000..d803837f7 Binary files /dev/null and b/data/pretrained_models/motion_tracker/soma-bones/LICENSE.pdf differ diff --git a/data/pretrained_models/motion_tracker/soma-bones/MODEL_CARD.md b/data/pretrained_models/motion_tracker/soma-bones/MODEL_CARD.md index dbf6de433..5eb3f8d0b 100644 --- a/data/pretrained_models/motion_tracker/soma-bones/MODEL_CARD.md +++ b/data/pretrained_models/motion_tracker/soma-bones/MODEL_CARD.md @@ -1,5 +1,17 @@ # SOMA BONES-SEED Motion Tracker +## License + +Model license: NVIDIA Open Model License + +The model weights, checkpoints, serialized model configurations and exported +models are licensed under the [NVIDIA Open Model License](LICENSE.pdf). +See [NOTICE.txt](NOTICE.txt) for attribution. Include both files when +redistributing these artifacts. + +ProtoMotions source code remains Apache-2.0. Training datasets and robot assets +have separate terms; this model license does not grant rights to those materials. + ## Overview This is a continuous-action general motion tracker for the SOMA 23-body diff --git a/data/pretrained_models/motion_tracker/soma-bones/NOTICE.txt b/data/pretrained_models/motion_tracker/soma-bones/NOTICE.txt new file mode 100644 index 000000000..559a7c65e --- /dev/null +++ b/data/pretrained_models/motion_tracker/soma-bones/NOTICE.txt @@ -0,0 +1 @@ +Licensed by NVIDIA Corporation under the NVIDIA Open Model License diff --git a/data/pretrained_models/motion_tracker/soma_bones_fsq/LICENSE.pdf b/data/pretrained_models/motion_tracker/soma_bones_fsq/LICENSE.pdf new file mode 100644 index 000000000..d803837f7 Binary files /dev/null and b/data/pretrained_models/motion_tracker/soma_bones_fsq/LICENSE.pdf differ diff --git a/data/pretrained_models/motion_tracker/soma_bones_fsq/MODEL_CARD.md b/data/pretrained_models/motion_tracker/soma_bones_fsq/MODEL_CARD.md index 8454c7153..926adae72 100644 --- a/data/pretrained_models/motion_tracker/soma_bones_fsq/MODEL_CARD.md +++ b/data/pretrained_models/motion_tracker/soma_bones_fsq/MODEL_CARD.md @@ -1,5 +1,17 @@ # SOMA BONES-SEED FSQ Motion Tracker +## License + +Model license: NVIDIA Open Model License + +The model weights, checkpoints, serialized model configurations and exported +models are licensed under the [NVIDIA Open Model License](LICENSE.pdf). +See [NOTICE.txt](NOTICE.txt) for attribution. Include both files when +redistributing these artifacts. + +ProtoMotions source code remains Apache-2.0. Training datasets and robot assets +have separate terms; this model license does not grant rights to those materials. + ## Overview This is a SOMA motion tracker with a finite-scalar-quantization (FSQ) bottleneck. diff --git a/data/pretrained_models/motion_tracker/soma_bones_fsq/NOTICE.txt b/data/pretrained_models/motion_tracker/soma_bones_fsq/NOTICE.txt new file mode 100644 index 000000000..559a7c65e --- /dev/null +++ b/data/pretrained_models/motion_tracker/soma_bones_fsq/NOTICE.txt @@ -0,0 +1 @@ +Licensed by NVIDIA Corporation under the NVIDIA Open Model License diff --git a/data/pretrained_models/motion_tracker/soma_bones_fsq_amp_muon/LICENSE.pdf b/data/pretrained_models/motion_tracker/soma_bones_fsq_amp_muon/LICENSE.pdf new file mode 100644 index 000000000..d803837f7 Binary files /dev/null and b/data/pretrained_models/motion_tracker/soma_bones_fsq_amp_muon/LICENSE.pdf differ diff --git a/data/pretrained_models/motion_tracker/soma_bones_fsq_amp_muon/MODEL_CARD.md b/data/pretrained_models/motion_tracker/soma_bones_fsq_amp_muon/MODEL_CARD.md index 0beae39fc..a04800cee 100644 --- a/data/pretrained_models/motion_tracker/soma_bones_fsq_amp_muon/MODEL_CARD.md +++ b/data/pretrained_models/motion_tracker/soma_bones_fsq_amp_muon/MODEL_CARD.md @@ -1,5 +1,17 @@ # SOMA BONES-SEED FSQ AMP Tracker +## License + +Model license: NVIDIA Open Model License + +The model weights, checkpoints, serialized model configurations and exported +models are licensed under the [NVIDIA Open Model License](LICENSE.pdf). +See [NOTICE.txt](NOTICE.txt) for attribution. Include both files when +redistributing these artifacts. + +ProtoMotions source code remains Apache-2.0. Training datasets and robot assets +have separate terms; this model license does not grant rights to those materials. + ## Overview This is a SOMA motion tracker that combines an FSQ actor bottleneck with diff --git a/data/pretrained_models/motion_tracker/soma_bones_fsq_amp_muon/NOTICE.txt b/data/pretrained_models/motion_tracker/soma_bones_fsq_amp_muon/NOTICE.txt new file mode 100644 index 000000000..559a7c65e --- /dev/null +++ b/data/pretrained_models/motion_tracker/soma_bones_fsq_amp_muon/NOTICE.txt @@ -0,0 +1 @@ +Licensed by NVIDIA Corporation under the NVIDIA Open Model License diff --git a/data/scripts/keypoint_utils.py b/data/scripts/keypoint_utils.py index 5347bd4a4..5f693b561 100644 --- a/data/scripts/keypoint_utils.py +++ b/data/scripts/keypoint_utils.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import torch +from protomotions.utils.retargeting_keypoints import SOMA_PELVIS_KEYPOINT_SHIFT from protomotions.utils.rotations import quaternion_to_matrix from typing import List, Dict, Tuple @@ -580,11 +581,13 @@ def extract_keypoints_from_motion_soma_skel( keypoint_orientations_quat, w_last=True ) - # Pelvis offset: shift backward in local +y to create forward bias on all - # Body keypoints relative to root. + # Pelvis offset: store the pelvis keypoint 7 cm behind the Hips joint, as + # every SOMA keypoint file so far has it. The retargeters move it back onto + # the joint when they load the file (soma_pelvis_keypoint_on_hips in + # protomotions/utils/retargeting_keypoints.py). root_idx = conceptual_keypoint_names.index("pelvis") root_offset = torch.tensor( - [0.0, 0.07, 0.0], device=device, dtype=keypoint_positions.dtype + SOMA_PELVIS_KEYPOINT_SHIFT, device=device, dtype=keypoint_positions.dtype ) R_root = keypoint_orientations_mat[:, root_idx, :, :] keypoint_positions[:, root_idx, :] += torch.einsum( @@ -671,9 +674,10 @@ def extract_keypoints_from_motion_soma_skel( dim=1, ) - # 0.225 * 0.8 = 0.18m, matching retargeter torso_aux x. - # A 0.30m offset would produce 0.24m after scaling, creating - # a 6cm x-mismatch with the torso_aux (0.18) that pushes pelvis backward. + # 0.225 m in front of the shifted pelvis keypoint, i.e. 0.155 m in front + # of the Hips joint (0.124 m after G1's 0.8 scaling, against G1's 0.18 m + # torso_aux). It was sized as 0.225 * 0.8 = 0.18 m for a G1 pelvis matched + # to the shifted keypoint, which the retargeters no longer do. pelvis_idx = conceptual_keypoint_names.index("pelvis") p_local_pelvis = torch.tensor( [0.0, -0.225, 0.0], device=device, dtype=keypoint_positions.dtype diff --git a/docs/source/api_reference/scripts.rst b/docs/source/api_reference/scripts.rst index 97260ce1e..95e352f20 100644 --- a/docs/source/api_reference/scripts.rst +++ b/docs/source/api_reference/scripts.rst @@ -70,10 +70,18 @@ For kinematic motion playback without physics training state: **Keyboard Controls:** +Click inside the simulation viewer to give it keyboard focus, then use the +task's key bindings. + * **W/A/S/D** - Move keyboard-controlled task targets when enabled * **R** - Reset environments * **Q** - Quit +**MuJoCo only:** press **.** to toggle between ProtoMotions key bindings +(default) and MuJoCo's native viewer key bindings. Only the selected set is +active. Switching back to ProtoMotions key bindings restores the initial +visualization settings. Mouse orbit and zoom work in either mode. + train_slurm.py -------------- @@ -106,4 +114,3 @@ See Also * :doc:`../getting_started/quickstart` - Getting started guide * :doc:`../user_guide/configuration` - Configuration system - diff --git a/docs/source/getting_started/pretrained_models.rst b/docs/source/getting_started/pretrained_models.rst index 633600823..a880a155b 100644 --- a/docs/source/getting_started/pretrained_models.rst +++ b/docs/source/getting_started/pretrained_models.rst @@ -4,6 +4,13 @@ Pre-trained Models Each shipped checkpoint has a model card beside it with the training recipe, input and output contract, artifact roles, intended use, and limitations. +NVIDIA-provided pretrained models use the **NVIDIA Open Model License** unless +their card explicitly specifies otherwise; ProtoMotions source code remains +Apache-2.0. Each model directory includes a full license and attribution notice. +Keep these files with redistributed checkpoints and exported models. See the +`model licensing policy `__ +for details. Training datasets and robot assets retain their separate terms. + Unless its model card says otherwise, a policy should be used on its training simulator. The recommended checkpoints were trained or fine-tuned in IsaacLab. The G1 deployment tracker is the only current policy trained with the full diff --git a/docs/source/tutorials/workflows/retargeting_pyroki.rst b/docs/source/tutorials/workflows/retargeting_pyroki.rst index b184dcf46..c996a4128 100644 --- a/docs/source/tutorials/workflows/retargeting_pyroki.rst +++ b/docs/source/tutorials/workflows/retargeting_pyroki.rst @@ -412,6 +412,20 @@ Comparing ``batch_retarget_to_g1_from_keypoints.py`` and * onp.array([1.1, 1.1, 1.0])[None, None, :] ) + Take the vertical factor for the pelvis and legs from the robot/source + pelvis-to-ankle drop, not from the hip-joint height. Knee bend is very + sensitive to it: near a straight leg, a pelvis target only 1 cm too low + already bends G1's knees by about 20 degrees. Keep the x and y factors + equal, so the scaling does not depend on which way the motion is heading. + + The robot matches the pelvis keypoint with its pelvis origin. Keypoint + extraction stores the SOMA pelvis keypoint 7 cm behind the Hips joint, so + ``load_motion_data`` moves it back onto the joint with + ``soma_pelvis_keypoint_on_hips`` from + ``protomotions/utils/retargeting_keypoints.py``. Call it in a new robot's + script too; otherwise every SOMA target lands about 7 cm too far forward + relative to the pelvis. + 3. **Auxiliary Point Offsets** Update hand and torso auxiliary point offsets in ``pc_alignment_cost()``: diff --git a/docs/source/user_guide/configuration.rst b/docs/source/user_guide/configuration.rst index b12a3f56c..038d07ca3 100644 --- a/docs/source/user_guide/configuration.rst +++ b/docs/source/user_guide/configuration.rst @@ -124,6 +124,28 @@ Robot configs in ``protomotions/robot_configs/`` define the robot: Using Configurations -------------------- +Marker Body Mass and Inertia +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Before importing a robot, the shared loading code resolves its MJCF masses. +A zero-mass body (typically a fixed tracking marker) receives **0.001 kg (1 g)** +and diagonal inertia of **4e-8 kg m²** at its local origin. These inertias +correspond to a 1 g solid sphere with a 1 cm radius; the visual marker geometry +does not need to change. This convention avoids the extremely small inertia +of a visualization-sized sphere. + +Source MJCF files are unchanged. The loader prepares a temporary asset with +explicit marker inertias, which IsaacGym, IsaacLab, Newton, Genesis and MuJoCo +all consume. IsaacLab performs this step before MJCF-to-USD conversion and +invalidates older conversion caches. PhysX articulation links cannot have zero +mass; leaving a marker massless can silently give it the simulator's default +mass instead. + +Ordinary physical bodies retain their authored or geometry-derived mass. +Disabling mass randomization does not make a body massless: it uses the MJCF +values after normalization. When enabled, body-mass randomization overrides +the selected bodies and scales their inertia from the positive source properties. + Basic Usage ~~~~~~~~~~~ diff --git a/docs/source/user_guide/experiments.rst b/docs/source/user_guide/experiments.rst index 057e43308..781f502fe 100644 --- a/docs/source/user_guide/experiments.rst +++ b/docs/source/user_guide/experiments.rst @@ -27,6 +27,19 @@ reproduce the motion while maintaining physical plausibility. * ``mlp_bm_l2c2.py`` - BeyondMimic policy with a torso anchor and L2C2 smoothness regularization, aimed at sim2real * ``fsq.py`` - FSQ tracker over max-coordinates proprioception and target poses +The BM recipe uses ``q_target = q_default + 0.25 * effort_limit / stiffness * action`` +without clipping the raw action or changing the robot's PD gains. Its action-rate +reward is ``-0.1 * sum((action - previous_action)**2)`` on raw policy outputs, +with no mean, square root, or timestep division. This is distinct from the L2C2 +model loss and from the processed-action L2 norm used by other recipes. + +The scale change applies to all newly built configs using +``make_bm_pd_action_config``. The raw action-rate reward is selected explicitly +by the BM L2C2 recipe; other recipes keep their existing reward bindings. +Resume and inference use the action +scales and reward bindings saved in ``resolved_configs.pt``; existing pretrained +configs are not rewritten. Use a new experiment name to train with this recipe. + **Example command**: .. code-block:: bash diff --git a/examples/experiments/masked_mimic/transformer.py b/examples/experiments/masked_mimic/transformer.py index d645e251b..e6542e602 100644 --- a/examples/experiments/masked_mimic/transformer.py +++ b/examples/experiments/masked_mimic/transformer.py @@ -205,7 +205,8 @@ def env_config(robot_cfg: RobotConfig, args: argparse.Namespace) -> EnvConfig: # Termination components configuration termination_components = { - "tracking_error": tracking_error_term_factory(threshold=0.25), + # Always ran at 0.5 m: the 0.25 m set here before never took effect. + "tracking_error": tracking_error_term_factory(threshold=0.5), } # Reward components diff --git a/examples/experiments/mimic/mlp_bm_l2c2.py b/examples/experiments/mimic/mlp_bm_l2c2.py index 47c6e1c37..6de5c857c 100644 --- a/examples/experiments/mimic/mlp_bm_l2c2.py +++ b/examples/experiments/mimic/mlp_bm_l2c2.py @@ -74,7 +74,7 @@ def env_config(robot_cfg: RobotConfig, args: argparse.Namespace) -> EnvConfig: max_coords_obs_factory, mimic_target_poses_max_coords_factory, previous_actions_factory, - action_smoothness_factory, + action_rate_l2_factory, global_anchor_ori_rew_factory, relative_body_pos_rew_factory, relative_body_ori_rew_factory, @@ -131,7 +131,7 @@ def env_config(robot_cfg: RobotConfig, args: argparse.Namespace) -> EnvConfig: with_velocities=True, with_relative=True, ), - # Common observations (processed actions after tanh/clamp) + # Common observations (joint position targets after PD action scaling) "historical_previous_processed_actions": previous_actions_factory( history_steps=1, processed=True ), @@ -164,7 +164,7 @@ def env_config(robot_cfg: RobotConfig, args: argparse.Namespace) -> EnvConfig: weight=1.0, sigma=3.14, ), - "action_rate": action_smoothness_factory(weight=-0.1), + "action_rate": action_rate_l2_factory(weight=-0.1), "limits_dof_pos": MdpComponent( compute_func=compute_soft_pos_limit_rew, dynamic_vars={ diff --git a/legal/MODEL-LICENSING.md b/legal/MODEL-LICENSING.md new file mode 100644 index 000000000..53b5b564b --- /dev/null +++ b/legal/MODEL-LICENSING.md @@ -0,0 +1,50 @@ +# Model release licensing + +ProtoMotions source code remains under [Apache-2.0](../LICENSE.md). +NVIDIA-provided pretrained model weights, checkpoints, serialized model +configurations and exported models are released under the **NVIDIA Open Model +License**, unless a model card explicitly specifies an approved exception. +This is a release policy for models distributed by this project, not a license +imposed on every model that someone trains using the Apache-licensed code. + +The bundled [agreement](NVIDIA-Open-Model-License.pdf) is the unmodified +[official NVIDIA PDF](https://www.nvidia.com/content/dam/en-zz/Solutions/license-agreements/enterprise-software/nvidia-open-model-license-agreements-24-10-2025.pdf), +version October 24, 2025. The [official license page](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/) +provides an HTML version. Source files, including model architectures and +experiment scripts, keep their Apache SPDX headers. Third-party datasets, +robot assets and other separately licensed components retain their own terms. +This policy does not revoke rights previously granted under another license. + +## Publishing a model + +Keep each model under `data/pretrained_models///` and include: + +1. A `MODEL_CARD.md` with a `## License` section containing + `Model license: NVIDIA Open Model License`. +2. A byte-for-byte copy of `legal/NVIDIA-Open-Model-License.pdf` named + `LICENSE.pdf` beside that card. +3. A `NOTICE.txt` containing the attribution: + `Licensed by NVIDIA Corporation under the NVIDIA Open Model License`. +4. Links to `LICENSE.pdf` and `NOTICE.txt` from the card's License section. + Keep the card's training-data provenance and applicable third-party notices. + +Include the license, notice and model card in standalone downloads and +exported-model packages too; a weights-only or ONNX-only upload is not a complete +release. Preserve applicable upstream terms when distributing fine-tuned models, +adapters or derivatives; do not infer dataset redistribution rights from the +model license. + +Run `python scripts/check_model_licenses.py` before publishing. CI runs this +without loading checkpoints or downloading LFS payloads, and the release +projection checks the actual outgoing tree. + +## Explicit exceptions + +Never silently relabel third-party weights. A maintainer-reviewed exception +must specify the alternative license and its rationale in both the card and +`legal/model-license-exceptions.json`. The JSON maps the model directory +(relative to the repository) to `license`, `license_file`, and `reason`. +The card uses `Model license: ` and +`License exception: ` in its License section. Bundle the full alternative +license under the specified filename and a nonempty `NOTICE.txt`, linking both +from that section. No current public model has an exception. diff --git a/legal/NVIDIA-Open-Model-License.pdf b/legal/NVIDIA-Open-Model-License.pdf new file mode 100644 index 000000000..d803837f7 Binary files /dev/null and b/legal/NVIDIA-Open-Model-License.pdf differ diff --git a/legal/model-license-exceptions.json b/legal/model-license-exceptions.json new file mode 100644 index 000000000..0967ef424 --- /dev/null +++ b/legal/model-license-exceptions.json @@ -0,0 +1 @@ +{} diff --git a/protomotions/agents/base_agent/agent.py b/protomotions/agents/base_agent/agent.py index 6a385f347..f067dd442 100644 --- a/protomotions/agents/base_agent/agent.py +++ b/protomotions/agents/base_agent/agent.py @@ -702,15 +702,23 @@ def fit(self): self.fit_start_time = time.time() self.fabric.call("on_fit_start", self) - # Marks the transition from setup (env/scene build, which is what fails or - # wedges before training) into the training loop. train_slurm.py's opt-in - # startup watchdog disarms on the first "Epoch " line: printing one here, - # at loop entry, means it disarms right after the scene build rather than - # only at the end of the first (potentially long, silent) data-collection - # epoch -- the rich progress line below flushes to a non-TTY log just once, - # at 100%, which would leave the watchdog armed through a healthy slow - # rollout. flush=True so it reaches the log before that silent window. - print(f"Epoch {self.current_epoch}: entering training loop", flush=True) + # Guarantee last.ckpt exists once training is reached: a resume point for a + # long first epoch, and the signal the cluster retry harness reads + # (scripts/cluster/retry_lib.sh) to tell "reached training" from a flaky + # scene build. Write it ONLY when absent. On a resume the loaded checkpoint + # already is the resume point, and re-writing here every attempt would + # advance last.ckpt's mtime with no real progress -- which the retry loop + # would misread as "training advanced", resetting its no-progress and + # exception-repeat counters and so masking a deterministic first-epoch crash + # instead of fast-failing it. Broadcast rank 0's view so every rank makes the + # SAME call: save() runs distributed collectives, and a split decision would + # deadlock. Safe to save here -- the forward pass during buffer registration + # above already materialized every LazyLinear parameter. + need_initial_ckpt = self.fabric.broadcast( + not (self.root_dir / "last.ckpt").exists() + ) + if need_initial_ckpt: + self.save(checkpoint_name="last.ckpt") while self.current_epoch < self.max_epochs: self.epoch_start_time = time.time() diff --git a/protomotions/agents/base_agent/model.py b/protomotions/agents/base_agent/model.py index f9567270b..f55e6bb54 100644 --- a/protomotions/agents/base_agent/model.py +++ b/protomotions/agents/base_agent/model.py @@ -167,6 +167,32 @@ def reset_rollout_state(self, env_ids=None) -> None: for child in self._proto_children(): child.reset_rollout_state(env_ids=env_ids) + def save_rollout_context(self) -> dict[str, torch.Tensor]: + """Copy every rollout-context buffer in this module tree. + + A caller that runs its own episodes on the training envs, such as + evaluation, saves this first and passes it to ``load_rollout_context`` + afterwards, so training resumes with the context it had. + """ + return { + self._rollout_context_name(name, key): getattr(module, key).clone() + for name, module in self.named_modules() + if isinstance(module, ProtoMotionsTensorDictModule) + for key in module.rollout_state_specs() + } + + def load_rollout_context(self, state: dict[str, torch.Tensor]) -> None: + """Write back rollout context returned by ``save_rollout_context``.""" + for name, module in self.named_modules(): + if not isinstance(module, ProtoMotionsTensorDictModule): + continue + for key in module.rollout_state_specs(): + getattr(module, key).copy_(state[self._rollout_context_name(name, key)]) + + @staticmethod + def _rollout_context_name(module_name: str, key: str) -> str: + return f"{module_name}.{key}" if module_name else key + def _reset_own_rollout_state(self, env_ids=None) -> None: """Reseed selected rows of this module's initialized rollout context.""" for key, spec in self.rollout_state_specs().items(): diff --git a/protomotions/agents/common/common.py b/protomotions/agents/common/common.py index 3c16af5b6..c5f2dccc9 100644 --- a/protomotions/agents/common/common.py +++ b/protomotions/agents/common/common.py @@ -249,6 +249,18 @@ def __init__(self, config: ModuleContainerConfig): self.models = nn.ModuleList() + # Cheap structural check first, so an ordering mistake is reported before + # any module is constructed. Building modules materializes lazy layers and, + # on a multi-rank job, an exception here kills one rank and leaves the rest + # blocked at the next collective -- which presents as a hang rather than as + # the error it is. See ModuleContainerConfig.find_key_flow_errors. + ordering_errors = self.config.find_key_flow_errors() + if ordering_errors: + raise AssertionError( + "ModuleContainer: module list is out of order.\n " + + "\n ".join(ordering_errors) + ) + # Build models and validate data flow # Container in_keys are "available" at the start (like sources) available_keys = set(self.config.in_keys) diff --git a/protomotions/agents/common/config.py b/protomotions/agents/common/config.py index 5209149ae..ef9ca2962 100644 --- a/protomotions/agents/common/config.py +++ b/protomotions/agents/common/config.py @@ -207,6 +207,12 @@ class ModuleContainerConfig: Modules are processed in order, with each module's outputs available to subsequent modules. Input keys are passed through, and all specified output keys must be produced by internal modules. + + **``models`` is execution order, not a dependency graph.** Nothing sorts this + list. A module that produces a key must sit before the module that consumes it; + appending a producer to the end of the list does not work, however clearly its + output names say what it feeds. ``find_key_flow_errors`` catches that specific + mistake without building anything -- see its docstring for why that matters. """ models: List[Any] = field( @@ -223,6 +229,57 @@ class ModuleContainerConfig: metadata={"help": "Output tensor keys produced by this container."}, ) + def find_key_flow_errors(self) -> List[str]: + """Return descriptions of producer-after-consumer ordering mistakes. + + A module needs a key, no earlier module declares it as an output, and a + LATER module does. That is unambiguous: the producer exists, it is simply + in the wrong place, and no amount of runtime context can fix it. + + Deliberately narrow. It does NOT report keys that are never produced at + all, because a few modules derive their real keys at construction time + rather than declaring them in config (``DiscreteLatent`` takes its + ``in_keys`` from its encoder, for instance), so "not declared anywhere" + does not imply "missing". Restricting the check to keys with a known + later producer makes false positives impossible, which is what lets it run + over every experiment config in the repo. + + The value is that this needs no torch, no lazy-layer materialization, no + GPU and no environment. ``ModuleContainer.__init__`` performs the + authoritative check, but only while building the modules -- so a config + that is wrong in this way otherwise gets discovered at model-construction + time, which in a distributed job means one rank dying and the rest + deadlocking at the next collective. + + Returns: + One human-readable string per mistake found; empty when the ordering + is sound. + """ + produced_at: Dict[str, int] = {} + for index, model in enumerate(self.models): + for key in getattr(model, "out_keys", None) or []: + produced_at.setdefault(key, index) + + errors: List[str] = [] + available = set(self.in_keys) + for index, model in enumerate(self.models): + name = type(model).__name__ + for key in getattr(model, "in_keys", None) or []: + if key in available: + continue + producer = produced_at.get(key) + if producer is not None and producer > index: + producer_name = type(self.models[producer]).__name__ + errors.append( + f"model {index} ({name}) needs input key {key!r}, but the " + f"module that produces it is model {producer} " + f"({producer_name}), which runs later. `models` is " + f"execution order: move the producer before the consumer." + ) + for key in getattr(model, "out_keys", None) or []: + available.add(key) + return errors + @dataclass class PretrainedModelConfig: diff --git a/protomotions/agents/evaluators/config.py b/protomotions/agents/evaluators/config.py index 15e584930..57be066f3 100644 --- a/protomotions/agents/evaluators/config.py +++ b/protomotions/agents/evaluators/config.py @@ -59,17 +59,3 @@ class MimicEvaluatorConfig(EvaluatorConfig): default_factory=MotionWeightsRulesConfig, metadata={"help": "Rules for updating motion sampling weights."} ) - eval_action_ema_alpha: Optional[float] = field( - default=None, - metadata={ - "help": ( - "EMA smoothing factor for actions during evaluation only. " - "Simulates deployment low-pass filtering. " - "a_applied = alpha * a_policy + (1-alpha) * a_prev. " - "None = disabled (raw actions). Typical values: 0.5-0.8." - "Smaller alpha = more smoothing." - ), - "min": 0.0, - "max": 1.0, - } - ) diff --git a/protomotions/agents/evaluators/mimic_evaluator.py b/protomotions/agents/evaluators/mimic_evaluator.py index 793e0e8dd..56748b792 100644 --- a/protomotions/agents/evaluators/mimic_evaluator.py +++ b/protomotions/agents/evaluators/mimic_evaluator.py @@ -70,15 +70,22 @@ def _create_metrics( return metrics def initialize_eval(self) -> Dict: - """Initialize evaluation tracking and cache env state for restoration.""" + """Initialize evaluation tracking and cache the state eval will overwrite. + + Saves the env state, the model's rollout context, and the motion manager's + clip assignment, all restored by ``cleanup_after_evaluation``. + """ num_motions = self.motion_lib.num_motions() motion_lengths = self.motion_lib.get_motion_length(None) motion_num_frames = (motion_lengths / self.env.dt).floor().long() motion_num_frames = motion_num_frames.clamp(max=self.config.max_eval_steps) self._init_eval_component_buffers(num_motions) - # Cache env + motion manager state (restored in cleanup_after_evaluation) + # Cache env, model rollout context, and motion manager state (restored in + # cleanup_after_evaluation). Evaluation reseeds the model context of every + # env it resets, so training would otherwise resume with eval's context. self._env_snapshot = self.env.save_state() + self._model_context_snapshot = self.agent.model.save_rollout_context() self._cached_motion_ids = self.motion_manager.motion_ids.clone() self._cached_motion_times = self.motion_manager.motion_times.clone() @@ -143,70 +150,51 @@ def _resolved_min_motion_weight(self, num_motions: int) -> float: return float(min_weight) return float(min_weight) - def _park_inactive_envs(self, active_env_ids: Tensor) -> None: - """Move envs not in ``active_env_ids`` far below the terrain. - - With scene-paired motions, ``_build_eval_batches`` returns a single - batch whose ``env_ids`` (from ``get_unique_fixed_motions``) can be - much smaller than ``num_envs`` -- the remaining envs would otherwise - keep running physics with their (replicated) scenes, contributing - substantially to the PhysX broadphase pair budget and triggering - ``foundLostPairsCapacity`` overflow at large num_envs (silent contact - drops -> tunneling -> phantom failures). - - Parking those envs at z << 0 removes their AABBs from the broadphase - active region without changing the policy's view of the rollout. - """ - if active_env_ids is None or active_env_ids.numel() >= self.num_envs: - return - all_mask = torch.ones(self.num_envs, dtype=torch.bool, device=self.device) - all_mask[active_env_ids] = False - inactive_env_ids = torch.nonzero(all_mask, as_tuple=False).flatten() - if inactive_env_ids.numel() == 0: - return - self.env.simulator.park_envs(inactive_env_ids) - def evaluate_episode(self, env_ids: torch.Tensor, max_steps: int) -> None: - """Run a single episode batch, optionally with EMA action smoothing. - - When eval_action_ema_alpha is set, actions are low-pass filtered to - simulate deployment conditions. Motions that fail under EMA get higher - sampling weight, creating curriculum pressure toward smooth policies. + """Run a single episode batch with the policy's own actions. + + Evaluation is an ordinary rollout in which only ``env_ids`` are scored. Every + other env plays its background clip (see ``_background_motion_ids``) and, like + a training env, restarts it whenever it falls or the clip ends. No env is + parked or left idle: all envs share one physics scene, so the unscored ones + should load it the way training does. Parking them below the floor instead + cost ~12.8 points of SAMP eval success at 8192 envs (75.96% vs 88.80%). + Scoring stays on the deduped ``env_ids`` subset -- widening it would corrupt + the motion-id-indexed buffers (duplicate indices drop writes). + + Every env reset here, scored or not, also gets fresh model rollout context + (for example MaskedMimic's per-episode latent noise), as training does after + an env finishes. ``env.reset`` already clears the env's own history. """ - ema_alpha = self.config.eval_action_ema_alpha - self._on_episode_start(env_ids) - - # Park envs that aren't part of this batch so they don't generate - # PhysX broadphase pairs / contacts during the eval. Pre-eval state is - # restored later in cleanup_after_evaluation via env.restore_state(). - self._park_inactive_envs(env_ids) - - obs, _ = self.env.reset(env_ids, **self._get_reset_kwargs()) + all_env_ids = torch.arange(self.num_envs, device=self.device) + obs, _ = self.env.reset(all_env_ids, **self._get_reset_kwargs()) + self.agent.model.reset_rollout_context(env_ids=all_env_ids) + unscored = torch.ones(self.num_envs, dtype=torch.bool, device=self.device) + unscored[env_ids] = False self.agent.pre_collect_step(0) obs = self.agent.add_agent_info_to_obs(obs) obs_td = self.agent.obs_dict_to_tensordict(obs) - prev_actions = None - for step_idx in range(max_steps): actions = self._policy_action(obs_td) + obs, rewards, dones, terminated, extras = self.env.step(actions) - # Apply EMA smoothing (deployment simulation) - if ema_alpha is not None: - if prev_actions is None: - prev_actions = actions.clone() - actions = ema_alpha * actions + (1.0 - ema_alpha) * prev_actions - prev_actions = actions.clone() + # Score before restarting anything: a reset discards the cached env + # context that the failure checks read. + self._check_eval_components(env_ids, step_idx) + self._on_episode_step(env_ids, extras, actions) + + restart_ids = (dones.bool() & unscored).nonzero(as_tuple=False).flatten() + if restart_ids.numel() > 0: + self.motion_manager.motion_times[restart_ids] = 0.0 + obs, _ = self.env.reset(restart_ids, **self._get_reset_kwargs()) + self.agent.model.reset_rollout_context(env_ids=restart_ids) - obs, rewards, dones, terminated, extras = self.env.step(actions) self.agent.pre_collect_step(step_idx + 1) obs = self.agent.add_agent_info_to_obs(obs) obs_td = self.agent.obs_dict_to_tensordict(obs) - self._check_eval_components(env_ids, step_idx) - self._on_episode_step(env_ids, extras, actions) - def run_evaluation(self) -> None: """Run evaluation across multiple motions.""" for env_ids, motion_ids in self._build_eval_batches(): @@ -251,10 +239,30 @@ def _build_eval_batches(self): # --- Hook overrides --- + def _background_motion_ids(self) -> Tensor: + """The clip each env plays while it is not being scored. + + An env with a fixed (scene-paired) clip plays that clip, the only one its + scene objects match. Any other env plays library clip ``env_id % num_motions``: + a fixed choice rather than a random draw, so evaluation stays reproducible, + and one that spreads the unscored envs over different clips. + """ + env_ids = torch.arange(self.num_envs, device=self.device) + motion_ids = env_ids % self.motion_lib.num_motions() + fixed = getattr(self.motion_manager, "_fixed_motion_ids_per_env", None) + if fixed is not None: + motion_ids = torch.where(fixed != -1, fixed, motion_ids) + return motion_ids + def _on_episode_start(self, env_ids: Tensor) -> None: - """Set motion_ids/times in the motion manager before reset.""" + """Set motion_ids/times in the motion manager before reset. + + Every env starts its background clip from frame 0, then the scored + subset is overlaid with this batch's motions. + """ + self.motion_manager.motion_ids[:] = self._background_motion_ids() + self.motion_manager.motion_times[:] = 0.0 self.motion_manager.motion_ids[env_ids] = self._episode_ctx.motion_ids - self.motion_manager.motion_times[env_ids] = 0.0 def _get_reset_kwargs(self) -> dict: """Customize env.reset() for mimic evaluation.""" @@ -316,12 +324,19 @@ def process_eval_results(self) -> Tuple[Dict, Optional[float], int]: return to_log, success_rate, num_eval_items def cleanup_after_evaluation(self) -> None: - """Restore env and motion manager state after evaluation.""" + """Restore the env, model rollout context, and motion manager state. + + ``env.restore_state`` first clears the simulator's contact memory (see + ``Simulator.flush_contact_cache``), then writes every env's saved state + back, so training continues the episodes it was running before eval. + """ self.motion_manager.motion_ids = self._cached_motion_ids self.motion_manager.motion_times = self._cached_motion_times self.env.restore_state(self._env_snapshot) + self.agent.model.load_rollout_context(self._model_context_snapshot) del self._env_snapshot + del self._model_context_snapshot del self._cached_motion_ids del self._cached_motion_times super().cleanup_after_evaluation() diff --git a/protomotions/envs/action/action_functions.py b/protomotions/envs/action/action_functions.py index 2eae2cef5..5c1600e77 100644 --- a/protomotions/envs/action/action_functions.py +++ b/protomotions/envs/action/action_functions.py @@ -359,9 +359,9 @@ def bm_pd_action( def make_bm_pd_action_config(robot_config) -> Dict[str, Any]: """Create action config dict for BeyondMimic-style PD action scaling. - Computes per-joint action scale as ``effort_limit / stiffness``, - which normalizes the policy output so ~O(1) values map to ~ the - maximum achievable deflection per joint. Uses ``robot_config.default_dof_pos`` + Computes per-joint action scale as ``0.25 * effort_limit / stiffness``, + so unit policy output produces a proportional torque of 25% of the effort + limit at the default pose. Uses ``robot_config.default_dof_pos`` as the action offset (default standing pose). Args: @@ -391,7 +391,7 @@ def make_bm_pd_action_config(robot_config) -> Dict[str, Any]: if torch.any(stiffness <= 0): raise ValueError("BM PD action stiffness must be positive for all joints") - action_scale = effort_limit / stiffness + action_scale = 0.25 * effort_limit / stiffness pd_action_offset = robot_config.default_dof_pos.clone() diff --git a/protomotions/envs/base_env/env.py b/protomotions/envs/base_env/env.py index 5f8d143c1..ac57439b0 100644 --- a/protomotions/envs/base_env/env.py +++ b/protomotions/envs/base_env/env.py @@ -984,6 +984,7 @@ def _build_context(self, env_ids: Optional[Tensor] = None) -> EnvContext: if self.state_history else None, # Actions (historical) + current_action=self._select_context_tensor(self._current_raw_action, env_ids), current_processed_action=self._select_context_tensor( self._current_processed_action, env_ids ), @@ -1901,6 +1902,13 @@ def restore_state(self, snapshot: dict) -> None: snapshot: Dictionary from save_state() containing state tensors """ env_ids = torch.arange(self.num_envs, device=self.device) + # Every env is about to be teleported back. Clear the simulator's + # memory of the contacts at the current poses first: otherwise a robot + # that was standing on the ground here can be launched on the first + # physics substep after the write-back. + flush_contact_cache = getattr(self.simulator, "flush_contact_cache", None) + if flush_contact_cache is not None: + flush_contact_cache() self.simulator.reset_envs( snapshot["robot_state"], snapshot.get("object_state"), env_ids ) diff --git a/protomotions/envs/base_env/utils.py b/protomotions/envs/base_env/utils.py index 71e63aeb1..a16b004b9 100644 --- a/protomotions/envs/base_env/utils.py +++ b/protomotions/envs/base_env/utils.py @@ -166,7 +166,8 @@ def combine_evaluation( Evaluation components return numeric values [num_envs]. If a component has a ``threshold`` in its static_params, the value is compared against it to determine failure. ``fail_above`` (default True) controls the comparison - direction. + direction. Evaluation kernels must not declare a ``threshold`` parameter, + otherwise MdpComponent would pass this failure cut to the kernel as well. Args: raw_values: Dict of {name: value_tensor} from ComponentManager.execute_all(). diff --git a/protomotions/envs/component_factories.py b/protomotions/envs/component_factories.py index 009500554..1e6155acf 100644 --- a/protomotions/envs/component_factories.py +++ b/protomotions/envs/component_factories.py @@ -758,6 +758,20 @@ def action_smoothness_factory(weight: float = -0.02) -> MdpComponent: ) +def action_rate_l2_factory(weight: float = -0.1) -> MdpComponent: + """Penalize squared raw action changes; requires at least one history step.""" + from protomotions.envs.rewards import compute_action_rate_l2 + + return MdpComponent( + compute_func=compute_action_rate_l2, + dynamic_vars={ + "raw_action": EnvContext.current_action, + "previous_raw_action": EnvContext.previous_action, + }, + static_params={"weight": weight}, + ) + + def gt_rew_factory(weight: float = 0.5, coefficient: float = -100.0) -> MdpComponent: """Factory for position tracking reward. @@ -1813,6 +1827,7 @@ def steering_velocity_error_factory( "path_obs_factory", # Reward factories "action_smoothness_factory", + "action_rate_l2_factory", "gt_rew_factory", "gr_rew_factory", "gv_rew_factory", diff --git a/protomotions/envs/context_views.py b/protomotions/envs/context_views.py index 58d2e93bc..de96f720d 100644 --- a/protomotions/envs/context_views.py +++ b/protomotions/envs/context_views.py @@ -144,6 +144,16 @@ def __init__(self, state: "RobotState", anchor_idx: int): self.anchor_vel = state.rigid_body_vel[:, anchor_idx, :] self.anchor_ang_vel = state.rigid_body_ang_vel[:, anchor_idx, :] + # NoisyObservations also carries root/anchor channels that were noised + # on their own (anchor_rot_noise, root_ang_vel_noise, + # anchor_ang_vel_noise). Re-deriving them from the whole-body tensors + # would silently drop that noise, so use them when present. RobotState + # has none of these fields, so clean views are unchanged. + for name in ("anchor_rot", "root_local_ang_vel", "anchor_local_ang_vel"): + value = getattr(state, name, None) + if value is not None: + setattr(self, name, value) + def _compute_root_local_ang_vel(self) -> Tensor: return compute_local_ang_vel(self.root_rot, self.root_ang_vel) @@ -591,6 +601,7 @@ class EnvContext: noisy_historical: Optional[HistoricalView] = NestedField(HistoricalView) # Actions (historical) + current_action: Optional[Tensor] = FieldPath() previous_action: Optional[Tensor] = FieldPath() current_processed_action: Optional[Tensor] = FieldPath() previous_processed_action: Optional[Tensor] = FieldPath() @@ -645,6 +656,7 @@ def __init__( dt: float, historical: Optional[HistoricalView] = None, noisy_historical: Optional[HistoricalView] = None, + current_action: Optional[Tensor] = None, previous_action: Optional[Tensor] = None, current_processed_action: Optional[Tensor] = None, previous_processed_action: Optional[Tensor] = None, @@ -679,6 +691,7 @@ def __init__( dt: Simulation timestep in seconds. historical: Clean historical state view (optional). noisy_historical: Noisy historical state view (optional). + current_action: Current raw policy action before PD processing (optional). previous_action: Previous raw action [num_envs, action_dim] (optional). current_processed_action: Current processed action (optional). previous_processed_action: Previous processed action (optional). @@ -711,6 +724,7 @@ def __init__( self.noisy_historical = noisy_historical # Actions + self.current_action = current_action self.previous_action = previous_action self.current_processed_action = current_processed_action self.previous_processed_action = previous_processed_action diff --git a/protomotions/envs/mdp_component.py b/protomotions/envs/mdp_component.py index cfd909a1f..09a2a6630 100644 --- a/protomotions/envs/mdp_component.py +++ b/protomotions/envs/mdp_component.py @@ -38,12 +38,20 @@ Metadata keys ------------- -Some ``static_params`` keys are *metadata* consumed by the reward/termination -combining logic, not by the compute_func itself. These are automatically filtered -out before calling the compute_func. Current metadata keys: +Some ``static_params`` keys are *metadata* read by the combining logic in +``base_env/utils.py`` (reward weighting and clamping, the evaluation failure +cut). Current metadata keys: weight, multiplicative, zero_during_grace_period, - min_value, max_value + min_value, max_value, threshold, fail_above, use_region_weights + +A metadata key is passed to the compute_func only when the compute_func +declares a parameter with that name; otherwise it is filtered out. Termination +kernels declare ``threshold``, so a termination's configured threshold reaches +its kernel. Evaluation kernels must not declare ``threshold``: for them it is +the failure cut that ``combine_evaluation`` applies to the kernel's output. A +compute_func that only takes ``**kwargs``, or whose signature cannot be +inspected (a TorchScript function, for example), receives no metadata keys. Example ------- @@ -80,7 +88,8 @@ from __future__ import annotations -from typing import Any, Callable, Dict, TYPE_CHECKING +import inspect +from typing import Any, Callable, Dict, FrozenSet, TYPE_CHECKING from torch import Tensor @@ -90,9 +99,9 @@ from protomotions.envs.context_views import EnvContext -# Keys in static_params that are reward/termination metadata, NOT compute_func arguments. -# These are consumed by combine_rewards / combine_terminations and must not -# be forwarded to the compute_func function. +# Keys in static_params that are read by combine_rewards / combine_evaluation. +# resolve_args drops them before calling the compute_func, unless the +# compute_func declares a parameter of the same name. _METADATA_KEYS = frozenset({ "weight", "multiplicative", "zero_during_grace_period", "min_value", "max_value", @@ -100,6 +109,27 @@ }) +def _declared_keyword_params(func: Callable[..., Any]) -> FrozenSet[str]: + """Return the names of the parameters ``func`` declares and accepts by keyword. + + ``*args`` and ``**kwargs`` do not count, so a wrapper that forwards + ``**kwargs`` declares nothing. A function whose signature cannot be + inspected declares nothing either. That includes TorchScript functions, so a + ``@torch.jit.script`` kernel never receives metadata keys: if one needs a + threshold, give its parameter a name that is not in ``_METADATA_KEYS``. + """ + try: + parameters = inspect.signature(func).parameters.values() + except (TypeError, ValueError): + return frozenset() + return frozenset( + param.name + for param in parameters + if param.kind + in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY) + ) + + class MdpComponent: """A compute function with explicit context bindings. @@ -157,15 +187,18 @@ def __init__( self.static_params = static_params or {} self.compile = compile self._device_ready = False + self._metadata_args = None def __getstate__(self): state = self.__dict__.copy() state.pop("_device_ready", None) + state.pop("_metadata_args", None) return state def __setstate__(self, state): self.__dict__.update(state) self._device_ready = False + self._metadata_args = None def resolve_args(self, ctx: "EnvContext") -> tuple: @@ -188,9 +221,15 @@ def resolve_args(self, ctx: "EnvContext") -> tuple: if not self._device_ready: self._ensure_device(resolved) + if self._metadata_args is None: + # Metadata keys the compute_func asks for by name, such as a + # termination kernel's threshold. + self._metadata_args = _METADATA_KEYS & _declared_keyword_params( + self.compute_func + ) func_params = { k: v for k, v in self.static_params.items() - if k not in _METADATA_KEYS + if k not in _METADATA_KEYS or k in self._metadata_args } return resolved, func_params diff --git a/protomotions/envs/rewards/__init__.py b/protomotions/envs/rewards/__init__.py index ce9507093..2b69c8f96 100644 --- a/protomotions/envs/rewards/__init__.py +++ b/protomotions/envs/rewards/__init__.py @@ -63,6 +63,7 @@ # Regularization reward kernels from protomotions.envs.rewards.regularization import ( compute_action_smoothness, + compute_action_rate_l2, compute_action_smoothness_logmeanexp, compute_pow_rew, compute_soft_pos_limit_rew, @@ -113,6 +114,7 @@ "compute_target_rew", # Regularization reward kernels "compute_action_smoothness", + "compute_action_rate_l2", "compute_action_smoothness_logmeanexp", "compute_pow_rew", "compute_soft_pos_limit_rew", diff --git a/protomotions/envs/rewards/regularization.py b/protomotions/envs/rewards/regularization.py index d7b1de659..704bdefb7 100644 --- a/protomotions/envs/rewards/regularization.py +++ b/protomotions/envs/rewards/regularization.py @@ -57,6 +57,15 @@ def compute_action_smoothness( return delta_norm(current_processed_action, previous_processed_action) +def compute_action_rate_l2(raw_action: Tensor, previous_raw_action: Tensor) -> Tensor: + """Sum squared raw policy-action changes over joints, per environment. + + Inputs are dimensionless policy outputs [num_envs, action_dim], before + clipping or PD scaling. No mean, square root, or timestep scaling is applied. + """ + return (raw_action - previous_raw_action).square().sum(dim=-1) + + def compute_action_smoothness_logmeanexp( current_processed_action: Tensor, previous_processed_action: Tensor, diff --git a/protomotions/inference_agent.py b/protomotions/inference_agent.py index aa05b659f..f58855fae 100644 --- a/protomotions/inference_agent.py +++ b/protomotions/inference_agent.py @@ -267,10 +267,12 @@ def main(): -3 ] # e.g., "isaacgym" from "protomotions.simulator.isaacgym.simulator.IsaacGymSimulator" + train_sim_params = None if args.simulator != current_simulator: log.info( f"Switching simulator from '{current_simulator}' (training) to '{args.simulator}' (inference)" ) + train_sim_params = simulator_config.sim from protomotions.simulator.factory import update_simulator_config_for_test simulator_config = update_simulator_config_for_test( @@ -341,6 +343,13 @@ def main(): log.info(f"CLI override: command_source = {args.command_source}") apply_command_source_overrides(env_config, args.command_source) + if train_sim_params is not None: + from protomotions.simulator.factory import warn_on_control_rate_change + + warn_on_control_rate_change( + train_sim_params, simulator_config.sim, args.simulator + ) + motion_lib_config.validate() # Create fabric config for inference (simplified) diff --git a/protomotions/robot_configs/amp.py b/protomotions/robot_configs/amp.py index 215d8fdfd..8ca9ed939 100644 --- a/protomotions/robot_configs/amp.py +++ b/protomotions/robot_configs/amp.py @@ -12,6 +12,7 @@ from protomotions.simulator.isaaclab.config import IsaacLabSimParams from protomotions.simulator.genesis.config import GenesisSimParams from protomotions.simulator.newton.config import NewtonSimParams +from protomotions.simulator.mujoco.config import MujocoSimParams from typing import List, Dict, Tuple from dataclasses import dataclass, field @@ -74,5 +75,13 @@ class AMPRobotConfig(RobotConfig): fps=120, decimation=4, ), + # Same 30 Hz control rate as the other simulators. implicitfast + # also treats the PD damping implicitly (Euler only does that for + # joint damping), which keeps stiff gains on light links stable. + mujoco=MujocoSimParams( + fps=960, + decimation=32, + integrator="implicitfast", + ), ) ) diff --git a/protomotions/robot_configs/smpl.py b/protomotions/robot_configs/smpl.py index 6b652a142..43ae804bc 100644 --- a/protomotions/robot_configs/smpl.py +++ b/protomotions/robot_configs/smpl.py @@ -15,6 +15,7 @@ ) from protomotions.simulator.genesis.config import GenesisSimParams from protomotions.simulator.newton.config import NewtonSimParams +from protomotions.simulator.mujoco.config import MujocoSimParams from protomotions.components.pose_lib import ControlInfo from typing import List, Dict, Tuple from dataclasses import dataclass, field @@ -123,5 +124,13 @@ class SmplRobotConfig(RobotConfig): fps=120, decimation=4, ), + # Same 30 Hz control rate as the other simulators. implicitfast + # also treats the PD damping implicitly (Euler only does that for + # joint damping), which keeps stiff gains on light links stable. + mujoco=MujocoSimParams( + fps=960, + decimation=32, + integrator="implicitfast", + ), ) ) diff --git a/protomotions/robot_configs/smplx.py b/protomotions/robot_configs/smplx.py index 81d9124d5..8ebf74c6c 100644 --- a/protomotions/robot_configs/smplx.py +++ b/protomotions/robot_configs/smplx.py @@ -15,6 +15,7 @@ ) from protomotions.simulator.genesis.config import GenesisSimParams from protomotions.simulator.newton.config import NewtonSimParams +from protomotions.simulator.mujoco.config import MujocoSimParams from protomotions.components.pose_lib import ControlInfo from typing import List, Dict, Tuple from dataclasses import dataclass, field @@ -208,5 +209,13 @@ class SMPLXRobotConfig(RobotConfig): fps=120, decimation=4, ), + # Same 30 Hz control rate as the other simulators. implicitfast + # also treats the PD damping implicitly (Euler only does that for + # joint damping), which keeps stiff gains on light links stable. + mujoco=MujocoSimParams( + fps=960, + decimation=32, + integrator="implicitfast", + ), ) ) diff --git a/protomotions/robot_configs/soma23.py b/protomotions/robot_configs/soma23.py index fa042c258..b225ab3c2 100644 --- a/protomotions/robot_configs/soma23.py +++ b/protomotions/robot_configs/soma23.py @@ -15,6 +15,7 @@ ) from protomotions.simulator.genesis.config import GenesisSimParams from protomotions.simulator.newton.config import NewtonSimParams +from protomotions.simulator.mujoco.config import MujocoSimParams from protomotions.components.pose_lib import ControlInfo from typing import List, Dict, Tuple from dataclasses import dataclass, field @@ -118,5 +119,13 @@ class Soma23RobotConfig(RobotConfig): fps=120, decimation=4, ), + # Same 30 Hz control rate as the other simulators. implicitfast + # also treats the PD damping implicitly (Euler only does that for + # joint damping), which keeps stiff gains on light links stable. + mujoco=MujocoSimParams( + fps=960, + decimation=32, + integrator="implicitfast", + ), ) ) diff --git a/protomotions/simulator/base_simulator/robot_mass_properties.py b/protomotions/simulator/base_simulator/robot_mass_properties.py new file mode 100644 index 000000000..4b29a6bc2 --- /dev/null +++ b/protomotions/simulator/base_simulator/robot_mass_properties.py @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Normalize mass and inertia of massless robot links before backend import.""" + +from copy import deepcopy +from pathlib import Path +import xml.etree.ElementTree as ET + +from dm_control import mjcf +import mujoco + + +MARKER_MASS = 0.001 # kg +MARKER_INERTIA = 4e-8 # kg m^2: a one-gram, one-centimetre-radius sphere + + +def _read_mjcf(path): + """Expand includes and make external asset references location independent.""" + path = Path(path).resolve() + root = ET.parse(path).getroot() + + def expand(parent, ancestors): + for child in list(parent): + if child.tag != "include": + expand(child, ancestors) + continue + included_path = (path.parent / child.attrib["file"]).resolve() + if included_path in ancestors: + raise ValueError(f"Recursive MJCF include: {included_path}") + included = ET.parse(included_path).getroot() + expand(included, ancestors | {included_path}) + index = list(parent).index(child) + parent.remove(child) + for offset, element in enumerate(included): + parent.insert(index + offset, element) + + expand(root, {path}) + compiler = root.find("compiler") + options = compiler.attrib if compiler is not None else {} + for element in root.findall("asset/*"): + for attribute in ( + "file", + "fileleft", + "fileright", + "filefront", + "fileback", + "fileup", + "filedown", + ): + filename = element.get(attribute) + if not filename: + continue + directory_key = {"mesh": "meshdir", "texture": "texturedir"}.get( + element.tag, "assetdir" + ) + directory = options.get(directory_key, options.get("assetdir", "")) + if options.get("strippath") == "true": + filename = Path(filename).name + element.set(attribute, str((path.parent / directory / filename).resolve())) + if compiler is not None: + for attribute in ("meshdir", "texturedir", "assetdir", "strippath"): + compiler.attrib.pop(attribute, None) + return root + + +def prepare_robot_mjcf(mjcf_path, output_dir): + """Return an importable MJCF with one-gram inertias for zero-mass bodies. + + The source is never edited. MuJoCo resolves defaults, geometry-derived mass + and includes, rather than guessing from missing XML attributes. Only robot + inertia is inspected; references to runtime sensors/ground are excluded from + the probe, not from the resulting asset. Assets with no massless links use + their original path. The caller owns the temporary output directory. + """ + root = _read_mjcf(mjcf_path) + probe = deepcopy(root) + for tag in ("sensor", "contact", "actuator", "tendon", "equality", "keyframe"): + for element in probe.findall(tag): + probe.remove(element) + compiler = probe.find("compiler") + if compiler is None: + compiler = ET.SubElement(probe, "compiler") + compiler.set("fusestatic", "false") + source = mjcf.from_xml_string(ET.tostring(probe, encoding="unicode")) + model = mujoco.MjModel.from_xml_string(source.to_xml_string(), source.get_assets()) + massless = {i for i in range(1, model.nbody) if model.body_mass[i] == 0} + if not massless: + return str(mjcf_path) + + compiler = root.find("compiler") + infer_from_geometry = ( + compiler is not None and compiler.get("inertiafromgeom") == "true" + ) + # With static fusion disabled, XML preorder is MuJoCo body order, including + # unnamed bodies. Body names and ordering are never rewritten. + for body_id, body in enumerate(root.findall(".//worldbody//body"), start=1): + if body_id not in massless and not infer_from_geometry: + continue + properties = model.body(body_id) + inertial = body.find("inertial") + if inertial is not None: + body.remove(inertial) + if body_id in massless: + attributes = dict( + pos="0 0 0", + mass=str(MARKER_MASS), + diaginertia=" ".join([str(MARKER_INERTIA)] * 3), + ) + else: + # Preserve geometry-derived physical properties when switching the + # compiler from forced inference to explicit/automatic inertia. + attributes = dict( + pos=" ".join(map(str, properties.ipos)), + quat=" ".join(map(str, properties.iquat)), + mass=str(properties.mass[0]), + diaginertia=" ".join(map(str, properties.inertia)), + ) + body.insert(0, ET.Element("inertial", attributes)) + if infer_from_geometry: + compiler.set("inertiafromgeom", "auto") + + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + output_path = output_dir / Path(mjcf_path).name + output_path.write_text(ET.tostring(root, encoding="unicode")) + return str(output_path) diff --git a/protomotions/simulator/base_simulator/simulator.py b/protomotions/simulator/base_simulator/simulator.py index 6a852f5d2..7c9f46e01 100644 --- a/protomotions/simulator/base_simulator/simulator.py +++ b/protomotions/simulator/base_simulator/simulator.py @@ -23,6 +23,8 @@ from abc import ABC, abstractmethod import logging import math +from pathlib import Path +from tempfile import TemporaryDirectory from typing import Dict, List, Optional, Any, Tuple, Callable @@ -55,6 +57,8 @@ from protomotions.robot_configs.base import ControlType, RobotConfig from protomotions.simulator.base_simulator.record import RecordingMixin from protomotions.simulator.base_simulator.user_interface import UserInterface +from protomotions.assets import resolve_asset_root +from protomotions.simulator.base_simulator.robot_mass_properties import prepare_robot_mjcf class Simulator(RecordingMixin, ABC): @@ -99,6 +103,23 @@ class Simulator(RecordingMixin, ABC): # ------------------------- # ⚙️ Group 1: Initialization & Configuration # ------------------------- + def _get_robot_asset_path(self) -> str: + """Prepare robot inertia once, retaining temporary assets for its lifetime.""" + if not hasattr(self, "_prepared_robot_asset_path"): + path = ( + resolve_asset_root(self.robot_config.asset.asset_root) + / self.robot_config.asset.asset_file_name + ) + self._robot_asset_directory = TemporaryDirectory( + prefix="protomotions-robot-" + ) + self._prepared_robot_asset_path = ( + prepare_robot_mjcf(path, self._robot_asset_directory.name) + if Path(path).suffix == ".xml" + else str(path) + ) + return self._prepared_robot_asset_path + def __init__( self, config: SimulatorConfig, @@ -792,85 +813,85 @@ def reset_envs( # Reset projectiles for reset environments self._reset_projectiles(env_ids) - def park_envs( + def flush_contact_cache(self) -> None: + """Make the simulator forget every env's current contacts. + + Called right before every env's state is written back + (``BaseEnv.restore_state``). The base implementation does nothing; + backends whose contact solver remembers contacts between steps override + it. See ``IsaacLabSimulator.flush_contact_cache`` for why this matters. + """ + return None + + def _move_all_envs_out_of_contact( self, - env_ids: torch.Tensor, hide_z: float = -50.0, + object_drop: float = 20.0, ) -> None: - """Move robot and scene objects for ``env_ids`` far below the terrain. - - Used during evaluation to disable physics for envs that are not being - evaluated, eliminating their contribution to the PhysX broadphase pair - budget. Parked bodies sit well below any terrain/object AABB so no - broadphase pairs are generated, no narrow-phase contacts are computed, - and no found/lost pair churn occurs. Velocities are zeroed so the - parked bodies stay put. - - Pre-eval state is restored later via ``BaseEnv.restore_state(snapshot)``, - which calls ``reset_envs`` over all envs with the saved snapshot. - - Args: - env_ids: Environment IDs to park. No-op if empty/None. - hide_z: World z-coordinate to teleport robot roots to. Object roots - are placed slightly below at ``hide_z - 1.0`` so their AABBs - cannot overlap with the parked robot's AABB. + """Move every robot and scene object far below the terrain, at rest. + + A helper for ``flush_contact_cache`` overrides, which then simulate one + substep here and write the real state back. Robots keep their x, y and + rotation, so each stays in its own grid cell instead of all landing at + the world origin (which could exceed PhysX's broadphase limits at large + env counts); they go to ``z = hide_z`` in the default joint pose with + zero velocity. Objects keep their x, y and rotation and go to + ``hide_z - object_drop``, far enough below that a robot cannot touch + its own furniture. Every object of an env goes to that same height, so + two objects that overlap when seen from above (a box on a table) end up + inside each other, and if either is dynamic the simulator records a + contact between them. + + This goes through ``reset_envs``, so it also clears the action buffers + and reschedules pushes and projectiles; the write-back that follows + does the same again for every env. """ from protomotions.simulator.base_simulator.simulator_state import ( StateConversion, ) - if env_ids is None or env_ids.numel() == 0: - return - - n = env_ids.numel() + n = self.num_envs device = self.device + env_ids = torch.arange(n, device=device) - # Preserve current root xy so parked envs stay in their own grid cell - # (avoids stacking all parked envs at world origin, which could exceed - # PhysX broadphase region limits at large num_envs). current_root = self.get_root_state(env_ids) - park_root_pos = current_root.root_pos.clone() - park_root_pos[:, 2] = hide_z - park_root_rot = current_root.root_rot.clone() - zero_root_vel = torch.zeros((n, 3), device=device, dtype=torch.float32) - - num_dofs = self.robot_config.kinematic_info.num_dofs - park_dof_pos = ( + root_pos = current_root.root_pos.clone() + root_pos[:, 2] = hide_z + zero_vel = torch.zeros((n, 3), device=device, dtype=torch.float32) + dof_pos = ( self.robot_config.default_dof_pos.unsqueeze(0) .repeat(n, 1) .to(device=device, dtype=torch.float32) ) - park_dof_vel = torch.zeros((n, num_dofs), device=device, dtype=torch.float32) - - park_state = ResetState( - root_pos=park_root_pos, - root_rot=park_root_rot, - root_vel=zero_root_vel, - root_ang_vel=zero_root_vel.clone(), - dof_pos=park_dof_pos, - dof_vel=park_dof_vel, + robot_state = ResetState( + root_pos=root_pos, + root_rot=current_root.root_rot.clone(), + root_vel=zero_vel, + root_ang_vel=zero_vel.clone(), + dof_pos=dof_pos, + dof_vel=torch.zeros_like(dof_pos), state_conversion=StateConversion.COMMON, ) - park_object_state = None + object_state = None if self.scene_lib.num_objects_per_scene > 0: current_obj = self.get_object_root_state(env_ids) - obj_root_pos = current_obj.root_pos.clone() - obj_root_pos[..., 2] = hide_z - 1.0 - obj_root_rot = current_obj.root_rot.clone() - m = self.scene_lib.num_objects_per_scene + obj_pos = current_obj.root_pos.clone() + obj_pos[..., 2] = hide_z - object_drop zero_obj_vel = torch.zeros( - (n, m, 3), device=device, dtype=torch.float32 + (n, self.scene_lib.num_objects_per_scene, 3), + device=device, + dtype=torch.float32, ) - park_object_state = ObjectState( - root_pos=obj_root_pos, - root_rot=obj_root_rot, + object_state = ObjectState( + root_pos=obj_pos, + root_rot=current_obj.root_rot.clone(), root_vel=zero_obj_vel, root_ang_vel=zero_obj_vel.clone(), state_conversion=StateConversion.COMMON, ) - self.reset_envs(park_state, park_object_state, env_ids) + self.reset_envs(robot_state, object_state, env_ids) @abstractmethod def _set_simulator_env_state( diff --git a/protomotions/simulator/factory.py b/protomotions/simulator/factory.py index 1af29e1b4..25c9e8a9d 100644 --- a/protomotions/simulator/factory.py +++ b/protomotions/simulator/factory.py @@ -78,6 +78,27 @@ def simulator_config( ) +def warn_on_control_rate_change( + train_sim_params, sim_params, simulator_name: str +) -> bool: + """Warn when inference runs the policy at a different control rate than training. + + The policy acts once per control step (fps / decimation), so a different + rate puts it in a regime it never trained in. Call this after all config + overrides are applied. Returns True if it warned. + """ + train_rate = train_sim_params.fps / train_sim_params.decimation + rate = sim_params.fps / sim_params.decimation + if abs(train_rate - rate) <= 1e-6 * train_rate: + return False + log.warning( + f"Control rate changes from {train_rate:g} Hz (training) to {rate:g} Hz " + f"({simulator_name}). Match it with " + f"--overrides simulator.sim.fps=... simulator.sim.decimation=..." + ) + return True + + def update_simulator_config_for_test( current_simulator_config: SimulatorConfig, new_simulator: str, diff --git a/protomotions/simulator/genesis/simulator.py b/protomotions/simulator/genesis/simulator.py index d9994fba8..0df68f842 100644 --- a/protomotions/simulator/genesis/simulator.py +++ b/protomotions/simulator/genesis/simulator.py @@ -6,7 +6,6 @@ from typing import Dict, Optional, Tuple from easydict import EasyDict -from protomotions.assets import resolve_asset_root from protomotions.utils import rotations from protomotions.simulator.base_simulator.simulator import Simulator from protomotions.simulator.base_simulator.config import ( @@ -126,11 +125,7 @@ def _create_simulation(self) -> None: # ===== Group 2: Environment Setup & Configuration ===== def _create_envs(self) -> None: """Creates the simulation environments and loads robot assets.""" - asset_root = resolve_asset_root(self.robot_config.asset.asset_root) - asset_file = self.robot_config.asset.asset_file_name - - asset_path = os.path.join(asset_root, asset_file) - asset_root = os.path.dirname(asset_path) + asset_path = self._get_robot_asset_path() asset_file = os.path.basename(asset_path) asset_file_ext = os.path.splitext(asset_file)[1] diff --git a/protomotions/simulator/isaacgym/simulator.py b/protomotions/simulator/isaacgym/simulator.py index 77ee761f0..7ed2fe901 100644 --- a/protomotions/simulator/isaacgym/simulator.py +++ b/protomotions/simulator/isaacgym/simulator.py @@ -12,7 +12,7 @@ import numpy as np from rich.progress import Progress import os -from protomotions.assets import get_asset_root, resolve_asset_root +from protomotions.assets import get_asset_root from protomotions.components.terrains.terrain import Terrain from protomotions.components.terrains.config import CombineMode from protomotions.components.scene_lib import ( @@ -346,9 +346,7 @@ def _load_humanoid_asset(self): Returns: Loaded asset handle (opaque gymapi handle) """ - asset_root = resolve_asset_root(self.robot_config.asset.asset_root) - asset_file = self.robot_config.asset.asset_file_name - asset_path = os.path.join(asset_root, asset_file) + asset_path = self._get_robot_asset_path() asset_root = os.path.dirname(asset_path) asset_file = os.path.basename(asset_path) diff --git a/protomotions/simulator/isaaclab/simulator.py b/protomotions/simulator/isaaclab/simulator.py index d36d37cbc..ec691038b 100644 --- a/protomotions/simulator/isaaclab/simulator.py +++ b/protomotions/simulator/isaaclab/simulator.py @@ -56,6 +56,55 @@ ) +def _force_sequential_cloning() -> None: + """Make IsaacLab assign scene ``i``'s object to env ``i`` instead of drawing at random. + + This gives the exact mapping only when every scene has one object. Each object + slot is spawned as its own group, and IsaacLab 3 enumerates combinations of the + groups' prototypes with the last group varying fastest, so with two objects per + scene env ``i`` gets combination ``(i // N, i % N)`` for ``N`` scenes: the last + slot is right, but every earlier slot holds scene 0's object in every env. That + is still better than the random default, which scrambles every slot. + + IsaacLab 3 picks which spawn prototype each env receives via a pluggable strategy, + and the default is ``random`` -- uniform sampling WITH REPLACEMENT. IsaacLab 2 used a + plain round-robin, which for one-prototype-per-env is the identity map ProtoMotions + relies on: env i must receive scene i's furniture, because env i is tracking the motion + authored against it. Under the random default the furniture is scrambled -- several envs + get the same mesh and others get none -- while positions, physics properties and the + motion are still applied per env index. The result looks like a plausible scene + containing the wrong object, and nothing warns. + + ``InteractiveScene`` builds its own ``CloneCfg`` and offers no injection point. Patching + ``CloneCfg`` does not work either: ``configclass`` turns the default into a + ``default_factory`` closure captured at class-creation time, so neither the class + attribute nor the field default is consulted afterwards. What ``InteractiveScene`` does + do is call ``cloner.ReplicateSession(...)`` by module-attribute lookup at call time -- + so wrapping that class is the one hook that takes effect. + """ + try: + import isaaclab.cloner as _cloner + from isaaclab.cloner.cloner_strategies import sequential + except ImportError: # pragma: no cover - older IsaacLab has no pluggable strategy + return + + session_cls = _cloner.ReplicateSession + if getattr(session_cls, "_protomotions_sequential", False): + return + + class _SequentialReplicateSession(session_cls): + _protomotions_sequential = True + + def __init__(self, *args, **kwargs): + kwargs["clone_strategy"] = sequential + super().__init__(*args, **kwargs) + + _cloner.ReplicateSession = _SequentialReplicateSession + log.info( + "Forced IsaacLab clone strategy to 'sequential' (IsaacLab 3 defaults to 'random')." + ) + + class IsaacLabSimulator(Simulator): config: IsaacLabSimulatorConfig @@ -125,6 +174,9 @@ def __init__( scene_cfg = self._get_scene_cfg() + # IsaacLab 3 defaults to a RANDOM clone strategy that scrambles which furniture + # each env receives; force sequential so env i gets scene i's object. + _force_sequential_cloning() self._scene = InteractiveScene(scene_cfg) if not self.headless: self._setup_keyboard() @@ -649,6 +701,33 @@ def _physics_step(self) -> None: self._sim.render() self._scene.update(dt=self._sim.get_physics_dt()) + def flush_contact_cache(self) -> None: + """Move every env out of contact and simulate one physics substep there. + + PhysX remembers each link/ground contact from the last simulated + substep, and writing a new root pose and joint state does not clear + that memory. After evaluation every env is written back to its saved + pre-evaluation state; a robot that was still walking when evaluation + ended is teleported into its saved mid-stride pose, and the first + substep can resolve the old contacts against the new pose and launch + it. Teleporting away alone is not enough: PhysX only drops a contact + after simulating a substep in which the bodies no longer touch. So every + robot goes 50 m underground and its scene objects 20 m further down, + one substep runs there, and the caller then overwrites all robot and + object state, so these poses never reach training. + + Measured with the real evaluator at 8192 envs: an Agile tracker had 47 + launches among 177,564 robots still walking when evaluation ended + (link speeds up to 4.5e15 m/s, occasionally a non-finite velocity that + crashed the run) and 0 of 177,386 with the flush; the G1 bones-deploy + tracker had 4 of 57,758 (up to 63 m/s) and 0 of 57,955 with it. The + flush costs about 0.1 s per evaluation. + """ + self._move_all_envs_out_of_contact(hide_z=-50.0, object_drop=20.0) + self._scene.write_data_to_sim() + self._sim.step(render=False) + self._scene.update(dt=self._sim.get_physics_dt()) + def _apply_simulator_pd_targets(self, pd_targets: torch.Tensor) -> None: """Applies PD position targets using IsaacLab's internal PD controller.""" self._robot.set_joint_position_target(pd_targets.detach(), joint_ids=None) diff --git a/protomotions/simulator/isaaclab/utils/mjcf_to_usd.py b/protomotions/simulator/isaaclab/utils/mjcf_to_usd.py index ae3dc105e..c78c76bdb 100644 --- a/protomotions/simulator/isaaclab/utils/mjcf_to_usd.py +++ b/protomotions/simulator/isaaclab/utils/mjcf_to_usd.py @@ -18,6 +18,7 @@ import os import xml.etree.ElementTree as ET from pathlib import Path +from tempfile import TemporaryDirectory from typing import Any, Callable, Dict, Mapping, MutableMapping, Optional, Tuple, Union from filelock import FileLock @@ -30,7 +31,7 @@ _CONVERSION_CACHE: Dict[Tuple[Any, ...], str] = {} ConverterFactory = Callable[..., str] -_MJCF_CONVERTER_CACHE_VERSION = "isaaclab3-d6-workaround-v2" +_MJCF_CONVERTER_CACHE_VERSION = "isaaclab3-d6-marker-inertia-v3" def _absolute_path(path: Union[os.PathLike, str]) -> str: @@ -256,7 +257,13 @@ def default_mjcf_converter_factory(**cfg_kwargs: Any) -> str: install_isaaclab_mjcf_d6_workaround() - converter = MjcfConverter(MjcfConverterCfg(**cfg_kwargs)) + from protomotions.simulator.base_simulator.robot_mass_properties import prepare_robot_mjcf + + with TemporaryDirectory(prefix="protomotions-robot-usd-") as directory: + prepared_path = prepare_robot_mjcf(cfg_kwargs["asset_path"], directory) + converter = MjcfConverter( + MjcfConverterCfg(**{**cfg_kwargs, "asset_path": prepared_path}) + ) return _absolute_path(converter.usd_path) diff --git a/protomotions/simulator/mujoco/config.py b/protomotions/simulator/mujoco/config.py index 5357ab917..a659615dd 100644 --- a/protomotions/simulator/mujoco/config.py +++ b/protomotions/simulator/mujoco/config.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field +from typing import Optional from protomotions.simulator.base_simulator.config import SimParams, SimulatorConfig @@ -11,7 +12,8 @@ class MujocoSimParams(SimParams): Defaults match RoboJuDo: 1000Hz physics, 20x decimation (50Hz control). No solver params -- we use MuJoCo's built-in defaults for everything - (solver, integrator, iterations, etc.) since RoboJuDo does the same. + (solver, iterations, etc.) since RoboJuDo does the same. The integrator + can be overridden per robot. """ fps: int = field( @@ -22,6 +24,15 @@ class MujocoSimParams(SimParams): default=20, metadata={"help": "Number of physics steps per control step."} ) + integrator: Optional[str] = field( + default=None, + metadata={ + "help": "MuJoCo integrator: 'euler', 'implicitfast', 'implicit' or " + "'rk4'. None keeps the MJCF's choice (MuJoCo's default is Euler). " + "Euler applies actuator damping explicitly, which can go unstable " + "with stiff PD gains on light links; 'implicitfast' does not." + }, + ) @dataclass @@ -39,10 +50,3 @@ class MujocoSimulatorConfig(SimulatorConfig): "False: explicit PD torque computation at each physics substep." }, ) - use_separate_keyboard_window: bool = field( - default=True, - metadata={ - "help": "Capture application keys in a separate Tk window instead " - "of competing with MuJoCo viewer keyboard shortcuts." - }, - ) diff --git a/protomotions/simulator/mujoco/keyboard.py b/protomotions/simulator/mujoco/keyboard.py index a9ec08ccb..0c4ddf5d2 100644 --- a/protomotions/simulator/mujoco/keyboard.py +++ b/protomotions/simulator/mujoco/keyboard.py @@ -1,160 +1,76 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Separate keyboard capture for the MuJoCo passive viewer. +"""Task keys in the MuJoCo viewer, without a second window.""" -MuJoCo owns keyboard focus for its viewer and reserves many keys for camera -and UI commands. This module keeps application controls in a small Tk window, -so the simulator can forward raw transitions without competing with viewer -callbacks. -""" +import logging +from queue import Empty, SimpleQueue +from typing import Iterator -from __future__ import annotations +import numpy as np -import multiprocessing as mp -import queue -import string -from typing import List, Tuple +log = logging.getLogger(__name__) -_CLOSE_EVENT = "__window_closed__" -_ERROR_EVENT = "__window_error__" -_KEYS = frozenset(string.ascii_letters + string.digits + string.punctuation + " ") +class MujocoViewerKeyboard: + """Queue viewer-thread presses; dispatch them on the simulation thread. - -def _printable_key(char: str, keysym: str) -> str: - """Return the one-character key represented by a Tk event. - - Tk normally fills ``event.char`` for key presses, but key-release events - often expose only ``event.keysym``. Falling back to ``keysym`` prevents a - dropped release from latching a key and suppressing later clicks. + MuJoCo also runs its own shortcuts before notifying our callback. Task + mode restores visualization settings so W/R/L and digits do not leave + wireframe, reflections or geom groups toggled. Period switches to native + mode for intentional viewer adjustments, without issuing task commands. """ - key = (char or keysym or "").lower() - if len(key) != 1 or key not in _KEYS: - return "" - return key - - -def _run_keyboard_window(events, ready) -> None: - """Run the Tk event loop in a process that owns its GUI resources.""" - try: - import tkinter as tk - - root = tk.Tk() - root.title("ProtoMotions MuJoCo controls") - root.geometry("420x100") - root.resizable(False, False) - label = tk.Label( - root, - text=( - "Click here, then press your registered ProtoMotions keys.\n" - "This window owns task keys; MuJoCo remains visible for viewing." - ), - padx=12, - pady=12, - ) - label.pack(fill=tk.BOTH, expand=True) - - def emit(event, pressed: bool) -> None: - key = _printable_key( - getattr(event, "char", ""), getattr(event, "keysym", "") - ) - if not key: - return - try: - events.put_nowait((key, pressed)) - except queue.Full: - # A held key can generate OS repeat events. Dropping a - # repeated transition is preferable to blocking the GUI. - pass - - root.bind_all("", lambda event: emit(event, True)) - root.bind_all("", lambda event: emit(event, False)) - - def close() -> None: - try: - events.put_nowait((_CLOSE_EVENT, False)) - except queue.Full: - pass - root.destroy() - - root.protocol("WM_DELETE_WINDOW", close) - # Keep task-key focus within this Tk application without grabbing the - # entire desktop. A global grab would also intercept window-manager - # shortcuts and make the user's session appear frozen. - root.deiconify() - root.update_idletasks() - root.update() - root.grab_set() - root.focus_force() - ready.set() - root.mainloop() - except BaseException as exc: - try: - events.put_nowait((_ERROR_EVENT, repr(exc))) - except queue.Full: - pass - ready.set() - - -class MujocoKeyboardWindow: - """Own a separate process that captures printable key transitions.""" def __init__(self) -> None: - context = mp.get_context("spawn") - self._events = context.Queue(maxsize=256) - self._ready = context.Event() - self._process = context.Process( - target=_run_keyboard_window, - args=(self._events, self._ready), - daemon=True, - ) - self._process.start() - if not self._ready.wait(timeout=5.0): - self.close() - raise RuntimeError("MuJoCo keyboard window did not start") - if not self._process.is_alive(): - error = self._read_startup_error() - self.close() - raise RuntimeError(f"MuJoCo keyboard window failed to start: {error}") + self._events = SimpleQueue() + self.native = False + self._baseline = None - def _read_startup_error(self) -> str: - try: - event = self._events.get_nowait() - except queue.Empty: - return "keyboard window process exited" - if event[0] == _ERROR_EVENT: - return event[1] - return repr(event) + def on_key(self, keycode: int) -> None: + # The passive viewer supplies discrete presses, not key releases. + # Never mutate physics or invoke application callbacks on its thread. + if 33 <= keycode <= 126: + self._events.put(chr(keycode)) - def drain_events(self) -> List[Tuple[str, bool]]: - """Return all queued ``(key, pressed)`` transitions.""" - events: List[Tuple[str, bool]] = [] + def drain_presses(self) -> Iterator[str]: while True: try: - event = self._events.get_nowait() - except queue.Empty: - break + yield self._events.get_nowait() + except Empty: + return - key, value = event - if key == _ERROR_EVENT: - raise RuntimeError(f"MuJoCo keyboard window failed: {value}") - if key == _CLOSE_EVENT: - raise RuntimeError("MuJoCo keyboard window was closed") - events.append((key, bool(value))) + def toggle_native(self) -> None: + self.native = not self.native + log.info( + "MuJoCo keyboard: %s", + "native viewer shortcuts" if self.native else "task controls", + ) - if not self._process.is_alive(): - raise RuntimeError("MuJoCo keyboard window exited unexpectedly") - return events + def capture(self, viewer) -> None: + with viewer.lock(): + self._baseline = ( + viewer.opt.flags.copy(), + viewer.opt.geomgroup.copy(), + viewer.user_scn.flags.copy(), + ) - def close(self) -> None: - """Stop the child process and release its IPC resources.""" - process = getattr(self, "_process", None) - if process is not None: - if process.is_alive(): - process.terminate() - process.join(timeout=1.0) - events = getattr(self, "_events", None) - if events is not None: - events.close() - events.join_thread() + def sync(self, viewer) -> None: + if self.native: + viewer.sync() + return + if self._baseline is None: + self.capture(viewer) + flags, groups, render_flags = self._baseline + with viewer.lock(): + viewer.opt.flags[:] = flags + viewer.opt.geomgroup[:] = groups + viewer.sync() + # sync first copies native render flags into user_scn. Only a *change* + # to user_scn is propagated back, so setting it before that first sync + # cannot undo a native shortcut. A second sync is needed only on change. + with viewer.lock(): + changed = not np.array_equal(viewer.user_scn.flags, render_flags) + if changed: + viewer.user_scn.flags[:] = render_flags + if changed: + viewer.sync() diff --git a/protomotions/simulator/mujoco/recording.py b/protomotions/simulator/mujoco/recording.py new file mode 100644 index 000000000..01ff283eb --- /dev/null +++ b/protomotions/simulator/mujoco/recording.py @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Offscreen capture of MuJoCo's live camera and debug geometry. + +Callers hold the passive viewer lock while reading its state and rendering. +Clip lifetime and video encoding remain in the shared recording mixin. +""" + +from __future__ import annotations + +import mujoco +import numpy as np + + +def recording_dimensions(viewer) -> tuple[int, int]: + """Return even viewport dimensions (width, height), or a 640×480 fallback.""" + viewport = getattr(viewer, "viewport", None) + width, height = 640, 480 + if viewport is not None and min(viewport.width, viewport.height) > 0: + width, height = viewport.width, viewport.height + # H.264's usual yuv420p pixel format requires even dimensions. + return max(2, int(width) // 2 * 2), max(2, int(height) // 2 * 2) + + +def render_viewport( + model: mujoco.MjModel, + data: mujoco.MjData, + viewer, + width: int, + height: int, +) -> np.ndarray: + """Render RGB pixels with the live camera, options, flags, and markers.""" + model.vis.global_.offwidth = max(model.vis.global_.offwidth, width) + model.vis.global_.offheight = max(model.vis.global_.offheight, height) + # update_scene rebuilds model geometry only; user_scn holds goal markers + # and arrows that must be appended to the offscreen scene separately. + debug_scene = viewer.user_scn if viewer is not None else None + renderer = mujoco.Renderer( + model, + height=height, + width=width, + max_geom=max(10000, model.ngeom) + + (debug_scene.maxgeom if debug_scene is not None else 0), + ) + try: + if viewer is not None: + renderer.update_scene(data, camera=viewer.cam, scene_option=viewer.opt) + scene = renderer.scene + scene.flags[:] = debug_scene.flags + for i in range(debug_scene.ngeom): + source = debug_scene.geoms[i] + dest = scene.geoms[scene.ngeom] + for field in ("size", "pos", "mat", "rgba"): + getattr(dest, field)[:] = getattr(source, field) + # Preserve material, label and segmentation metadata too. + for field in ( + "type", + "dataid", + "objtype", + "objid", + "category", + "matid", + "texcoord", + "segid", + "emission", + "specular", + "shininess", + "reflectance", + "camdist", + "modelrbound", + "transparent", + "label", + ): + setattr(dest, field, getattr(source, field)) + scene.ngeom += 1 + else: + renderer.update_scene(data) + return renderer.render() + finally: + renderer.close() diff --git a/protomotions/simulator/mujoco/simulator.py b/protomotions/simulator/mujoco/simulator.py index 908001604..ec9e420c1 100644 --- a/protomotions/simulator/mujoco/simulator.py +++ b/protomotions/simulator/mujoco/simulator.py @@ -17,7 +17,6 @@ log = logging.getLogger(__name__) -from protomotions.assets import resolve_asset_root from protomotions.simulator.base_simulator.config import ( MarkerState, ProjectileConfig, @@ -33,7 +32,17 @@ StateConversion, ) from protomotions.simulator.mujoco.config import MujocoSimulatorConfig -from protomotions.simulator.mujoco.keyboard import MujocoKeyboardWindow +from protomotions.simulator.mujoco.keyboard import MujocoViewerKeyboard +from protomotions.simulator.mujoco.recording import recording_dimensions, render_viewport + + +# Names accepted by MujocoSimParams.integrator +_INTEGRATORS = { + "euler": mujoco.mjtIntegrator.mjINT_EULER, + "rk4": mujoco.mjtIntegrator.mjINT_RK4, + "implicit": mujoco.mjtIntegrator.mjINT_IMPLICIT, + "implicitfast": mujoco.mjtIntegrator.mjINT_IMPLICITFAST, +} def _to_torch_f32(arr: np.ndarray) -> torch.Tensor: @@ -41,6 +50,28 @@ def _to_torch_f32(arr: np.ndarray) -> torch.Tensor: return torch.from_numpy(arr.astype(np.float32)) +# A free joint's qvel[3:6] is angular velocity in the BODY frame, while its qvel[0:3] +# is linear velocity in the world frame. ProtoMotions' root_ang_vel is world-frame, so +# every read and write of qvel[3:6] goes through these two helpers. +def _world_to_body(quat_wxyz: np.ndarray, vec: np.ndarray) -> np.ndarray: + """Express a world-frame vector in the frame of a body oriented by quat_wxyz.""" + quat = np.asarray(quat_wxyz, dtype=np.float64) + inverse = quat * np.array([1.0, -1.0, -1.0, -1.0]) / np.linalg.norm(quat) + out = np.zeros(3) + mujoco.mju_rotVecQuat(out, np.asarray(vec, dtype=np.float64), inverse) + return out + + +def _body_to_world(quat_wxyz: np.ndarray, vec: np.ndarray) -> np.ndarray: + """Express a body-frame vector in the world frame, for a body oriented by quat_wxyz.""" + quat = np.asarray(quat_wxyz, dtype=np.float64) + out = np.zeros(3) + mujoco.mju_rotVecQuat( + out, np.asarray(vec, dtype=np.float64), quat / np.linalg.norm(quat) + ) + return out + + class MujocoSimulator(Simulator): """MuJoCo CPU-only simulator backend. @@ -88,7 +119,7 @@ def __init__( self.data: Optional[mujoco.MjData] = None self.viewer = None self._viewer_initialized = False - self._keyboard_window = None + self._viewer_keyboard = MujocoViewerKeyboard() # Cached control parameters self._kp = None # [num_dofs] stiffness in common DOF order @@ -292,9 +323,7 @@ def _add_ground_and_light(root: ET.Element) -> None: def _create_simulation(self) -> None: """Create the MuJoCo simulation environment.""" # Load MJCF model - asset_root = resolve_asset_root(self.robot_config.asset.asset_root) - asset_file = self.robot_config.asset.asset_file_name - asset_path = os.path.join(asset_root, asset_file) + asset_path = self._get_robot_asset_path() # MJCF injection below needs _proj_config before _init_projectiles runs self._resolve_proj_config() @@ -310,12 +339,28 @@ def _create_simulation(self) -> None: f"MuJoCo timestep: {self.model.opt.timestep:.4f}s ({self.config.sim.fps}Hz)" ) + # Integrator: None keeps whatever the MJCF's