├── .github └── workflows │ └── ci.yml ├── .gitignore ├── .pre-commit-config.yaml ├── CHANGELOG.md ├── CITATION.cff ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── assets └── banner.png ├── learning ├── README.md ├── notebooks │ ├── dm_control_suite.ipynb │ ├── locomotion.ipynb │ ├── manipulation.ipynb │ ├── training_vision_1.ipynb │ ├── training_vision_1_t4.ipynb │ └── training_vision_2.ipynb ├── train_jax_ppo.py └── train_rsl_rl.py ├── mujoco_playground ├── __init__.py ├── _src │ ├── __init__.py │ ├── collision.py │ ├── dm_control_suite │ │ ├── README.md │ │ ├── __init__.py │ │ ├── acrobot.py │ │ ├── ball_in_cup.py │ │ ├── cartpole.py │ │ ├── cheetah.py │ │ ├── common.py │ │ ├── dm_control_suite_test.py │ │ ├── finger.py │ │ ├── fish.py │ │ ├── hopper.py │ │ ├── humanoid.py │ │ ├── pendulum.py │ │ ├── point_mass.py │ │ ├── reacher.py │ │ ├── swimmer.py │ │ ├── walker.py │ │ └── xmls │ │ │ ├── README.md │ │ │ ├── acrobot.xml │ │ │ ├── ball_in_cup.xml │ │ │ ├── cartpole.xml │ │ │ ├── cheetah.xml │ │ │ ├── common │ │ │ ├── materials.xml │ │ │ ├── skybox.xml │ │ │ └── visual.xml │ │ │ ├── finger.xml │ │ │ ├── fish.xml │ │ │ ├── hopper.xml │ │ │ ├── humanoid.xml │ │ │ ├── manipulator.xml │ │ │ ├── pendulum.xml │ │ │ ├── point_mass.xml │ │ │ ├── reacher.xml │ │ │ ├── swimmer.xml │ │ │ └── walker.xml │ ├── gait.py │ ├── locomotion │ │ ├── __init__.py │ │ ├── apollo │ │ │ ├── __init__.py │ │ │ ├── base.py │ │ │ ├── constants.py │ │ │ ├── joystick.py │ │ │ └── xmls │ │ │ │ ├── apollo_mjx_feetonly.xml │ │ │ │ └── scene_mjx_feetonly_flat_terrain.xml │ │ ├── barkour │ │ │ ├── __init__.py │ │ │ └── joystick.py │ │ ├── berkeley_humanoid │ │ │ ├── __init__.py │ │ │ ├── base.py │ │ │ ├── berkeley_humanoid_constants.py │ │ │ ├── joystick.py │ │ │ ├── randomize.py │ │ │ └── xmls │ │ │ │ ├── assets │ │ │ │ ├── hfield.png │ │ │ │ └── rocky_texture.png │ │ │ │ ├── berkeley_humanoid_mjx_feetonly.xml │ │ │ │ ├── scene_mjx_feetonly_flat_terrain.xml │ │ │ │ └── scene_mjx_feetonly_rough_terrain.xml │ │ ├── g1 │ │ │ ├── __init__.py │ │ │ ├── base.py │ │ │ ├── g1_constants.py │ │ │ ├── joystick.py │ │ │ ├── randomize.py │ │ │ └── xmls │ │ │ │ ├── assets │ │ │ │ ├── hfield.png │ │ │ │ └── rocky_texture.png │ │ │ │ ├── g1_mjx_feetonly.xml │ │ │ │ ├── g1_mjx_feetonly_restricted.xml │ │ │ │ ├── scene_mjx_feetonly.xml │ │ │ │ ├── scene_mjx_feetonly_flat_terrain.xml │ │ │ │ └── scene_mjx_feetonly_rough_terrain.xml │ │ ├── go1 │ │ │ ├── README.md │ │ │ ├── __init__.py │ │ │ ├── base.py │ │ │ ├── getup.py │ │ │ ├── go1_constants.py │ │ │ ├── handstand.py │ │ │ ├── joystick.py │ │ │ ├── randomize.py │ │ │ └── xmls │ │ │ │ ├── assets │ │ │ │ ├── hfield.png │ │ │ │ └── rocky_texture.png │ │ │ │ ├── go1_mjx.xml │ │ │ │ ├── go1_mjx_feetonly.xml │ │ │ │ ├── go1_mjx_fullcollisions.xml │ │ │ │ ├── scene_mjx_feetonly_bowl.xml │ │ │ │ ├── scene_mjx_feetonly_flat_terrain.xml │ │ │ │ ├── scene_mjx_feetonly_rough_terrain.xml │ │ │ │ ├── scene_mjx_feetonly_stairs.xml │ │ │ │ ├── scene_mjx_flat_terrain.xml │ │ │ │ └── scene_mjx_fullcollisions_flat_terrain.xml │ │ ├── h1 │ │ │ ├── __init__.py │ │ │ ├── base.py │ │ │ ├── h1_constants.py │ │ │ ├── inplace_gait_tracking.py │ │ │ ├── joystick.py │ │ │ ├── joystick_gait_tracking.py │ │ │ └── xmls │ │ │ │ ├── h1_mjx_feetonly.xml │ │ │ │ └── scene_mjx_feetonly.xml │ │ ├── locomotion_test.py │ │ ├── op3 │ │ │ ├── __init__.py │ │ │ ├── base.py │ │ │ ├── joystick.py │ │ │ ├── op3_constants.py │ │ │ └── xmls │ │ │ │ ├── op3_mjx_feetonly.xml │ │ │ │ └── scene_mjx_feetonly.xml │ │ ├── spot │ │ │ ├── __init__.py │ │ │ ├── base.py │ │ │ ├── getup.py │ │ │ ├── joystick.py │ │ │ ├── joystick_gait_tracking.py │ │ │ ├── spot_constants.py │ │ │ └── xmls │ │ │ │ ├── scene_mjx_feetonly_flat_terrain.xml │ │ │ │ ├── scene_mjx_flat_terrain.xml │ │ │ │ ├── spot_mjx.xml │ │ │ │ └── spot_mjx_feetonly.xml │ │ └── t1 │ │ │ ├── __init__.py │ │ │ ├── base.py │ │ │ ├── joystick.py │ │ │ ├── randomize.py │ │ │ ├── t1_constants.py │ │ │ └── xmls │ │ │ ├── assets │ │ │ ├── hfield.png │ │ │ └── rocky_texture.png │ │ │ ├── scene_mjx_feetonly_flat_terrain.xml │ │ │ ├── scene_mjx_feetonly_rough_terrain.xml │ │ │ └── t1_mjx_feetonly.xml │ ├── manipulation │ │ ├── __init__.py │ │ ├── aloha │ │ │ ├── __init__.py │ │ │ ├── aloha_constants.py │ │ │ ├── base.py │ │ │ ├── handover.py │ │ │ ├── single_peg_insertion.py │ │ │ └── xmls │ │ │ │ ├── assets │ │ │ │ ├── socket_blue.png │ │ │ │ └── socket_white.png │ │ │ │ ├── joint_position_actuators.xml │ │ │ │ ├── mjx_aloha.xml │ │ │ │ ├── mjx_hand_over.xml │ │ │ │ ├── mjx_scene.xml │ │ │ │ └── mjx_single_peg_insertion.xml │ │ ├── franka_emika_panda │ │ │ ├── __init__.py │ │ │ ├── open_cabinet.py │ │ │ ├── panda.py │ │ │ ├── panda_kinematics.py │ │ │ ├── pick.py │ │ │ ├── pick_cartesian.py │ │ │ ├── randomize_vision.py │ │ │ └── xmls │ │ │ │ ├── mjx_cabinet.xml │ │ │ │ ├── mjx_scene.xml │ │ │ │ ├── mjx_single_cube.xml │ │ │ │ └── mjx_single_cube_camera.xml │ │ ├── franka_emika_panda_robotiq │ │ │ ├── __init__.py │ │ │ ├── assets │ │ │ │ ├── fileback.png │ │ │ │ ├── fileback_light.png │ │ │ │ ├── filedown.png │ │ │ │ ├── filedown_light.png │ │ │ │ ├── filefront.png │ │ │ │ ├── filefront_light.png │ │ │ │ ├── fileleft.png │ │ │ │ ├── fileleft_light.png │ │ │ │ ├── fileright.png │ │ │ │ ├── fileright_light.png │ │ │ │ ├── fileup.png │ │ │ │ └── fileup_light.png │ │ │ ├── panda_robotiq.py │ │ │ ├── push_cube.py │ │ │ └── xmls │ │ │ │ ├── panda_updated_robotiq_2f85.xml │ │ │ │ └── scene_panda_robotiq_cube.xml │ │ ├── leap_hand │ │ │ ├── README.md │ │ │ ├── __init__.py │ │ │ ├── base.py │ │ │ ├── leap_hand_constants.py │ │ │ ├── reorient.py │ │ │ ├── rotate_z.py │ │ │ └── xmls │ │ │ │ ├── leap_rh_mjx.xml │ │ │ │ ├── meshes │ │ │ │ ├── dex_cube.obj │ │ │ │ └── leap_mount.obj │ │ │ │ ├── reorientation_cube.xml │ │ │ │ ├── reorientation_cube_textures │ │ │ │ ├── dex_cube.png │ │ │ │ ├── fileback.png │ │ │ │ ├── filedown.png │ │ │ │ ├── filefront.png │ │ │ │ ├── fileleft.png │ │ │ │ ├── fileright.png │ │ │ │ ├── fileup.png │ │ │ │ ├── grayback.png │ │ │ │ ├── graydown.png │ │ │ │ ├── grayfront.png │ │ │ │ ├── grayleft.png │ │ │ │ ├── grayright.png │ │ │ │ └── grayup.png │ │ │ │ └── scene_mjx_cube.xml │ │ └── manipulation_test.py │ ├── mjx_env.py │ ├── registry.py │ ├── registry_test.py │ ├── reward.py │ ├── wrapper.py │ ├── wrapper_test.py │ └── wrapper_torch.py ├── config │ ├── README.md │ ├── __init__.py │ ├── dm_control_suite_params.py │ ├── locomotion_params.py │ └── manipulation_params.py ├── experimental │ ├── brax_network_to_onnx.ipynb │ ├── learning │ │ ├── README.md │ │ ├── ablate_device_topology.ipynb │ │ ├── apollo_joystick.ipynb │ │ ├── cube_reorient.ipynb │ │ ├── cube_rotate.ipynb │ │ ├── dm_control_suite.ipynb │ │ ├── gait_tracking.ipynb │ │ ├── getup.ipynb │ │ ├── handstand.ipynb │ │ ├── humanoid_joystick.ipynb │ │ ├── humanoid_joystick_gait_tracking.ipynb │ │ ├── inference.ipynb │ │ ├── joystick.ipynb │ │ ├── joystick_gait_tracking.ipynb │ │ ├── open_cabinet.ipynb │ │ ├── phase.ipynb │ │ ├── position_cube.ipynb │ │ └── viz_sim_policies.ipynb │ ├── madrona_benchmarking │ │ ├── README.md │ │ ├── benchmark.py │ │ ├── data │ │ │ ├── madrona_mjx.csv │ │ │ └── maniskill3 │ │ │ │ ├── README.md │ │ │ │ ├── isaac_lab.csv │ │ │ │ ├── isaac_lab_state.csv │ │ │ │ ├── maniskill.csv │ │ │ │ └── maniskill_state.csv │ │ ├── figures │ │ │ ├── cartpole_benchmark.png │ │ │ ├── cartpole_benchmark_combined.png │ │ │ ├── cartpole_benchmark_full.png │ │ │ ├── cartpole_benchmark_stacked.png │ │ │ ├── cartpole_benchmark_stacked_full.png │ │ │ ├── cartpole_benchmark_state.png │ │ │ └── madrona_mjx_main.png │ │ ├── get_data.sh │ │ ├── make_plots.py │ │ └── print_tables.py │ ├── sim2sim │ │ ├── README.md │ │ ├── assets │ │ │ └── f710.jpg │ │ ├── gamepad_reader.py │ │ ├── onnx │ │ │ ├── apollo_policy.onnx │ │ │ ├── bh_policy.onnx │ │ │ ├── g1_policy.onnx │ │ │ ├── go1_policy.onnx │ │ │ ├── leap_reorient_policy.onnx │ │ │ └── t1_policy.onnx │ │ ├── play_apollo_joystick.py │ │ ├── play_bh_joystick.py │ │ ├── play_g1_joystick.py │ │ ├── play_go1_joystick.py │ │ ├── play_leap_reorient.py │ │ └── play_t1_joystick.py │ └── utils │ │ └── plotting.py └── py.typed ├── pylintrc └── pyproject.toml /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: build 2 | 3 | on: 4 | push: 5 | branches: [main] 6 | pull_request: 7 | branches: [main] 8 | 9 | jobs: 10 | build: 11 | runs-on: ubuntu-latest 12 | strategy: 13 | matrix: 14 | python-version: ["3.10", "3.11", "3.12"] 15 | 16 | steps: 17 | - uses: actions/checkout@v2 18 | - name: Set up Python ${{ matrix.python-version }} 19 | uses: actions/setup-python@v1 20 | with: 21 | python-version: ${{ matrix.python-version }} 22 | - name: Install dependencies 23 | run: | 24 | python -m pip install --upgrade pip 25 | pip install uv 26 | uv pip install --system -e ".[test]" 27 | - name: Trigger git clone 28 | run: | 29 | python -c "import mujoco_playground" 30 | - name: Test with pytest 31 | run: | 32 | pytest -n auto mujoco_playground/_src/ 33 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | __pycache__ 3 | *.egg-info 4 | .venv 5 | MUJOCO_LOG.TXT 6 | mujoco_menagerie 7 | checkpoints/ 8 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | default_stages: [pre-commit] 2 | exclude: '.*__init__\.py$' 3 | 4 | # Install 5 | # 1. pip install -e . 6 | # 2. pre-commit install 7 | # 3. pre-commit run --all-files # make sure all files are clean 8 | repos: 9 | - repo: https://github.com/google/pyink 10 | rev: 24.10.0 11 | hooks: 12 | - id: pyink 13 | name: Pyink (Formatting) 14 | # pyink will automatically read configuration from pyproject.toml 15 | 16 | - repo: https://github.com/pycqa/isort 17 | rev: 5.13.2 18 | hooks: 19 | - id: isort 20 | name: Isort (Import Sorting) 21 | # isort will automatically read configuration from pyproject.toml 22 | 23 | # - repo: local 24 | # hooks: 25 | # - id: pylint 26 | # name: Pylint (Google Style) 27 | # entry: pylint 28 | # args: ['--rcfile=pylintrc'] 29 | # language: system 30 | # types: [python] 31 | 32 | # - repo: local # re-using locally installed libraries 33 | # hooks: 34 | # - id: pytype 35 | # name: Pytype (Type Checking) 36 | # entry: pytype 37 | # language: system 38 | # types: [python] 39 | # always_run: true 40 | 41 | # - repo: local 42 | # hooks: 43 | # - id: pytest 44 | # name: Run Pytest 45 | # entry: pytest 46 | # language: system 47 | # pass_filenames: false 48 | # always_run: true 49 | # args: ["-n", "auto"] # Execute tests in parallel 50 | # verbose: true 51 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | 3 | All notable changes to this project will be documented in this file. 4 | 5 | ## [0.0.4] - 2025-02-07 6 | 7 | ### Added 8 | 9 | - Added ALOHA handover task (thanks to @Andrew-Luo1). 10 | - Added Booster T1 joystick task. 11 | 12 | ### Changed 13 | 14 | - Fixed foot friction randomization for G1 tasks. 15 | - Fix various bugs in `train_jax_ppo.py` (thanks to @vincentzhang). 16 | - Fixed a small bug in the privileged state of the go1 joystick task. 17 | 18 | ## [0.0.3] - 2025-01-18 19 | 20 | ### Changed 21 | 22 | - Updated supported Python versions to 3.10-3.12. 23 | 24 | ## [0.0.2] - 2025-01-16 25 | 26 | Initial release. 27 | -------------------------------------------------------------------------------- /CITATION.cff: -------------------------------------------------------------------------------- 1 | cff-version: 1.2.0 2 | message: "If you use this software, please cite it using the following reference:" 3 | title: "MuJoCo Playground: An open-source framework for GPU-accelerated robot learning and sim-to-real transfer." 4 | authors: 5 | - family-names: Zakka 6 | given-names: Kevin 7 | - family-names: Tabanpour 8 | given-names: Baruch 9 | - family-names: Liao 10 | given-names: Qiayuan 11 | - family-names: Haiderbhai 12 | given-names: Mustafa 13 | - family-names: Holt 14 | given-names: Samuel 15 | - family-names: Luo 16 | given-names: Jing Yuan 17 | - family-names: Allshire 18 | given-names: Arthur 19 | - family-names: Frey 20 | given-names: Erik 21 | - family-names: Sreenath 22 | given-names: Koushil 23 | - family-names: Kahrs 24 | given-names: Lueder A. 25 | - family-names: Sferrazza 26 | given-names: Carlo 27 | - family-names: Tassa 28 | given-names: Yuval 29 | - family-names: Abbeel 30 | given-names: Pieter 31 | year: 2025 32 | repository-code: "https://github.com/google-deepmind/mujoco_playground" 33 | url: "https://github.com/google-deepmind/mujoco_playground" 34 | license: "Apache-2.0" 35 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # How to Contribute 2 | 3 | ## Contributor License Agreement 4 | 5 | Contributions to this project must be accompanied by a Contributor License 6 | Agreement. You (or your employer) retain the copyright to your contribution, 7 | this simply gives us permission to use and redistribute your contributions as 8 | part of the project. Head over to to see 9 | your current agreements on file or to sign a new one. 10 | 11 | You generally only need to submit a CLA once, so if you've already submitted one 12 | (even if it was for a different project), you probably don't need to do it 13 | again. 14 | 15 | ## Adding New Tasks 16 | 17 | We welcome contributions of new tasks, particularly those that demonstrate: 18 | 19 | - Impressive capabilities in simulation 20 | - Tasks that have been successfully transferred to real robots 21 | 22 | When submitting a new task, please ensure: 23 | 24 | 1. The task is well-documented with clear objectives and reward structure 25 | 2. Add the relevant RL hyperparameters to the config file so that it can be 26 | independently reproduced 27 | 3. Ensure that it works across at least 3 seeds 28 | 4. Show a video of the behavior 29 | 5. Make sure your new task passes all the tests 30 | 31 | For an example of a well-structured task contribution, see @Andrew-Luo1's 32 | excellent [ALOHA Handover Task 33 | PR](https://github.com/google-deepmind/mujoco_playground/pull/29). 34 | 35 | ## Code reviews 36 | 37 | All submissions, including submissions by project members, require review. We 38 | use GitHub pull requests for this purpose. Consult 39 | [GitHub Help](https://help.github.com/articles/about-pull-requests/) for more 40 | information on using pull requests. 41 | 42 | ## Community Guidelines 43 | 44 | This project follows [Google's Open Source Community 45 | Guidelines](https://opensource.google/conduct/). 46 | 47 | ## Linting and Code Health 48 | 49 | Before submitting a PR, please run: 50 | 51 | ```shell 52 | pip install -e ".[dev]" 53 | pre-commit install 54 | pre-commit run --all-files 55 | ``` 56 | 57 | or you can run manually 58 | 59 | ```shell 60 | pyink . 61 | isort . 62 | pylint . --rcfile=pylintrc 63 | pytype . 64 | ``` 65 | 66 | and resolve any issues that pop up. 67 | 68 | ## Testing 69 | 70 | To run the tests, use the following command: 71 | 72 | ```shell 73 | pytest 74 | ``` 75 | 76 | -------------------------------------------------------------------------------- /assets/banner.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/assets/banner.png -------------------------------------------------------------------------------- /learning/README.md: -------------------------------------------------------------------------------- 1 | # Learning RL Agents 2 | 3 | In this directory, we demonstrate learning RL agents from MuJoCo Playground environments using [Brax](https://github.com/google/brax) and [RSL-RL](https://github.com/leggedrobotics/rsl_rl). We provide two entrypoints from the command line: `python train_jax_ppo.py` and `python train_rsl_rl.py`. 4 | 5 | For more detailed tutorials on using MuJoCo Playground for RL, see: 6 | 7 | 1. Intro. to the Playground with DM Control Suite [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/dm_control_suite.ipynb) 8 | 2. Locomotion Environments [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/locomotion.ipynb) 9 | 3. Manipulation Environments [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/manipulation.ipynb) 10 | 4. Training CartPole from Vision [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/training_vision_1.ipynb) 11 | 5. Robotic Manipulation from Vision [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-deepmind/mujoco_playground/blob/main/learning/notebooks/training_vision_2.ipynb) 12 | 13 | ## Training with brax PPO 14 | 15 | To train with brax PPO, you can use the `train_jax_ppo.py` script. This script uses the brax PPO algorithm to train an agent on a given environment. 16 | 17 | ```bash 18 | python train_jax_ppo.py --env_name=CartpoleBalance 19 | ``` 20 | 21 | To train a vision-based policy using pixel observations: 22 | ```bash 23 | python train_jax_ppo.py --env_name=CartpoleBalance --vision 24 | ``` 25 | 26 | Use `python train_jax_ppo.py --help` to see possible options and usage. Logs and checkpoints are saved in `logs` directory. 27 | 28 | ## Training with RSL-RL 29 | 30 | To train with RSL-RL, you can use the `train_rsl_rl.py` script. This script uses the RSL-RL algorithm to train an agent on a given environment. 31 | 32 | ```bash 33 | python train_rsl_rl.py --env_name=LeapCubeReorient 34 | ``` 35 | 36 | To render the behaviour from the resulting policy: 37 | ```bash 38 | python learning/train_rsl_rl.py --env_name LeapCubeReorient --play_only --load_run_name 39 | ``` 40 | 41 | where `run_name` is the name of the run you want to load (will be printed in the console when the training run is started). 42 | 43 | Logs and checkpoints are saved in `logs` directory. 44 | -------------------------------------------------------------------------------- /mujoco_playground/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """MuJoCo Playground.""" 16 | from mujoco_playground._src import dm_control_suite 17 | from mujoco_playground._src import locomotion 18 | from mujoco_playground._src import manipulation 19 | from mujoco_playground._src import registry 20 | from mujoco_playground._src import wrapper 21 | from mujoco_playground._src import wrapper_torch 22 | # pylint: disable=g-importing-member 23 | from mujoco_playground._src.mjx_env import init 24 | from mujoco_playground._src.mjx_env import MjxEnv 25 | from mujoco_playground._src.mjx_env import render_array 26 | from mujoco_playground._src.mjx_env import State 27 | from mujoco_playground._src.mjx_env import step 28 | 29 | # pylint: enable=g-importing-member 30 | 31 | __all__ = [ 32 | "dm_control_suite", 33 | "init", 34 | "locomotion", 35 | "manipulation", 36 | "MjxEnv", 37 | "registry", 38 | "render_array", 39 | "State", 40 | "step", 41 | "wrapper", 42 | "wrapper_torch", 43 | ] 44 | -------------------------------------------------------------------------------- /mujoco_playground/_src/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/_src/collision.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Utilities for extracting collision information.""" 16 | 17 | from typing import Any, Tuple 18 | 19 | import jax 20 | import jax.numpy as jnp 21 | from mujoco import mjx 22 | 23 | 24 | def get_collision_info( 25 | contact: Any, geom1: int, geom2: int 26 | ) -> Tuple[jax.Array, jax.Array]: 27 | """Get the distance and normal of the collision between two geoms.""" 28 | mask = (jnp.array([geom1, geom2]) == contact.geom).all(axis=1) 29 | mask |= (jnp.array([geom2, geom1]) == contact.geom).all(axis=1) 30 | idx = jnp.where(mask, contact.dist, 1e4).argmin() 31 | dist = contact.dist[idx] * mask[idx] 32 | normal = (dist < 0) * contact.frame[idx, 0, :3] 33 | return dist, normal 34 | 35 | 36 | def geoms_colliding(state: mjx.Data, geom1: int, geom2: int) -> jax.Array: 37 | """Return True if the two geoms are colliding.""" 38 | return get_collision_info(state.contact, geom1, geom2)[0] < 0 39 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/README.md: -------------------------------------------------------------------------------- 1 | # DeepMind Control Suite 2 | 3 | A re-implementation of the [DeepMind Control Suite](https://arxiv.org/abs/1801.00690) using [MJX](https://mujoco.readthedocs.io/en/latest/mjx.html). 4 | 5 | ## Contents 6 | 7 | * ✅: Implemented 8 | * ❌: Not implemented 9 | 10 | | Env | Implemented | 11 | | ------------------------- | ------------- | 12 | | acrobot-swingup | ✅ | 13 | | acrobot-swingup_sparse | ✅ | 14 | | ball_in_cup-catch | ✅ | 15 | | cartpole-balance | ✅ | 16 | | cartpole-balance_sparse | ✅ | 17 | | cartpole-swingup | ✅ | 18 | | cartpole-swingup_sparse | ✅ | 19 | | cheetah-run | ✅ | 20 | | finger-spin | ✅ | 21 | | finger_turn_easy | ✅ | 22 | | finger_turn_hard | ✅ | 23 | | fish-upright | ✅ | 24 | | fish-swim | ✅ | 25 | | hopper-stand | ✅ | 26 | | hopper-hop | ✅ | 27 | | humanoid-stand | ✅ | 28 | | humanoid-walk | ✅ | 29 | | humanoid-run | ✅ | 30 | | pendulum-swingup | ✅ | 31 | | point_mass-easy | ✅ | 32 | | reacher-easy | ✅ | 33 | | reacher-hard | ✅ | 34 | | swimmer-swimmer6 | ✅ | 35 | | swimmer-swimmer15 | ✅ | 36 | | walker-stand | ✅ | 37 | | walker-walk | ✅ | 38 | | walker-run | ✅ | 39 | | manipulator-bring_ball | ❌ | 40 | | manipulator-bring_peg | ❌ | 41 | | manipulator-insert_ball | ❌ | 42 | | manipulator-insert_peg | ❌ | 43 | | dog-stand | ❌ | 44 | | dog-walk | ❌ | 45 | | dog-trot | ❌ | 46 | | dog-run | ❌ | 47 | | dog-fetch | ❌ | 48 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/common.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Common functions for the DM Control Suite.""" 16 | 17 | from typing import Dict 18 | 19 | from mujoco_playground._src import mjx_env 20 | 21 | _XML_PATH = mjx_env.ROOT_PATH / "dm_control_suite" / "xmls" / "common" 22 | 23 | 24 | def get_assets() -> Dict[str, bytes]: 25 | assets = {} 26 | for f in _XML_PATH.glob("*"): 27 | if f.is_file(): 28 | assets[f.name] = f.read_bytes() 29 | return assets 30 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/dm_control_suite_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Tests for the DM Control Suite.""" 16 | from absl.testing import absltest 17 | from absl.testing import parameterized 18 | import jax 19 | from jax import numpy as jp 20 | 21 | from mujoco_playground._src import dm_control_suite 22 | 23 | 24 | class TestSuite(parameterized.TestCase): 25 | """Tests for the DM Control Suite.""" 26 | 27 | @parameterized.named_parameters( 28 | {"testcase_name": f"test_can_create_{env_name}", "env_name": env_name} 29 | for env_name in dm_control_suite.ALL_ENVS 30 | ) 31 | def test_can_create_all_environments(self, env_name: str) -> None: 32 | env = dm_control_suite.load(env_name) 33 | state = jax.jit(env.reset)(jax.random.PRNGKey(42)) 34 | state = jax.jit(env.step)(state, jp.zeros(env.action_size)) 35 | self.assertIsNotNone(state) 36 | self.assertEqual(state.obs.shape[0], env.observation_size) 37 | self.assertFalse(jp.isnan(state.data.qpos).any()) 38 | 39 | 40 | if __name__ == "__main__": 41 | absltest.main() 42 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/README.md: -------------------------------------------------------------------------------- 1 | # XML Changes Compared to DM Control 2 | 3 | Several changes are made to DM Control envs for them to be performant on XPU, although changes are kept to a minimum. For all envs, we disable `eulerdamp` and reduce solver `iterations` and `ls_iterations`. For a few envs, we increase the simulator timestep, and set `max_contact_points` and `max_geom_pairs` for contact culling. The full set of changes are listed below: 4 | 5 | * acrobot 6 | * `iterations`="2", `ls_iterations`="4` 7 | * ball_in_cup 8 | * `iterations`="1", `ls_iterations`="4" 9 | * cartpole 10 | * `iterations`="1", `ls_iterations`="4" 11 | * cheetah 12 | * `iterations`="4", `ls_iterations`="8" 13 | * `max_contact_points`=6, `max_geom_pairs`=4 14 | * finger 15 | * `iterations`="2", `ls_iterations`="8" 16 | * `max_contact_points`=4, `max_geom_pairs`=2 17 | * change cylinder contype/conaffinity to 0 18 | * fish 19 | * `iterations`="2", `ls_iterations`="6" 20 | * contacts disabled 21 | * hopper 22 | * `iterations`="4", `ls_iterations`="8" 23 | * `max_contact_points`=6, `max_geom_pairs`=2 24 | * humanoid 25 | * Removed touch sensors 26 | * `timestep`="0.005" compared to 0.0025 in DM Control 27 | * `max_contact_points`=8, `max_geom_pairs`=8 28 | * manipulator 29 | * `timestep`="0.005" from "0.002" in DM Control 30 | * `max_contact_points`=8, `max_geom_pairs`=8 31 | * pendulum 32 | * `timestep`="0.01" compared to "0.02" in DM Control 33 | * `iterations`="4", `ls_iterations`="8" 34 | * point_mass 35 | * `iterations`="1", `ls_iterations`="4" 36 | * reacher 37 | * `timestep`="0.005" from "0.02" in DM Control 38 | * `iterations`="1", `ls_iterations`="6" 39 | * swimmer 40 | * `timestep`="0.003" from "0.002" in DM Control 41 | * `iterations`="4", `ls_iterations`="8" 42 | * geom `contype`/`conaffinity` set to zero, since contacts are disabled 43 | * walker 44 | * `timestep`="0.005", from "0.0025" in DM Control 45 | * `iterations`="2", `ls_iterations`="5" 46 | * `max_contact_points`=4, `max_geom_pairs`=4 47 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/acrobot.xml: -------------------------------------------------------------------------------- 1 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/ball_in_cup.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/cartpole.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/cheetah.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/common/materials.xml: -------------------------------------------------------------------------------- 1 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/common/skybox.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 5 | 6 | 7 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/common/visual.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/finger.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 9 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | 79 | 80 | 81 | 82 | 83 | 84 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/hopper.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/pendulum.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/point_mass.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/reacher.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/swimmer.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | -------------------------------------------------------------------------------- /mujoco_playground/_src/dm_control_suite/xmls/walker.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 58 | 59 | 60 | 61 | 62 | 63 | 64 | 65 | 66 | 67 | 68 | 69 | 70 | 71 | 72 | 73 | 74 | 75 | 76 | 77 | 78 | -------------------------------------------------------------------------------- /mujoco_playground/_src/gait.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | # pylint: disable=line-too-long 16 | """Gait-related utilities. 17 | 18 | References: 19 | - https://arxiv.org/pdf/2201.00206 20 | - https://github.com/google-deepmind/mujoco_mpc/blob/main/mjpc/tasks/quadruped/quadruped.h#L88 21 | - https://colab.research.google.com/drive/1Rwu7fu9Wmko_GygRrHvxb8A_0zRDhhSp?usp=sharing 22 | """ 23 | # pylint: enable=line-too-long 24 | 25 | from typing import Union 26 | 27 | import jax 28 | import jax.numpy as jp 29 | import mujoco 30 | import numpy as np 31 | 32 | 33 | def get_rz( 34 | phi: Union[jax.Array, float], swing_height: Union[jax.Array, float] = 0.08 35 | ) -> jax.Array: 36 | def cubic_bezier_interpolation(y_start, y_end, x): 37 | y_diff = y_end - y_start 38 | bezier = x**3 + 3 * (x**2 * (1 - x)) 39 | return y_start + y_diff * bezier 40 | 41 | x = (phi + jp.pi) / (2 * jp.pi) 42 | stance = cubic_bezier_interpolation(0, swing_height, 2 * x) 43 | swing = cubic_bezier_interpolation(swing_height, 0, 2 * x - 1) 44 | return jp.where(x <= 0.5, stance, swing) 45 | 46 | 47 | # Foot order:"FR", "FL", "RR", "RL". 48 | GAIT_PHASES = { 49 | # trot (diagonals together). 50 | 0: np.array([0, np.pi, np.pi, 0]), 51 | # walk (staggered diagonals). 52 | 1: np.array([0, 0.5 * np.pi, np.pi, 1.5 * np.pi]), 53 | # pace (same side legs together). 54 | 2: np.array([0, np.pi, 0, np.pi]), 55 | # bound (front and back legs together). 56 | 3: np.array([0, 0, np.pi, np.pi]), 57 | # pronk (all legs together). 58 | 4: np.array([0, 0, 0, 0]), 59 | } 60 | 61 | 62 | def draw_joystick_command( 63 | scn, 64 | cmd, 65 | xyz, 66 | theta, 67 | rgba=None, 68 | radius=0.02, 69 | scl=1.0, 70 | ): 71 | if rgba is None: 72 | rgba = [0.2, 0.2, 0.6, 0.3] 73 | scn.ngeom += 1 74 | scn.geoms[scn.ngeom - 1].category = mujoco.mjtCatBit.mjCAT_DECOR 75 | 76 | vx, vy, vtheta = cmd 77 | 78 | angle = theta + vtheta 79 | rotation_matrix = np.array( 80 | [[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]] 81 | ) 82 | 83 | arrow_from = xyz 84 | rotated_velocity = rotation_matrix @ np.array([vx, vy]) 85 | to = np.asarray([rotated_velocity[0], rotated_velocity[1], 0]) 86 | to = to / (np.linalg.norm(to) + 1e-6) 87 | arrow_to = arrow_from + to * scl 88 | 89 | mujoco.mjv_initGeom( 90 | geom=scn.geoms[scn.ngeom - 1], 91 | type=mujoco.mjtGeom.mjGEOM_ARROW.value, 92 | size=np.zeros(3), 93 | pos=np.zeros(3), 94 | mat=np.zeros(9), 95 | rgba=np.asarray(rgba).astype(np.float32), 96 | ) 97 | mujoco.mjv_connector( 98 | geom=scn.geoms[scn.ngeom - 1], 99 | type=mujoco.mjtGeom.mjGEOM_ARROW.value, 100 | width=radius, 101 | from_=arrow_from, 102 | to=arrow_to, 103 | ) 104 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/apollo/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/apollo/constants.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Constants for Apollo.""" 16 | 17 | from etils import epath 18 | 19 | from mujoco_playground._src import mjx_env 20 | 21 | XML_DIR = mjx_env.ROOT_PATH / "locomotion" / "apollo" / "xmls" 22 | 23 | FEET_ONLY_FLAT_TERRAIN_XML = XML_DIR / "scene_mjx_feetonly_flat_terrain.xml" 24 | 25 | 26 | def task_to_xml(task_name: str) -> epath.Path: 27 | return { 28 | "flat_terrain": FEET_ONLY_FLAT_TERRAIN_XML, 29 | }[task_name] 30 | 31 | 32 | FEET_SITES = [ 33 | "l_foot", 34 | "r_foot", 35 | ] 36 | 37 | HAND_SITES = [ 38 | "left_palm", 39 | "right_palm", 40 | ] 41 | 42 | LEFT_FEET_GEOMS = ["collision_l_sole"] 43 | RIGHT_FEET_GEOMS = ["collision_r_sole"] 44 | FEET_GEOMS = LEFT_FEET_GEOMS + RIGHT_FEET_GEOMS 45 | 46 | ROOT_BODY = "torso_link" 47 | 48 | GRAVITY_SENSOR = "upvector" 49 | GLOBAL_LINVEL_SENSOR = "global_linvel" 50 | GLOBAL_ANGVEL_SENSOR = "global_angvel" 51 | LOCAL_LINVEL_SENSOR = "local_linvel" 52 | ACCELEROMETER_SENSOR = "accelerometer" 53 | GYRO_SENSOR = "gyro" 54 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/apollo/xmls/scene_mjx_feetonly_flat_terrain.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 45 | 59 | 60 | 61 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/barkour/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/berkeley_humanoid/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/berkeley_humanoid/berkeley_humanoid_constants.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Constants for Berkeley Humanoid.""" 16 | 17 | from etils import epath 18 | 19 | from mujoco_playground._src import mjx_env 20 | 21 | ROOT_PATH = mjx_env.ROOT_PATH / "locomotion" / "berkeley_humanoid" 22 | FEET_ONLY_FLAT_TERRAIN_XML = ( 23 | ROOT_PATH / "xmls" / "scene_mjx_feetonly_flat_terrain.xml" 24 | ) 25 | FEET_ONLY_ROUGH_TERRAIN_XML = ( 26 | ROOT_PATH / "xmls" / "scene_mjx_feetonly_rough_terrain.xml" 27 | ) 28 | 29 | 30 | def task_to_xml(task_name: str) -> epath.Path: 31 | return { 32 | "flat_terrain": FEET_ONLY_FLAT_TERRAIN_XML, 33 | "rough_terrain": FEET_ONLY_ROUGH_TERRAIN_XML, 34 | }[task_name] 35 | 36 | 37 | FEET_SITES = [ 38 | "l_foot", 39 | "r_foot", 40 | ] 41 | 42 | LEFT_FEET_GEOMS = [ 43 | "l_foot1", 44 | ] 45 | 46 | RIGHT_FEET_GEOMS = [ 47 | "r_foot1", 48 | ] 49 | 50 | FEET_GEOMS = LEFT_FEET_GEOMS + RIGHT_FEET_GEOMS 51 | 52 | FEET_POS_SENSOR = [f"{site}_pos" for site in FEET_SITES] 53 | 54 | ROOT_BODY = "torso" 55 | 56 | GRAVITY_SENSOR = "upvector" 57 | GLOBAL_LINVEL_SENSOR = "global_linvel" 58 | GLOBAL_ANGVEL_SENSOR = "global_angvel" 59 | LOCAL_LINVEL_SENSOR = "local_linvel" 60 | ACCELEROMETER_SENSOR = "accelerometer" 61 | GYRO_SENSOR = "gyro" 62 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/berkeley_humanoid/randomize.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Domain randomization for the Berkeley Humanoid environment.""" 16 | 17 | import jax 18 | from mujoco import mjx 19 | 20 | FLOOR_GEOM_ID = 0 21 | TORSO_BODY_ID = 1 22 | 23 | 24 | def domain_randomize(model: mjx.Model, rng: jax.Array): 25 | @jax.vmap 26 | def rand_dynamics(rng): 27 | # Floor friction: =U(0.4, 1.0). 28 | rng, key = jax.random.split(rng) 29 | geom_friction = model.geom_friction.at[FLOOR_GEOM_ID, 0].set( 30 | jax.random.uniform(key, minval=0.4, maxval=1.0) 31 | ) 32 | 33 | # Scale static friction: *U(0.9, 1.1). 34 | rng, key = jax.random.split(rng) 35 | frictionloss = model.dof_frictionloss[6:] * jax.random.uniform( 36 | key, shape=(12,), minval=0.9, maxval=1.1 37 | ) 38 | dof_frictionloss = model.dof_frictionloss.at[6:].set(frictionloss) 39 | 40 | # Scale armature: *U(1.0, 1.05). 41 | rng, key = jax.random.split(rng) 42 | armature = model.dof_armature[6:] * jax.random.uniform( 43 | key, shape=(12,), minval=1.0, maxval=1.05 44 | ) 45 | dof_armature = model.dof_armature.at[6:].set(armature) 46 | 47 | # Scale all link masses: *U(0.9, 1.1). 48 | rng, key = jax.random.split(rng) 49 | dmass = jax.random.uniform( 50 | key, shape=(model.nbody,), minval=0.9, maxval=1.1 51 | ) 52 | body_mass = model.body_mass.at[:].set(model.body_mass * dmass) 53 | 54 | # Add mass to torso: +U(-1.0, 1.0). 55 | rng, key = jax.random.split(rng) 56 | dmass = jax.random.uniform(key, minval=-1.0, maxval=1.0) 57 | body_mass = body_mass.at[TORSO_BODY_ID].set( 58 | body_mass[TORSO_BODY_ID] + dmass 59 | ) 60 | 61 | # Jitter qpos0: +U(-0.05, 0.05). 62 | rng, key = jax.random.split(rng) 63 | qpos0 = model.qpos0 64 | qpos0 = qpos0.at[7:].set( 65 | qpos0[7:] 66 | + jax.random.uniform(key, shape=(12,), minval=-0.05, maxval=0.05) 67 | ) 68 | 69 | return ( 70 | geom_friction, 71 | dof_frictionloss, 72 | dof_armature, 73 | body_mass, 74 | qpos0, 75 | ) 76 | 77 | ( 78 | friction, 79 | frictionloss, 80 | armature, 81 | body_mass, 82 | qpos0, 83 | ) = rand_dynamics(rng) 84 | 85 | in_axes = jax.tree_util.tree_map(lambda x: None, model) 86 | in_axes = in_axes.tree_replace({ 87 | "geom_friction": 0, 88 | "dof_frictionloss": 0, 89 | "dof_armature": 0, 90 | "body_mass": 0, 91 | "qpos0": 0, 92 | }) 93 | 94 | model = model.tree_replace({ 95 | "geom_friction": friction, 96 | "dof_frictionloss": frictionloss, 97 | "dof_armature": armature, 98 | "body_mass": body_mass, 99 | "qpos0": qpos0, 100 | }) 101 | 102 | return model, in_axes 103 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/berkeley_humanoid/xmls/assets/hfield.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/locomotion/berkeley_humanoid/xmls/assets/hfield.png -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/berkeley_humanoid/xmls/assets/rocky_texture.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/locomotion/berkeley_humanoid/xmls/assets/rocky_texture.png -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/berkeley_humanoid/xmls/scene_mjx_feetonly_flat_terrain.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 19 | 20 | 21 | 22 | 23 | 25 | 26 | 27 | 28 | 37 | 38 | 39 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/berkeley_humanoid/xmls/scene_mjx_feetonly_rough_terrain.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 24 | 25 | 26 | 27 | 36 | 37 | 38 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/g1/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/g1/g1_constants.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Constants for G1.""" 16 | 17 | from etils import epath 18 | 19 | from mujoco_playground._src import mjx_env 20 | 21 | ROOT_PATH = mjx_env.ROOT_PATH / "locomotion" / "g1" 22 | FEET_ONLY_FLAT_TERRAIN_XML = ( 23 | ROOT_PATH / "xmls" / "scene_mjx_feetonly_flat_terrain.xml" 24 | ) 25 | FEET_ONLY_ROUGH_TERRAIN_XML = ( 26 | ROOT_PATH / "xmls" / "scene_mjx_feetonly_rough_terrain.xml" 27 | ) 28 | 29 | 30 | def task_to_xml(task_name: str) -> epath.Path: 31 | return { 32 | "flat_terrain": FEET_ONLY_FLAT_TERRAIN_XML, 33 | "rough_terrain": FEET_ONLY_ROUGH_TERRAIN_XML, 34 | }[task_name] 35 | 36 | 37 | FEET_SITES = [ 38 | "left_foot", 39 | "right_foot", 40 | ] 41 | 42 | HAND_SITES = [ 43 | "left_palm", 44 | "right_palm", 45 | ] 46 | 47 | LEFT_FEET_GEOMS = ["left_foot"] 48 | RIGHT_FEET_GEOMS = ["right_foot"] 49 | FEET_GEOMS = LEFT_FEET_GEOMS + RIGHT_FEET_GEOMS 50 | 51 | ROOT_BODY = "torso_link" 52 | 53 | GRAVITY_SENSOR = "upvector" 54 | GLOBAL_LINVEL_SENSOR = "global_linvel" 55 | GLOBAL_ANGVEL_SENSOR = "global_angvel" 56 | LOCAL_LINVEL_SENSOR = "local_linvel" 57 | ACCELEROMETER_SENSOR = "accelerometer" 58 | GYRO_SENSOR = "gyro" 59 | 60 | RESTRICTED_JOINT_RANGE = ( 61 | # Left leg. 62 | (-1.57, 1.57), 63 | (-0.5, 0.5), 64 | (-0.7, 0.7), 65 | (0, 1.57), 66 | (-0.4, 0.4), 67 | (-0.2, 0.2), 68 | # Right leg. 69 | (-1.57, 1.57), 70 | (-0.5, 0.5), 71 | (-0.7, 0.7), 72 | (0, 1.57), 73 | (-0.4, 0.4), 74 | (-0.2, 0.2), 75 | # Waist. 76 | (-2.618, 2.618), 77 | (-0.52, 0.52), 78 | (-0.52, 0.52), 79 | # Left shoulder. 80 | (-3.0892, 2.6704), 81 | (-1.5882, 2.2515), 82 | (-2.618, 2.618), 83 | (-1.0472, 2.0944), 84 | (-1.97222, 1.97222), 85 | (-1.61443, 1.61443), 86 | (-1.61443, 1.61443), 87 | # Right shoulder. 88 | (-3.0892, 2.6704), 89 | (-2.2515, 1.5882), 90 | (-2.618, 2.618), 91 | (-1.0472, 2.0944), 92 | (-1.97222, 1.97222), 93 | (-1.61443, 1.61443), 94 | (-1.61443, 1.61443), 95 | ) 96 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/g1/randomize.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Utilities for randomization.""" 16 | import jax 17 | from mujoco import mjx 18 | 19 | FLOOR_GEOM_ID = 0 20 | TORSO_BODY_ID = 16 21 | 22 | 23 | def domain_randomize(model: mjx.Model, rng: jax.Array): 24 | @jax.vmap 25 | def rand_dynamics(rng): 26 | # Floor / foot friction: =U(0.4, 1.0). 27 | rng, key = jax.random.split(rng) 28 | friction = jax.random.uniform(key, minval=0.4, maxval=1.0) 29 | pair_friction = model.pair_friction.at[0:2, 0:2].set(friction) 30 | 31 | # Scale static friction: *U(0.9, 1.1). 32 | rng, key = jax.random.split(rng) 33 | frictionloss = model.dof_frictionloss[6:] * jax.random.uniform( 34 | key, shape=(29,), minval=0.5, maxval=2.0 35 | ) 36 | dof_frictionloss = model.dof_frictionloss.at[6:].set(frictionloss) 37 | 38 | # Scale armature: *U(1.0, 1.05). 39 | rng, key = jax.random.split(rng) 40 | armature = model.dof_armature[6:] * jax.random.uniform( 41 | key, shape=(29,), minval=1.0, maxval=1.05 42 | ) 43 | dof_armature = model.dof_armature.at[6:].set(armature) 44 | 45 | # Scale all link masses: *U(0.9, 1.1). 46 | rng, key = jax.random.split(rng) 47 | dmass = jax.random.uniform( 48 | key, shape=(model.nbody,), minval=0.9, maxval=1.1 49 | ) 50 | body_mass = model.body_mass.at[:].set(model.body_mass * dmass) 51 | 52 | # Add mass to torso: +U(-1.0, 1.0). 53 | rng, key = jax.random.split(rng) 54 | dmass = jax.random.uniform(key, minval=-1.0, maxval=1.0) 55 | body_mass = body_mass.at[TORSO_BODY_ID].set( 56 | body_mass[TORSO_BODY_ID] + dmass 57 | ) 58 | 59 | # Jitter qpos0: +U(-0.05, 0.05). 60 | rng, key = jax.random.split(rng) 61 | qpos0 = model.qpos0 62 | qpos0 = qpos0.at[7:].set( 63 | qpos0[7:] 64 | + jax.random.uniform(key, shape=(29,), minval=-0.05, maxval=0.05) 65 | ) 66 | 67 | return ( 68 | pair_friction, 69 | dof_frictionloss, 70 | dof_armature, 71 | body_mass, 72 | qpos0, 73 | ) 74 | 75 | ( 76 | pair_friction, 77 | frictionloss, 78 | armature, 79 | body_mass, 80 | qpos0, 81 | ) = rand_dynamics(rng) 82 | 83 | in_axes = jax.tree_util.tree_map(lambda x: None, model) 84 | in_axes = in_axes.tree_replace({ 85 | "pair_friction": 0, 86 | "dof_frictionloss": 0, 87 | "dof_armature": 0, 88 | "body_mass": 0, 89 | "qpos0": 0, 90 | }) 91 | 92 | model = model.tree_replace({ 93 | "pair_friction": pair_friction, 94 | "dof_frictionloss": frictionloss, 95 | "dof_armature": armature, 96 | "body_mass": body_mass, 97 | "qpos0": qpos0, 98 | }) 99 | 100 | return model, in_axes 101 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/g1/xmls/assets/hfield.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/locomotion/g1/xmls/assets/hfield.png -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/g1/xmls/assets/rocky_texture.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/locomotion/g1/xmls/assets/rocky_texture.png -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/g1/xmls/scene_mjx_feetonly.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/g1/xmls/scene_mjx_feetonly_flat_terrain.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 44 | 61 | 62 | 63 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/g1/xmls/scene_mjx_feetonly_rough_terrain.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 44 | 61 | 62 | 63 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/go1/README.md: -------------------------------------------------------------------------------- 1 | # Unitree Go1 environments 2 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/go1/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/go1/go1_constants.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Defines Unitree Go1 quadruped constants.""" 16 | 17 | from etils import epath 18 | 19 | from mujoco_playground._src import mjx_env 20 | 21 | ROOT_PATH = mjx_env.ROOT_PATH / "locomotion" / "go1" 22 | FEET_ONLY_FLAT_TERRAIN_XML = ( 23 | ROOT_PATH / "xmls" / "scene_mjx_feetonly_flat_terrain.xml" 24 | ) 25 | FEET_ONLY_ROUGH_TERRAIN_XML = ( 26 | ROOT_PATH / "xmls" / "scene_mjx_feetonly_rough_terrain.xml" 27 | ) 28 | FULL_FLAT_TERRAIN_XML = ROOT_PATH / "xmls" / "scene_mjx_flat_terrain.xml" 29 | FULL_COLLISIONS_FLAT_TERRAIN_XML = ( 30 | ROOT_PATH / "xmls" / "scene_mjx_fullcollisions_flat_terrain.xml" 31 | ) 32 | 33 | 34 | def task_to_xml(task_name: str) -> epath.Path: 35 | return { 36 | "flat_terrain": FEET_ONLY_FLAT_TERRAIN_XML, 37 | "rough_terrain": FEET_ONLY_ROUGH_TERRAIN_XML, 38 | }[task_name] 39 | 40 | 41 | FEET_SITES = [ 42 | "FR", 43 | "FL", 44 | "RR", 45 | "RL", 46 | ] 47 | 48 | FEET_GEOMS = [ 49 | "FR", 50 | "FL", 51 | "RR", 52 | "RL", 53 | ] 54 | 55 | FEET_POS_SENSOR = [f"{site}_pos" for site in FEET_SITES] 56 | 57 | ROOT_BODY = "trunk" 58 | 59 | UPVECTOR_SENSOR = "upvector" 60 | GLOBAL_LINVEL_SENSOR = "global_linvel" 61 | GLOBAL_ANGVEL_SENSOR = "global_angvel" 62 | LOCAL_LINVEL_SENSOR = "local_linvel" 63 | ACCELEROMETER_SENSOR = "accelerometer" 64 | GYRO_SENSOR = "gyro" 65 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/go1/randomize.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Domain randomization for the Go1 environment.""" 16 | 17 | import jax 18 | from mujoco import mjx 19 | 20 | FLOOR_GEOM_ID = 0 21 | TORSO_BODY_ID = 1 22 | 23 | 24 | def domain_randomize(model: mjx.Model, rng: jax.Array): 25 | @jax.vmap 26 | def rand_dynamics(rng): 27 | # Floor friction: =U(0.4, 1.0). 28 | rng, key = jax.random.split(rng) 29 | geom_friction = model.geom_friction.at[FLOOR_GEOM_ID, 0].set( 30 | jax.random.uniform(key, minval=0.4, maxval=1.0) 31 | ) 32 | 33 | # Scale static friction: *U(0.9, 1.1). 34 | rng, key = jax.random.split(rng) 35 | frictionloss = model.dof_frictionloss[6:] * jax.random.uniform( 36 | key, shape=(12,), minval=0.9, maxval=1.1 37 | ) 38 | dof_frictionloss = model.dof_frictionloss.at[6:].set(frictionloss) 39 | 40 | # Scale armature: *U(1.0, 1.05). 41 | rng, key = jax.random.split(rng) 42 | armature = model.dof_armature[6:] * jax.random.uniform( 43 | key, shape=(12,), minval=1.0, maxval=1.05 44 | ) 45 | dof_armature = model.dof_armature.at[6:].set(armature) 46 | 47 | # Jitter center of mass positiion: +U(-0.05, 0.05). 48 | rng, key = jax.random.split(rng) 49 | dpos = jax.random.uniform(key, (3,), minval=-0.05, maxval=0.05) 50 | body_ipos = model.body_ipos.at[TORSO_BODY_ID].set( 51 | model.body_ipos[TORSO_BODY_ID] + dpos 52 | ) 53 | 54 | # Scale all link masses: *U(0.9, 1.1). 55 | rng, key = jax.random.split(rng) 56 | dmass = jax.random.uniform( 57 | key, shape=(model.nbody,), minval=0.9, maxval=1.1 58 | ) 59 | body_mass = model.body_mass.at[:].set(model.body_mass * dmass) 60 | 61 | # Add mass to torso: +U(-1.0, 1.0). 62 | rng, key = jax.random.split(rng) 63 | dmass = jax.random.uniform(key, minval=-1.0, maxval=1.0) 64 | body_mass = body_mass.at[TORSO_BODY_ID].set( 65 | body_mass[TORSO_BODY_ID] + dmass 66 | ) 67 | 68 | # Jitter qpos0: +U(-0.05, 0.05). 69 | rng, key = jax.random.split(rng) 70 | qpos0 = model.qpos0 71 | qpos0 = qpos0.at[7:].set( 72 | qpos0[7:] 73 | + jax.random.uniform(key, shape=(12,), minval=-0.05, maxval=0.05) 74 | ) 75 | 76 | return ( 77 | geom_friction, 78 | body_ipos, 79 | body_mass, 80 | qpos0, 81 | dof_frictionloss, 82 | dof_armature, 83 | ) 84 | 85 | ( 86 | friction, 87 | body_ipos, 88 | body_mass, 89 | qpos0, 90 | dof_frictionloss, 91 | dof_armature, 92 | ) = rand_dynamics(rng) 93 | 94 | in_axes = jax.tree_util.tree_map(lambda x: None, model) 95 | in_axes = in_axes.tree_replace({ 96 | "geom_friction": 0, 97 | "body_ipos": 0, 98 | "body_mass": 0, 99 | "qpos0": 0, 100 | "dof_frictionloss": 0, 101 | "dof_armature": 0, 102 | }) 103 | 104 | model = model.tree_replace({ 105 | "geom_friction": friction, 106 | "body_ipos": body_ipos, 107 | "body_mass": body_mass, 108 | "qpos0": qpos0, 109 | "dof_frictionloss": dof_frictionloss, 110 | "dof_armature": dof_armature, 111 | }) 112 | 113 | return model, in_axes 114 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/go1/xmls/assets/hfield.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/locomotion/go1/xmls/assets/hfield.png -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/go1/xmls/assets/rocky_texture.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/locomotion/go1/xmls/assets/rocky_texture.png -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/go1/xmls/scene_mjx_feetonly_bowl.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 29 | 32 | 33 | 34 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/go1/xmls/scene_mjx_feetonly_flat_terrain.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 19 | 20 | 21 | 22 | 23 | 25 | 26 | 27 | 28 | 36 | 38 | 45 | 53 | 56 | 57 | 58 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/go1/xmls/scene_mjx_feetonly_rough_terrain.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 24 | 25 | 26 | 27 | 35 | 37 | 38 | 39 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/go1/xmls/scene_mjx_feetonly_stairs.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 56 | 57 | 59 | 60 | 61 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/go1/xmls/scene_mjx_flat_terrain.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 35 | 37 | 40 | 43 | 46 | 47 | 48 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/go1/xmls/scene_mjx_fullcollisions_flat_terrain.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 35 | 37 | 40 | 43 | 46 | 47 | 48 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/h1/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/h1/base.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """H1 base classes.""" 16 | 17 | from typing import Any, Dict, Optional, Union 18 | 19 | from etils import epath 20 | import jax 21 | from ml_collections import config_dict 22 | import mujoco 23 | from mujoco import mjx 24 | 25 | from mujoco_playground._src import mjx_env 26 | from mujoco_playground._src.locomotion.h1 import h1_constants as consts 27 | 28 | 29 | def get_assets() -> Dict[str, bytes]: 30 | assets = {} 31 | mjx_env.update_assets(assets, consts.ROOT_PATH, "*.xml") 32 | path = mjx_env.MENAGERIE_PATH / "unitree_h1" 33 | mjx_env.update_assets(assets, path, "*.xml") 34 | mjx_env.update_assets(assets, path / "assets") 35 | return assets 36 | 37 | 38 | class H1Env(mjx_env.MjxEnv): 39 | """Base class for H1 environments.""" 40 | 41 | def __init__( 42 | self, 43 | xml_path: str, 44 | config: config_dict.ConfigDict, 45 | config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None, 46 | ) -> None: 47 | super().__init__(config, config_overrides) 48 | 49 | self._mj_model = mujoco.MjModel.from_xml_string( 50 | epath.Path(xml_path).read_text(), assets=get_assets() 51 | ) 52 | self._mj_model.opt.timestep = self._config.sim_dt 53 | 54 | # Increase offscreen framebuffer size to render at higher resolutions. 55 | # TODO(kevin): Consider moving this somewhere else. 56 | self._mj_model.vis.global_.offwidth = 3840 57 | self._mj_model.vis.global_.offheight = 2160 58 | 59 | self._mjx_model = mjx.put_model(self._mj_model) 60 | self._xml_path = xml_path 61 | 62 | # Sensor readings. 63 | 64 | def get_gravity(self, data: mjx.Data) -> jax.Array: 65 | """Return the gravity vector in the world frame.""" 66 | return mjx_env.get_sensor_data(self.mj_model, data, consts.GRAVITY_SENSOR) 67 | 68 | def get_global_linvel(self, data: mjx.Data) -> jax.Array: 69 | """Return the linear velocity of the robot in the world frame.""" 70 | return mjx_env.get_sensor_data( 71 | self.mj_model, data, consts.GLOBAL_LINVEL_SENSOR 72 | ) 73 | 74 | def get_global_angvel(self, data: mjx.Data) -> jax.Array: 75 | """Return the angular velocity of the robot in the world frame.""" 76 | return mjx_env.get_sensor_data( 77 | self.mj_model, data, consts.GLOBAL_ANGVEL_SENSOR 78 | ) 79 | 80 | def get_local_linvel(self, data: mjx.Data) -> jax.Array: 81 | """Return the linear velocity of the robot in the local frame.""" 82 | return mjx_env.get_sensor_data( 83 | self.mj_model, data, consts.LOCAL_LINVEL_SENSOR 84 | ) 85 | 86 | def get_accelerometer(self, data: mjx.Data) -> jax.Array: 87 | """Return the accelerometer readings in the local frame.""" 88 | return mjx_env.get_sensor_data( 89 | self.mj_model, data, consts.ACCELEROMETER_SENSOR 90 | ) 91 | 92 | def get_gyro(self, data: mjx.Data) -> jax.Array: 93 | """Return the gyroscope readings in the local frame.""" 94 | return mjx_env.get_sensor_data(self.mj_model, data, consts.GYRO_SENSOR) 95 | 96 | # Accessors. 97 | 98 | @property 99 | def xml_path(self) -> str: 100 | return self._xml_path 101 | 102 | @property 103 | def action_size(self) -> int: 104 | return self._mjx_model.nu 105 | 106 | @property 107 | def mj_model(self) -> mujoco.MjModel: 108 | return self._mj_model 109 | 110 | @property 111 | def mjx_model(self) -> mjx.Model: 112 | return self._mjx_model 113 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/h1/h1_constants.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """H1 constants.""" 16 | 17 | from mujoco_playground._src import mjx_env 18 | 19 | ROOT_PATH = mjx_env.ROOT_PATH / "locomotion" / "h1" / "xmls" 20 | FEET_ONLY_XML = ROOT_PATH / "scene_mjx_feetonly.xml" 21 | 22 | FEET_SITES = [ 23 | "left_foot", 24 | "right_foot", 25 | ] 26 | 27 | LEFT_FEET_GEOMS = [ 28 | "left_foot1", 29 | "left_foot2", 30 | "left_foot3", 31 | ] 32 | RIGHT_FEET_GEOMS = [ 33 | "right_foot1", 34 | "right_foot2", 35 | "right_foot3", 36 | ] 37 | 38 | ROOT_BODY = "torso_link" 39 | 40 | GRAVITY_SENSOR = "upvector" 41 | GLOBAL_LINVEL_SENSOR = "global_linvel" 42 | GLOBAL_ANGVEL_SENSOR = "global_angvel" 43 | LOCAL_LINVEL_SENSOR = "local_linvel" 44 | ACCELEROMETER_SENSOR = "accelerometer" 45 | GYRO_SENSOR = "gyro" 46 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/h1/xmls/scene_mjx_feetonly.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 41 | 42 | 43 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/locomotion_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Tests for the locomotion environments.""" 16 | from absl.testing import absltest 17 | from absl.testing import parameterized 18 | import jax 19 | import jax.numpy as jp 20 | 21 | from mujoco_playground._src import locomotion 22 | 23 | 24 | class TestSuite(parameterized.TestCase): 25 | """Tests for the locomotion environments.""" 26 | 27 | @parameterized.named_parameters( 28 | {"testcase_name": f"test_can_create_{env_name}", "env_name": env_name} 29 | for env_name in locomotion.ALL_ENVS 30 | ) 31 | def test_can_create_all_environments(self, env_name: str) -> None: 32 | env = locomotion.load(env_name) 33 | state = jax.jit(env.reset)(jax.random.PRNGKey(42)) 34 | state = jax.jit(env.step)(state, jp.zeros(env.action_size)) 35 | self.assertIsNotNone(state) 36 | obs_shape = jax.tree_util.tree_map(lambda x: x.shape, state.obs) 37 | obs_shape = obs_shape[0] if isinstance(obs_shape, tuple) else obs_shape 38 | self.assertEqual(obs_shape, env.observation_size) 39 | self.assertFalse(jp.isnan(state.data.qpos).any()) 40 | 41 | 42 | if __name__ == "__main__": 43 | absltest.main() 44 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/op3/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/op3/op3_constants.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Constants for OP3.""" 16 | 17 | from mujoco_playground._src import mjx_env 18 | 19 | ROOT_PATH = mjx_env.ROOT_PATH / "locomotion" / "op3" 20 | FEET_ONLY_XML = ROOT_PATH / "xmls" / "scene_mjx_feetonly.xml" 21 | 22 | FEET_SITES = [ 23 | "left_foot", 24 | "right_foot", 25 | ] 26 | 27 | LEFT_FEET_GEOMS = [ 28 | "l_foot1", 29 | "l_foot2", 30 | ] 31 | RIGHT_FEET_GEOMS = [ 32 | "r_foot1", 33 | "r_foot2", 34 | ] 35 | 36 | ROOT_BODY = "body_link" 37 | 38 | GRAVITY_SENSOR = "upvector" 39 | GLOBAL_LINVEL_SENSOR = "global_linvel" 40 | GLOBAL_ANGVEL_SENSOR = "global_angvel" 41 | LOCAL_LINVEL_SENSOR = "local_linvel" 42 | ACCELEROMETER_SENSOR = "accelerometer" 43 | GYRO_SENSOR = "gyro" 44 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/op3/xmls/scene_mjx_feetonly.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/spot/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/spot/spot_constants.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Defines Spot quadruped constants.""" 16 | 17 | from etils import epath 18 | 19 | from mujoco_playground._src import mjx_env 20 | 21 | ROOT_PATH = mjx_env.ROOT_PATH / "locomotion" / "spot" 22 | FEET_ONLY_FLAT_TERRAIN_XML = ( 23 | ROOT_PATH / "xmls" / "scene_mjx_feetonly_flat_terrain.xml" 24 | ) 25 | FULL_FLAT_TERRAIN_XML = ROOT_PATH / "xmls" / "scene_mjx_flat_terrain.xml" 26 | 27 | 28 | def task_to_xml(task_name: str) -> epath.Path: 29 | return { 30 | "flat_terrain": FEET_ONLY_FLAT_TERRAIN_XML, 31 | }[task_name] 32 | 33 | 34 | FEET_SITES = [ 35 | "FL", 36 | "FR", 37 | "HL", 38 | "HR", 39 | ] 40 | 41 | FEET_GEOMS = [ 42 | "FL", 43 | "FR", 44 | "HL", 45 | "HR", 46 | ] 47 | 48 | FEET_POS_SENSOR = [f"{site}_pos" for site in FEET_SITES] 49 | 50 | ROOT_BODY = "body" 51 | 52 | GRAVITY_SENSOR = "upvector" 53 | GLOBAL_LINVEL_SENSOR = "global_linvel" 54 | GLOBAL_ANGVEL_SENSOR = "global_angvel" 55 | LOCAL_LINVEL_SENSOR = "local_linvel" 56 | ACCELEROMETER_SENSOR = "accelerometer" 57 | GYRO_SENSOR = "gyro" 58 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/spot/xmls/scene_mjx_feetonly_flat_terrain.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 19 | 20 | 21 | 22 | 23 | 25 | 26 | 27 | 28 | 40 | 41 | 42 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/spot/xmls/scene_mjx_flat_terrain.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 19 | 20 | 21 | 22 | 23 | 25 | 26 | 27 | 28 | 40 | 41 | 42 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/t1/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/t1/base.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Base classes for T1.""" 16 | 17 | from typing import Any, Dict, Optional, Union 18 | 19 | from etils import epath 20 | import jax 21 | from ml_collections import config_dict 22 | import mujoco 23 | from mujoco import mjx 24 | 25 | from mujoco_playground._src import mjx_env 26 | from mujoco_playground._src.locomotion.t1 import t1_constants as consts 27 | 28 | 29 | def get_assets() -> Dict[str, bytes]: 30 | assets = {} 31 | mjx_env.update_assets(assets, consts.ROOT_PATH / "xmls", "*.xml") 32 | mjx_env.update_assets(assets, consts.ROOT_PATH / "xmls" / "assets") 33 | path = mjx_env.MENAGERIE_PATH / "booster_t1" 34 | mjx_env.update_assets(assets, path, "*.xml") 35 | mjx_env.update_assets(assets, path / "assets") 36 | return assets 37 | 38 | 39 | class T1Env(mjx_env.MjxEnv): 40 | """Base class for T1 environments.""" 41 | 42 | def __init__( 43 | self, 44 | xml_path: str, 45 | config: config_dict.ConfigDict, 46 | config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None, 47 | ) -> None: 48 | super().__init__(config, config_overrides) 49 | 50 | self._mj_model = mujoco.MjModel.from_xml_string( 51 | epath.Path(xml_path).read_text(), assets=get_assets() 52 | ) 53 | self._mj_model.opt.timestep = self.sim_dt 54 | 55 | self._mj_model.vis.global_.offwidth = 3840 56 | self._mj_model.vis.global_.offheight = 2160 57 | 58 | self._mjx_model = mjx.put_model(self._mj_model) 59 | self._xml_path = xml_path 60 | 61 | # Sensor readings. 62 | 63 | def get_gravity(self, data: mjx.Data) -> jax.Array: 64 | """Return the gravity vector in the world frame.""" 65 | return mjx_env.get_sensor_data( 66 | self.mj_model, data, f"{consts.GRAVITY_SENSOR}" 67 | ) 68 | 69 | def get_global_linvel(self, data: mjx.Data) -> jax.Array: 70 | """Return the linear velocity of the robot in the world frame.""" 71 | return mjx_env.get_sensor_data( 72 | self.mj_model, data, f"{consts.GLOBAL_LINVEL_SENSOR}" 73 | ) 74 | 75 | def get_global_angvel(self, data: mjx.Data) -> jax.Array: 76 | """Return the angular velocity of the robot in the world frame.""" 77 | return mjx_env.get_sensor_data( 78 | self.mj_model, data, f"{consts.GLOBAL_ANGVEL_SENSOR}" 79 | ) 80 | 81 | def get_local_linvel(self, data: mjx.Data) -> jax.Array: 82 | """Return the linear velocity of the robot in the local frame.""" 83 | return mjx_env.get_sensor_data( 84 | self.mj_model, data, f"{consts.LOCAL_LINVEL_SENSOR}" 85 | ) 86 | 87 | def get_accelerometer(self, data: mjx.Data) -> jax.Array: 88 | """Return the accelerometer readings in the local frame.""" 89 | return mjx_env.get_sensor_data( 90 | self.mj_model, data, f"{consts.ACCELEROMETER_SENSOR}" 91 | ) 92 | 93 | def get_gyro(self, data: mjx.Data) -> jax.Array: 94 | """Return the gyroscope readings in the local frame.""" 95 | return mjx_env.get_sensor_data(self.mj_model, data, f"{consts.GYRO_SENSOR}") 96 | 97 | # Accessors. 98 | 99 | @property 100 | def xml_path(self) -> str: 101 | return self._xml_path 102 | 103 | @property 104 | def action_size(self) -> int: 105 | return self._mjx_model.nu 106 | 107 | @property 108 | def mj_model(self) -> mujoco.MjModel: 109 | return self._mj_model 110 | 111 | @property 112 | def mjx_model(self) -> mjx.Model: 113 | return self._mjx_model 114 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/t1/t1_constants.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Constants for T1.""" 16 | 17 | from etils import epath 18 | 19 | from mujoco_playground._src import mjx_env 20 | 21 | ROOT_PATH = mjx_env.ROOT_PATH / "locomotion" / "t1" 22 | FEET_ONLY_FLAT_TERRAIN_XML = ( 23 | ROOT_PATH / "xmls" / "scene_mjx_feetonly_flat_terrain.xml" 24 | ) 25 | FEET_ONLY_ROUGH_TERRAIN_XML = ( 26 | ROOT_PATH / "xmls" / "scene_mjx_feetonly_rough_terrain.xml" 27 | ) 28 | 29 | 30 | def task_to_xml(task_name: str) -> epath.Path: 31 | return { 32 | "flat_terrain": FEET_ONLY_FLAT_TERRAIN_XML, 33 | "rough_terrain": FEET_ONLY_ROUGH_TERRAIN_XML, 34 | }[task_name] 35 | 36 | 37 | FEET_SITES = [ 38 | "left_foot", 39 | "right_foot", 40 | ] 41 | 42 | HAND_SITES = [ 43 | "left_hand", 44 | "right_hand", 45 | ] 46 | 47 | LEFT_FEET_GEOMS = [f"left_foot_{i}" for i in range(1, 5)] 48 | RIGHT_FEET_GEOMS = [f"right_foot_{i}" for i in range(1, 5)] 49 | FEET_GEOMS = LEFT_FEET_GEOMS + RIGHT_FEET_GEOMS 50 | 51 | ROOT_BODY = "Trunk" 52 | 53 | GRAVITY_SENSOR = "upvector" 54 | GLOBAL_LINVEL_SENSOR = "global_linvel" 55 | GLOBAL_ANGVEL_SENSOR = "global_angvel" 56 | LOCAL_LINVEL_SENSOR = "local_linvel" 57 | ACCELEROMETER_SENSOR = "accelerometer" 58 | GYRO_SENSOR = "gyro" 59 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/t1/xmls/assets/hfield.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/locomotion/t1/xmls/assets/hfield.png -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/t1/xmls/assets/rocky_texture.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/locomotion/t1/xmls/assets/rocky_texture.png -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/t1/xmls/scene_mjx_feetonly_flat_terrain.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 46 | 47 | 48 | -------------------------------------------------------------------------------- /mujoco_playground/_src/locomotion/t1/xmls/scene_mjx_feetonly_rough_terrain.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 45 | 46 | 47 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/aloha/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/aloha/aloha_constants.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Constants for ALOHA.""" 16 | 17 | from mujoco_playground._src import mjx_env 18 | 19 | XML_PATH = mjx_env.ROOT_PATH / "manipulation" / "aloha" / "xmls" 20 | 21 | ARM_JOINTS = [ 22 | "left/waist", 23 | "left/shoulder", 24 | "left/elbow", 25 | "left/forearm_roll", 26 | "left/wrist_angle", 27 | "left/wrist_rotate", 28 | "right/waist", 29 | "right/shoulder", 30 | "right/elbow", 31 | "right/forearm_roll", 32 | "right/wrist_angle", 33 | "right/wrist_rotate", 34 | ] 35 | 36 | FINGER_GEOMS = [ 37 | "left/left_finger_top", 38 | "left/left_finger_bottom", 39 | "left/right_finger_top", 40 | "left/right_finger_bottom", 41 | "right/left_finger_top", 42 | "right/left_finger_bottom", 43 | "right/right_finger_top", 44 | "right/right_finger_bottom", 45 | ] 46 | 47 | FINGER_JOINTS = [ 48 | "left/left_finger", 49 | "left/right_finger", 50 | "right/left_finger", 51 | "right/right_finger", 52 | ] 53 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/aloha/base.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Base class for ALOHA.""" 16 | 17 | from typing import Any, Dict, Optional, Union 18 | 19 | from etils import epath 20 | import jax.numpy as jp 21 | from ml_collections import config_dict 22 | import mujoco 23 | from mujoco import mjx 24 | import numpy as np 25 | 26 | from mujoco_playground._src import collision 27 | from mujoco_playground._src import mjx_env 28 | from mujoco_playground._src.manipulation.aloha import aloha_constants as consts 29 | 30 | 31 | def get_assets() -> Dict[str, bytes]: 32 | """Returns a dictionary of all assets used by the environment.""" 33 | assets = {} 34 | path = mjx_env.MENAGERIE_PATH / "aloha" 35 | mjx_env.update_assets(assets, path, "*.xml") 36 | mjx_env.update_assets(assets, path / "assets") 37 | path = mjx_env.ROOT_PATH / "manipulation" / "aloha" / "xmls" 38 | mjx_env.update_assets(assets, path, "*.xml") 39 | mjx_env.update_assets(assets, path / "assets") 40 | return assets 41 | 42 | 43 | class AlohaEnv(mjx_env.MjxEnv): 44 | """Base class for ALOHA environments.""" 45 | 46 | def __init__( 47 | self, 48 | xml_path: str, 49 | config: config_dict.ConfigDict, 50 | config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None, 51 | ) -> None: 52 | super().__init__(config, config_overrides) 53 | 54 | self._mj_model = mujoco.MjModel.from_xml_string( 55 | epath.Path(xml_path).read_text(), assets=get_assets() 56 | ) 57 | self._mj_model.opt.timestep = self._config.sim_dt 58 | 59 | self._mj_model.vis.global_.offwidth = 3840 60 | self._mj_model.vis.global_.offheight = 2160 61 | 62 | self._mjx_model = mjx.put_model(self._mj_model) 63 | self._xml_path = xml_path 64 | 65 | def _post_init_aloha(self, keyframe: str = "home"): 66 | """Initializes helpful robot properties.""" 67 | self._left_gripper_site = self._mj_model.site("left/gripper").id 68 | self._right_gripper_site = self._mj_model.site("right/gripper").id 69 | self._table_geom = self._mj_model.geom("table").id 70 | self._finger_geoms = [ 71 | self._mj_model.geom(geom_id).id for geom_id in consts.FINGER_GEOMS 72 | ] 73 | self._init_q = jp.array(self._mj_model.keyframe(keyframe).qpos) 74 | self._init_ctrl = jp.array(self._mj_model.keyframe(keyframe).ctrl) 75 | self._lowers, self._uppers = self.mj_model.actuator_ctrlrange.T 76 | arm_joint_ids = [self._mj_model.joint(j).id for j in consts.ARM_JOINTS] 77 | self._arm_qadr = jp.array( 78 | [self._mj_model.jnt_qposadr[joint_id] for joint_id in arm_joint_ids] 79 | ) 80 | self._finger_qposadr = np.array([ 81 | self._mj_model.jnt_qposadr[self._mj_model.joint(j).id] 82 | for j in consts.FINGER_JOINTS 83 | ]) 84 | 85 | @property 86 | def xml_path(self) -> str: 87 | return self._xml_path 88 | 89 | @property 90 | def action_size(self) -> int: 91 | return self._mjx_model.nu 92 | 93 | @property 94 | def mj_model(self) -> mujoco.MjModel: 95 | return self._mj_model 96 | 97 | @property 98 | def mjx_model(self) -> mjx.Model: 99 | return self._mjx_model 100 | 101 | def hand_table_collision(self, data) -> jp.ndarray: 102 | # Check for collisions with the floor. 103 | hand_table_collisions = [ 104 | collision.geoms_colliding(data, self._table_geom, g) 105 | for g in self._finger_geoms 106 | ] 107 | return (sum(hand_table_collisions) > 0).astype(float) 108 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/aloha/xmls/assets/socket_blue.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/aloha/xmls/assets/socket_blue.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/aloha/xmls/assets/socket_white.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/aloha/xmls/assets/socket_white.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/aloha/xmls/joint_position_actuators.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/aloha/xmls/mjx_hand_over.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 38 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/aloha/xmls/mjx_single_peg_insertion.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 43 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda/xmls/mjx_cabinet.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 10 | 12 | 13 | 14 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 26 | 29 | 30 | 31 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda/xmls/mjx_scene.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda/xmls/mjx_single_cube.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 19 | 22 | 25 | 26 | 27 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda/xmls/mjx_single_cube_camera.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 57 | 60 | 61 | 62 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileback.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileback.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileback_light.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileback_light.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/filedown.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/filedown.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/filedown_light.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/filedown_light.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/filefront.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/filefront.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/filefront_light.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/filefront_light.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileleft.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileleft.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileleft_light.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileleft_light.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileright.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileright.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileright_light.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileright_light.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileup.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileup.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileup_light.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/assets/fileup_light.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/franka_emika_panda_robotiq/xmls/scene_panda_robotiq_cube.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 28 | 29 | 32 | 33 | 34 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 50 | 51 | 52 | 54 | 55 | 56 | 57 | 58 | 62 | 63 | 64 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/README.md: -------------------------------------------------------------------------------- 1 | # LEAP Hand environments 2 | 3 | ## Hardware 4 | 5 | * [Dynamixel XL330-M288-T](https://emanual.robotis.com/docs/en/dxl/x/xl330-m288/) 6 | * [PD gain table](https://emanual.robotis.com/docs/en/dxl/x/xc330-m288/#position-pid-gain80-82-84) 7 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/leap_hand_constants.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Constants for leap hand.""" 16 | 17 | from mujoco_playground._src import mjx_env 18 | 19 | ROOT_PATH = mjx_env.ROOT_PATH / "manipulation" / "leap_hand" 20 | CUBE_XML = ROOT_PATH / "xmls" / "scene_mjx_cube.xml" 21 | 22 | NQ = 16 23 | NV = 16 24 | NU = 16 25 | 26 | JOINT_NAMES = [ 27 | # index 28 | "if_mcp", 29 | "if_rot", 30 | "if_pip", 31 | "if_dip", 32 | # middle 33 | "mf_mcp", 34 | "mf_rot", 35 | "mf_pip", 36 | "mf_dip", 37 | # ring 38 | "rf_mcp", 39 | "rf_rot", 40 | "rf_pip", 41 | "rf_dip", 42 | # thumb 43 | "th_cmc", 44 | "th_axl", 45 | "th_mcp", 46 | "th_ipl", 47 | ] 48 | 49 | ACTUATOR_NAMES = [ 50 | # index 51 | "if_mcp_act", 52 | "if_rot_act", 53 | "if_pip_act", 54 | "if_dip_act", 55 | # middle 56 | "mf_mcp_act", 57 | "mf_rot_act", 58 | "mf_pip_act", 59 | "mf_dip_act", 60 | # ring 61 | "rf_mcp_act", 62 | "rf_rot_act", 63 | "rf_pip_act", 64 | "rf_dip_act", 65 | # thumb 66 | "th_cmc_act", 67 | "th_axl_act", 68 | "th_mcp_act", 69 | "th_ipl_act", 70 | ] 71 | 72 | FINGERTIP_NAMES = [ 73 | "th_tip", 74 | "if_tip", 75 | "mf_tip", 76 | "rf_tip", 77 | ] 78 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/meshes/dex_cube.obj: -------------------------------------------------------------------------------- 1 | # Blender v3.0.1 OBJ File: 'dex_cube.blend' 2 | # www.blender.org 3 | mtllib dex_cube.mtl 4 | o Cube_Cube.002 5 | v -1.000000 -1.000000 1.000000 6 | v -1.000000 1.000000 1.000000 7 | v -1.000000 -1.000000 -1.000000 8 | v -1.000000 1.000000 -1.000000 9 | v 1.000000 -1.000000 1.000000 10 | v 1.000000 1.000000 1.000000 11 | v 1.000000 -1.000000 -1.000000 12 | v 1.000000 1.000000 -1.000000 13 | vt 1.000000 0.333333 14 | vt 1.000000 0.666667 15 | vt 0.500000 0.666667 16 | vt 0.500000 0.333333 17 | vt 0.500000 0.666667 18 | vt 0.500000 1.000000 19 | vt 0.000000 1.000000 20 | vt -0.000000 0.666667 21 | vt 0.500000 0.333333 22 | vt 0.500000 0.666667 23 | vt -0.000000 0.666667 24 | vt -0.000000 0.333333 25 | vt 1.000000 -0.000000 26 | vt 1.000000 0.333333 27 | vt 0.500000 0.333333 28 | vt 0.500000 0.000000 29 | vt 0.500000 0.000000 30 | vt 0.000000 0.000000 31 | vt 1.000000 0.666667 32 | vt 1.000000 1.000000 33 | vt 0.500000 1.000000 34 | vt 0.500000 0.666667 35 | vn -1.0000 0.0000 0.0000 36 | vn 0.0000 0.0000 -1.0000 37 | vn 1.0000 0.0000 0.0000 38 | vn 0.0000 0.0000 1.0000 39 | vn 0.0000 -1.0000 0.0000 40 | vn 0.0000 1.0000 0.0000 41 | usemtl Material.002 42 | s off 43 | f 1/1/1 2/2/1 4/3/1 3/4/1 44 | f 3/5/2 4/6/2 8/7/2 7/8/2 45 | f 7/9/3 8/10/3 6/11/3 5/12/3 46 | f 5/13/4 6/14/4 2/15/4 1/16/4 47 | f 3/17/5 7/9/5 5/12/5 1/18/5 48 | f 8/19/6 4/20/6 2/21/6 6/22/6 49 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 13 | 14 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/dex_cube.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/dex_cube.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/fileback.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/fileback.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/filedown.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/filedown.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/filefront.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/filefront.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/fileleft.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/fileleft.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/fileright.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/fileright.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/fileup.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/fileup.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/grayback.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/grayback.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/graydown.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/graydown.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/grayfront.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/grayfront.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/grayleft.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/grayleft.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/grayright.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/grayright.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/grayup.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/_src/manipulation/leap_hand/xmls/reorientation_cube_textures/grayup.png -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/leap_hand/xmls/scene_mjx_cube.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 20 | 21 | 22 | 23 | 24 | 25 | 26 | 27 | 28 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 36 | 37 | 38 | 39 | 40 | 41 | 42 | 43 | 44 | 45 | 46 | 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | 68 | 69 | 70 | -------------------------------------------------------------------------------- /mujoco_playground/_src/manipulation/manipulation_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Tests for the manipulation environments.""" 16 | from absl.testing import absltest 17 | from absl.testing import parameterized 18 | import jax 19 | import jax.numpy as jp 20 | 21 | from mujoco_playground._src import manipulation 22 | 23 | 24 | class TestSuite(parameterized.TestCase): 25 | """Tests for the manipulation environments.""" 26 | 27 | @parameterized.named_parameters( 28 | {"testcase_name": f"test_can_create_{env_name}", "env_name": env_name} 29 | for env_name in manipulation.ALL_ENVS 30 | ) 31 | def test_can_create_all_environments(self, env_name: str) -> None: 32 | env = manipulation.load(env_name) 33 | state = jax.jit(env.reset)(jax.random.PRNGKey(42)) 34 | state = jax.jit(env.step)(state, jp.zeros(env.action_size)) 35 | self.assertIsNotNone(state) 36 | obs_shape = jax.tree_util.tree_map(lambda x: x.shape, state.obs) 37 | obs_shape = obs_shape[0] if isinstance(obs_shape, tuple) else obs_shape 38 | self.assertEqual(obs_shape, env.observation_size) 39 | self.assertFalse(jp.isnan(state.data.qpos).any()) 40 | 41 | 42 | if __name__ == "__main__": 43 | absltest.main() 44 | -------------------------------------------------------------------------------- /mujoco_playground/_src/registry.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Registry for all environments.""" 16 | from typing import Any, Callable, Dict, Optional, Tuple, Union 17 | 18 | import jax 19 | import ml_collections 20 | from mujoco import mjx 21 | 22 | from mujoco_playground._src import dm_control_suite 23 | from mujoco_playground._src import locomotion 24 | from mujoco_playground._src import manipulation 25 | from mujoco_playground._src import mjx_env 26 | 27 | DomainRandomizer = Optional[ 28 | Callable[[mjx.Model, jax.Array], Tuple[mjx.Model, mjx.Model]] 29 | ] 30 | 31 | 32 | # A tuple containing all available environment names across all suites. 33 | ALL_ENVS = ( 34 | dm_control_suite.ALL_ENVS + locomotion.ALL_ENVS + manipulation.ALL_ENVS 35 | ) 36 | 37 | 38 | def get_default_config(env_name: str): 39 | if env_name in manipulation.ALL_ENVS: 40 | return manipulation.get_default_config(env_name) 41 | elif env_name in locomotion.ALL_ENVS: 42 | return locomotion.get_default_config(env_name) 43 | elif env_name in dm_control_suite.ALL_ENVS: 44 | return dm_control_suite.get_default_config(env_name) 45 | 46 | raise ValueError(f"Env '{env_name}' not found in default configs.") 47 | 48 | 49 | def load( 50 | env_name: str, 51 | config: Optional[ml_collections.ConfigDict] = None, 52 | config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None, 53 | ) -> mjx_env.MjxEnv: 54 | if env_name in manipulation.ALL_ENVS: 55 | return manipulation.load(env_name, config, config_overrides) 56 | elif env_name in locomotion.ALL_ENVS: 57 | return locomotion.load(env_name, config, config_overrides) 58 | elif env_name in dm_control_suite.ALL_ENVS: 59 | return dm_control_suite.load(env_name, config, config_overrides) 60 | 61 | raise ValueError(f"Env '{env_name}' not found. Available envs: {ALL_ENVS}") 62 | 63 | 64 | def get_domain_randomizer(env_name: str) -> Optional[DomainRandomizer]: 65 | if env_name in manipulation.ALL_ENVS: 66 | return manipulation.get_domain_randomizer(env_name) 67 | 68 | if env_name in locomotion.ALL_ENVS: 69 | return locomotion.get_domain_randomizer(env_name) 70 | 71 | return None 72 | -------------------------------------------------------------------------------- /mujoco_playground/_src/registry_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Tests for the registry module.""" 16 | 17 | from absl.testing import absltest 18 | from ml_collections import config_dict 19 | 20 | from mujoco_playground._src import dm_control_suite 21 | from mujoco_playground._src import locomotion 22 | from mujoco_playground._src import manipulation 23 | from mujoco_playground._src import registry 24 | 25 | 26 | class RegistryTest(absltest.TestCase): 27 | 28 | def test_new_env(self): 29 | class DemoEnv: 30 | 31 | def __init__(self, config, config_overrides): 32 | pass 33 | 34 | def demo_default_config() -> None: 35 | return config_dict.ConfigDict() 36 | 37 | dm_control_suite.register_environment( 38 | 'DemoEnv', DemoEnv, demo_default_config 39 | ) 40 | env = registry.load('DemoEnv') 41 | self.assertIsInstance(env, DemoEnv) 42 | config = registry.get_default_config('DemoEnv') 43 | self.assertEqual(config, config_dict.ConfigDict()) 44 | 45 | manipulation.register_environment('DemoEnv', DemoEnv, demo_default_config) 46 | env = registry.load('DemoEnv') 47 | self.assertIsInstance(env, DemoEnv) 48 | config = registry.get_default_config('DemoEnv') 49 | self.assertEqual(config, config_dict.ConfigDict()) 50 | 51 | locomotion.register_environment('DemoEnv', DemoEnv, demo_default_config) 52 | env = registry.load('DemoEnv') 53 | self.assertIsInstance(env, DemoEnv) 54 | config = registry.get_default_config('DemoEnv') 55 | self.assertEqual(config, config_dict.ConfigDict()) 56 | 57 | def test_constants(self) -> None: 58 | self.assertNotEmpty(dm_control_suite.ALL_ENVS) 59 | self.assertNotEmpty(locomotion.ALL_ENVS) 60 | self.assertNotEmpty(manipulation.ALL_ENVS) 61 | self.assertNotEmpty(registry.ALL_ENVS) 62 | self.assertEqual( 63 | sorted(registry.ALL_ENVS), 64 | sorted( 65 | dm_control_suite.ALL_ENVS 66 | + locomotion.ALL_ENVS 67 | + manipulation.ALL_ENVS 68 | ), 69 | ) 70 | 71 | 72 | if __name__ == '__main__': 73 | absltest.main() 74 | -------------------------------------------------------------------------------- /mujoco_playground/_src/wrapper_test.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | """Tests for the wrapper module.""" 16 | import functools 17 | 18 | from absl.testing import absltest 19 | import jax 20 | import jax.numpy as jp 21 | import numpy as np 22 | 23 | from mujoco_playground._src import dm_control_suite 24 | from mujoco_playground._src import wrapper 25 | 26 | 27 | class WrapperTest(absltest.TestCase): 28 | 29 | def test_auto_reset_wrapper(self): 30 | class DoneEnv: 31 | 32 | def __init__(self, env): 33 | self._env = env 34 | 35 | def reset(self, key): 36 | return self._env.reset(key) 37 | 38 | def step(self, state, action): 39 | state = self._env.step(state, jp.ones_like(action)) 40 | state = state.replace(done=action[0] > 0) 41 | return state 42 | 43 | env = wrapper.BraxAutoResetWrapper( 44 | DoneEnv(dm_control_suite.load('CartpoleBalance')) 45 | ) 46 | 47 | jit_reset = jax.jit(env.reset) 48 | jit_step = jax.jit(env.step) 49 | state = jit_reset(jax.random.PRNGKey(0)) 50 | first_qpos = state.info['first_state'].qpos 51 | 52 | state = jit_step(state, -jp.ones(env._env.action_size)) 53 | self.assertGreater(np.linalg.norm(state.data.qpos - first_qpos), 1e-3) 54 | state = jit_step(state, jp.ones(env._env.action_size)) 55 | np.testing.assert_allclose(state.data.qpos, first_qpos, atol=1e-6) 56 | 57 | def test_domain_randomization_wrapper(self): 58 | def randomization_fn(model, rng): 59 | @jax.vmap 60 | def get_gravity(rng): 61 | dg = jax.random.uniform(rng, shape=(3,), minval=-10.0, maxval=10.0) 62 | return model.opt.gravity + dg 63 | 64 | model_v = model.tree_replace({'opt.gravity': get_gravity(rng)}) 65 | in_axes = jax.tree.map(lambda x: None, model) 66 | in_axes = in_axes.tree_replace({'opt.gravity': 0}) 67 | return model_v, in_axes 68 | 69 | env = dm_control_suite.load('CartpoleBalance') 70 | rng = jax.random.PRNGKey(0) 71 | rng = jax.random.split(rng, 256) 72 | env = wrapper.wrap_for_brax_training( 73 | env, 74 | episode_length=200, 75 | randomization_fn=functools.partial(randomization_fn, rng=rng), 76 | ) 77 | 78 | # set the same key across the batch for env.reset so that only the 79 | # randomization wrapper creates variability in the env.step 80 | key = jp.zeros((256, 2), dtype=jp.uint32) 81 | state = jax.jit(env.reset)(key) 82 | self.assertEqual(state.data.qpos[:, 0].shape[0], 256) 83 | self.assertEqual(np.unique(state.data.qpos[:, 0]).shape[0], 1) 84 | 85 | # test that the DomainRandomizationWrapper creates variability in env.step 86 | state = jax.jit(env.step)(state, jp.zeros((256, env.action_size))) 87 | self.assertEqual(state.data.qpos[:, 0].shape[0], 256) 88 | self.assertEqual(np.unique(state.data.qpos[:, 0]).shape[0], 256) 89 | 90 | 91 | if __name__ == '__main__': 92 | absltest.main() 93 | -------------------------------------------------------------------------------- /mujoco_playground/config/README.md: -------------------------------------------------------------------------------- 1 | # Learning 2 | 3 | This directory contains parameters for re-producing RL results for canonical envs in MuJoCo Playground. 4 | -------------------------------------------------------------------------------- /mujoco_playground/config/__init__.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | -------------------------------------------------------------------------------- /mujoco_playground/experimental/learning/README.md: -------------------------------------------------------------------------------- 1 | # learning 2 | 3 | Experimental code for learning RL agents. 4 | -------------------------------------------------------------------------------- /mujoco_playground/experimental/madrona_benchmarking/README.md: -------------------------------------------------------------------------------- 1 | A section for studying the performance of [Madrona MJX](https://github.com/shacklettbp/madrona_mjx) for the CartpoleBalance and PandaPickCubeCartesian envs from Mujoco Playground. 2 | 3 | #### Data Collection 4 | `benchmark.py`: Instantiates the selected env + configuration and creates / benchmarks an unroll as specified in the CLI arguments. Each run appends a row to data/madrona_mjx.csv. 5 | `get_data.sh`: Runs multiple trials of benchmark.py with different arguments. Collects all data required for the two analysis scripts except for PPO training curves, which are loaded from ../data. 6 | 7 | #### Data Analysis 8 | `make_plots.py`: Analyzes pixels Cartpole performance with random actions as a function of image size and num_envs. Produces various plots under figures/ 9 | `print_tables.py`: Analyzes the bottlenecks in PPO training of Cartpole and PandaPickCubeCartesian. Prints a table to stdout. 10 | -------------------------------------------------------------------------------- /mujoco_playground/experimental/madrona_benchmarking/data/maniskill3/README.md: -------------------------------------------------------------------------------- 1 | isaac_lab.csv and maniskill.csv are pulled from: 2 | ``` 3 | Tao, Stone, Fanbo Xiang, Arth Shukla, Yuzhe Qin, Xander Hinrichsen, Xiaodi Yuan, Chen Bao, et al. “ManiSkill3: GPU Parallelized Robotics Simulation and Rendering for Generalizable Embodied AI.” arXiv, October 1, 2024. https://doi.org/10.48550/arXiv.2410.00425. 4 | ``` 5 | 6 | Repository: https://github.com/haosulab/ManiSkill/ 7 | Commit: 03f7315ae1dd58f2bb0499b9cd1b2fc0d0a7f8a8 8 | 9 | isaac_lab_state.csv and maniskill_state.csv are generated using Maniskill3's benchmark script on a RTX4090 GPU. 10 | -------------------------------------------------------------------------------- /mujoco_playground/experimental/madrona_benchmarking/data/maniskill3/isaac_lab_state.csv: -------------------------------------------------------------------------------- 1 | env.step/dt,env.step/fps,env.step/psps,env.step/total_steps,env.step/cpu_mem_use,env.step/gpu_mem_use,env.step+env.reset/dt,env.step+env.reset/fps,env.step+env.reset/psps,env.step+env.reset/total_steps,env.step+env.reset/cpu_mem_use,env.step+env.reset/gpu_mem_use,env_id,obs_mode,num_envs,num_cameras,camera_width,camera_height,sensor_settings,gpu_type 2 | 1.4912841320037842,2682.252103511184,670.5630258777959,1000,2652700672,2732118016,1.477783203125,2706.756980009912,676.6892450024781,1000,2653487104,2732118016,CartpoleBalanceBenchmark-v1,state,4,,,,,NVIDIA GeForce RTX 4090 3 | 1.499715805053711,10668.687991473807,666.792999467113,1000,2653700096,2732118016,1.4709534645080566,10877.298559102286,679.8311599438929,1000,2654224384,2732118016,CartpoleBalanceBenchmark-v1,state,16,,,,,NVIDIA GeForce RTX 4090 4 | 1.4984760284423828,21355.029638520784,667.3446762037745,1000,2679107584,2732118016,1.504070520401001,21275.59816242424,664.8624425757573,1000,2679631872,2732118016,CartpoleBalanceBenchmark-v1,state,32,,,,,NVIDIA GeForce RTX 4090 5 | 1.5351598262786863,41689.470310814206,651.397973606472,1000,2661003264,2732118016,1.572298288345337,40704.744433292384,636.0116317701935,1000,2661265408,2732118016,CartpoleBalanceBenchmark-v1,state,64,,,,,NVIDIA GeForce RTX 4090 6 | 1.5434703826904297,82930.00075380952,647.8906308891369,1000,2672857088,2732118016,1.5958075523376465,80210.17309543182,626.6419773080611,1000,2673381376,2732118016,CartpoleBalanceBenchmark-v1,state,128,,,,,NVIDIA GeForce RTX 4090 7 | 1.5089356899261477,169656.00436724347,662.7187670595448,1000,2687447040,2732118016,1.5033161640167236,170290.19319262248,665.1960671586816,1000,2687709184,2732118016,CartpoleBalanceBenchmark-v1,state,256,,,,,NVIDIA GeForce RTX 4090 8 | 1.5538675785064695,329500.40729475726,643.5554829975728,1000,2707533824,2799226880,1.5294928550720217,334751.4820367636,653.8114883530538,1000,2708058112,2799226880,CartpoleBalanceBenchmark-v1,state,512,,,,,NVIDIA GeForce RTX 4090 9 | 1.5957462787628174,641706.0240891852,626.6660391495949,1000,2788483072,2866335744,1.5697743892669678,652323.0389038095,637.0342176795015,1000,2788745216,2866335744,CartpoleBalanceBenchmark-v1,state,1024,,,,,NVIDIA GeForce RTX 4090 10 | 1.6796798706054688,1219279.9567585243,595.3515413859983,1000,2861690880,3002650624,1.6607286930084229,1233193.6026769262,602.1453138070929,1000,2861953024,3002650624,CartpoleBalanceBenchmark-v1,state,2048,,,,,NVIDIA GeForce RTX 4090 11 | 1.898465871810913,2157531.542082923,526.7410991413386,1000,3040579584,3275280384,1.884688138961792,2173303.856125683,530.5917617494343,1000,3041103872,3275280384,CartpoleBalanceBenchmark-v1,state,4096,,,,,NVIDIA GeForce RTX 4090 12 | 2.3809170722961426,3440691.0241941703,420.0062285393274,1000,3466051584,3820539904,2.371659994125366,3454120.7509894734,421.6455994860197,1000,3467624448,3820539904,CartpoleBalanceBenchmark-v1,state,8192,,,,,NVIDIA GeForce RTX 4090 13 | 3.38368821144104,4842053.692950157,295.5355037201024,1000,4301082624,4955099136,3.3572747707366943,4880148.667844908,297.86063646514333,1000,4302917632,4955099136,CartpoleBalanceBenchmark-v1,state,16384,,,,,NVIDIA GeForce RTX 4090 14 | -------------------------------------------------------------------------------- /mujoco_playground/experimental/madrona_benchmarking/data/maniskill3/maniskill_state.csv: -------------------------------------------------------------------------------- 1 | env.step/dt,env.step/fps,env.step/psps,env.step/total_steps,env.step/cpu_mem_use,env.step/gpu_mem_use,env.step+env.reset/dt,env.step+env.reset/fps,env.step+env.reset/psps,env.step+env.reset/total_steps,env.step+env.reset/cpu_mem_use,env.step+env.reset/gpu_mem_use,env_id,obs_mode,num_envs,control_mode,gpu_type,num_cameras,camera_width,camera_height 2 | 0.9577691555023192,4176.371703995967,1044.092925998992,1000,1879986176,1664315392,0.960362434387207,4165.094194414571,1041.273548603643,1000,1879986176,1664315392,CartpoleBalanceBenchmark-v1,state,4,pd_joint_delta_pos,NVIDIA GeForce RTX 4090,,, 3 | 0.9753913879394532,16403.67159054022,1025.2294744087635,1000,1883738112,1664315392,0.9705688953399658,16485.176968705146,1030.3235605440716,1000,1884262400,1664315392,CartpoleBalanceBenchmark-v1,state,16,pd_joint_delta_pos,NVIDIA GeForce RTX 4090,,, 4 | 0.9899818897247314,32323.823629640065,1010.119488426252,1000,1887064064,1664315392,0.9753108024597168,32810.053902096195,1025.314184440506,1000,1887326208,1664315392,CartpoleBalanceBenchmark-v1,state,32,pd_joint_delta_pos,NVIDIA GeForce RTX 4090,,, 5 | 1.001758337020874,63887.663955290256,998.2447493014104,1000,1912295424,1664315392,0.9975407123565674,64157.78244158864,1002.4653506498224,1000,1912295424,1664315392,CartpoleBalanceBenchmark-v1,state,64,pd_joint_delta_pos,NVIDIA GeForce RTX 4090,,, 6 | 1.001823902130127,127766.96555935644,998.1794184324722,1000,1894322176,1664315392,1.0042462348937988,127458.78008050116,995.7717193789152,1000,1894322176,1664315392,CartpoleBalanceBenchmark-v1,state,128,pd_joint_delta_pos,NVIDIA GeForce RTX 4090,,, 7 | 1.0364587306976318,246994.8801798297,964.82375070246,1000,1906331648,1664315392,1.0155222415924072,252087.04400070527,984.715015627755,1000,1906331648,1664315392,CartpoleBalanceBenchmark-v1,state,256,pd_joint_delta_pos,NVIDIA GeForce RTX 4090,,, 8 | 1.0597212314605713,483145.93008043343,943.6443946883466,1000,1943822336,1664315392,1.0565576553344729,484592.5798889953,946.469882595694,1000,1944346624,1664315392,CartpoleBalanceBenchmark-v1,state,512,pd_joint_delta_pos,NVIDIA GeForce RTX 4090,,, 9 | 1.111661434173584,921143.7660075416,899.5544589917399,1000,2010894336,1731424256,1.1108291149139404,921833.958303598,900.2284749058574,1000,2010894336,1731424256,CartpoleBalanceBenchmark-v1,state,1024,pd_joint_delta_pos,NVIDIA GeForce RTX 4090,,, 10 | 1.2031493186950684,1702199.3597779316,831.1520311415682,1000,2073571328,1731424256,1.189162015914917,1722221.171371935,840.9283063339526,1000,2074095616,1731424256,CartpoleBalanceBenchmark-v1,state,2048,pd_joint_delta_pos,NVIDIA GeForce RTX 4090,,, 11 | 1.359541654586792,3012780.069062985,735.5420090485803,1000,2273992704,1798533120,1.3539693355560305,3025179.295008117,738.569163820341,1000,2274516992,1798533120,CartpoleBalanceBenchmark-v1,state,4096,pd_joint_delta_pos,NVIDIA GeForce RTX 4090,,, 12 | 1.8960976600646973,4320452.565571163,527.3989948206986,1000,2677620736,2004054016,1.891533374786377,4330877.85243291,528.6716128458143,1000,2678407168,2004054016,CartpoleBalanceBenchmark-v1,state,8192,pd_joint_delta_pos,NVIDIA GeForce RTX 4090,,, 13 | 2.6573262214660645,6165596.029440766,376.31811703129677,1000,3506118656,2280878080,2.699517011642456,6069233.840475615,370.43663577121674,1000,3507830784,2282975232,CartpoleBalanceBenchmark-v1,state,16384,pd_joint_delta_pos,NVIDIA GeForce RTX 4090,,, 14 | -------------------------------------------------------------------------------- /mujoco_playground/experimental/madrona_benchmarking/figures/cartpole_benchmark.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/experimental/madrona_benchmarking/figures/cartpole_benchmark.png -------------------------------------------------------------------------------- /mujoco_playground/experimental/madrona_benchmarking/figures/cartpole_benchmark_combined.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/experimental/madrona_benchmarking/figures/cartpole_benchmark_combined.png -------------------------------------------------------------------------------- /mujoco_playground/experimental/madrona_benchmarking/figures/cartpole_benchmark_full.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/experimental/madrona_benchmarking/figures/cartpole_benchmark_full.png -------------------------------------------------------------------------------- /mujoco_playground/experimental/madrona_benchmarking/figures/cartpole_benchmark_stacked.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/experimental/madrona_benchmarking/figures/cartpole_benchmark_stacked.png -------------------------------------------------------------------------------- /mujoco_playground/experimental/madrona_benchmarking/figures/cartpole_benchmark_stacked_full.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/experimental/madrona_benchmarking/figures/cartpole_benchmark_stacked_full.png -------------------------------------------------------------------------------- /mujoco_playground/experimental/madrona_benchmarking/figures/cartpole_benchmark_state.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/experimental/madrona_benchmarking/figures/cartpole_benchmark_state.png -------------------------------------------------------------------------------- /mujoco_playground/experimental/madrona_benchmarking/figures/madrona_mjx_main.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/experimental/madrona_benchmarking/figures/madrona_mjx_main.png -------------------------------------------------------------------------------- /mujoco_playground/experimental/madrona_benchmarking/get_data.sh: -------------------------------------------------------------------------------- 1 | # Main Paper Figure. Use img_size = 64 data from Bottlenecks Figure. 2 | for img_size in 128 256 512 3 | do 4 | for env in "CartpoleBalance" "PandaPickCubeCartesian" 5 | do 6 | for ((run=0; run<5; run++)); do 7 | python benchmark.py \ 8 | --num-envs 1024 --img-size $img_size --env-name $env \ 9 | --measurement-mode 1 --bottleneck-mode 10 | done 11 | done 12 | done 13 | 14 | # Benchmark against Maniskill3 tables (state-based) 15 | for num_envs in 4 16 32 64 128 256 512 1024 2048 4096 8192 16384 16 | do 17 | for ((run=0; run<5; run++)); do 18 | python benchmark.py \ 19 | --num-envs $num_envs --img-size 0 --env-name CartpoleBalance \ 20 | --measurement-mode 0 21 | done 22 | done 23 | 24 | # Benchmark against Maniskill3 tables (pixel-based) 25 | for img_size in 80 128 224 256 26 | do 27 | for num_envs in 128 256 512 1024 28 | do 29 | for ((run=0; run<5; run++)); do 30 | python benchmark.py \ 31 | --num-envs $num_envs --img-size $img_size --env-name CartpoleBalance \ 32 | --measurement-mode 1 33 | done 34 | done 35 | done 36 | 37 | # Bottlenecks Figure 38 | for env in "CartpoleBalance" "PandaPickCubeCartesian" 39 | do 40 | for mode in 0 1 2 41 | do 42 | for ((run=0; run<5; run++)); do 43 | python benchmark.py \ 44 | --num-envs 1024 --env-name $env --img-size 64 --measurement-mode $mode \ 45 | --bottleneck-mode 46 | done 47 | done 48 | done 49 | 50 | -------------------------------------------------------------------------------- /mujoco_playground/experimental/sim2sim/README.md: -------------------------------------------------------------------------------- 1 | # Sim2Sim Transfer 2 | 3 | In this directory, we demonstrate how to deploy a Go1 joystick controller trained with Playground in native C MuJoCo and interact with it using a joystick. 4 | 5 | ## Usage 6 | 7 | ```bash 8 | python play_go1_joystick.py 9 | ``` 10 | 11 | 12 | Watch the video 13 | 14 | 15 | ## Requirements 16 | 17 | We'll need 2 additional dependencies: 18 | 19 | 1. `onnxruntime` for running the ONNX model. 20 | 2. `hidapi` for reading the joystick. 21 | 22 | ```bash 23 | uv pip install onnxruntime hidapi 24 | ``` 25 | 26 | On macOS, you'll need to install `hidapi` with brew and correctly set the `DYLD_LIBRARY_PATH` environment variable. 27 | 28 | ```bash 29 | brew install hidapi 30 | export DYLD_LIBRARY_PATH=/opt/homebrew/Cellar/hidapi/0.14.0/lib:$DYLD_LIBRARY_PATH 31 | ``` 32 | 33 | ## Joystick 34 | 35 | We use a Logitech G F710 Wireless Gamepad in this example. You can buy one on [Amazon](https://www.amazon.com/Logitech-Wireless-Nano-Receiver-Controller-Vibration/dp/B0041RR0TW) for $40. In principle, you can use any joystick of your choice, but you'll need to modify `gamepad_reader.py` to support it. 36 | 37 | Logitech F710 Gamepad 38 | 39 | - Why not use `inputs`? It didn't seem to read any joystick on macOS. 40 | - Why not use `pygame`? PyGame and MuJoCo's viewer don't place nice together. pygame needs to run on the main thread on macOS and we use the managed viewer in MuJoCo which runs the policy in a callback thread. 41 | 42 | ## Exporting a trained policy 43 | 44 | We have a notebook for exporting trained policies to ONNX format. See `mujoco_playground/experimental/brax_network_to_onnx.ipynb`. 45 | -------------------------------------------------------------------------------- /mujoco_playground/experimental/sim2sim/assets/f710.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/experimental/sim2sim/assets/f710.jpg -------------------------------------------------------------------------------- /mujoco_playground/experimental/sim2sim/gamepad_reader.py: -------------------------------------------------------------------------------- 1 | # Copyright 2025 DeepMind Technologies Limited 2 | # 3 | # Licensed under the Apache License, Version 2.0 (the "License"); 4 | # you may not use this file except in compliance with the License. 5 | # You may obtain a copy of the License at 6 | # 7 | # http://www.apache.org/licenses/LICENSE-2.0 8 | # 9 | # Unless required by applicable law or agreed to in writing, software 10 | # distributed under the License is distributed on an "AS IS" BASIS, 11 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | # See the License for the specific language governing permissions and 13 | # limitations under the License. 14 | # ============================================================================== 15 | # pylint: disable=line-too-long 16 | """Logitech F710 Gamepad class that uses HID under the hood. 17 | 18 | Adapted from motion_imitation: https://github.com/erwincoumans/motion_imitation/tree/master/motion_imitation/robots/gamepad/gamepad_reader.py. 19 | """ 20 | import threading 21 | import time 22 | 23 | import hid 24 | import numpy as np 25 | 26 | 27 | def _interpolate(value, old_max, new_scale, deadzone=0.01): 28 | ret = value * new_scale / old_max 29 | if abs(ret) < deadzone: 30 | return 0.0 31 | return ret 32 | 33 | 34 | class Gamepad: 35 | """Gamepad class that reads from a Logitech F710 gamepad.""" 36 | 37 | def __init__( 38 | self, 39 | vendor_id=0x046D, 40 | product_id=0xC219, 41 | vel_scale_x=0.4, 42 | vel_scale_y=0.4, 43 | vel_scale_rot=1.0, 44 | ): 45 | self._vendor_id = vendor_id 46 | self._product_id = product_id 47 | self._vel_scale_x = vel_scale_x 48 | self._vel_scale_y = vel_scale_y 49 | self._vel_scale_rot = vel_scale_rot 50 | 51 | self.vx = 0.0 52 | self.vy = 0.0 53 | self.wz = 0.0 54 | self.is_running = True 55 | 56 | self._device = None 57 | 58 | self.read_thread = threading.Thread(target=self.read_loop, daemon=True) 59 | self.read_thread.start() 60 | 61 | def _connect_device(self): 62 | try: 63 | self._device = hid.device() 64 | self._device.open(self._vendor_id, self._product_id) 65 | self._device.set_nonblocking(True) 66 | print( 67 | "Connected to" 68 | f" {self._device.get_manufacturer_string()} " 69 | f"{self._device.get_product_string()}" 70 | ) 71 | return True 72 | except (hid.HIDException, OSError) as e: 73 | print(f"Error connecting to device: {e}") 74 | return False 75 | 76 | def read_loop(self): 77 | if not self._connect_device(): 78 | self.is_running = False 79 | return 80 | 81 | while self.is_running: 82 | try: 83 | data = self._device.read(64) 84 | if data: 85 | self.update_command(data) 86 | except (hid.HIDException, OSError) as e: 87 | print(f"Error reading from device: {e}") 88 | 89 | self._device.close() 90 | 91 | def update_command(self, data): 92 | left_x = -(data[1] - 128) / 128.0 93 | left_y = -(data[2] - 128) / 128.0 94 | right_x = -(data[3] - 128) / 128.0 95 | 96 | self.vx = _interpolate(left_y, 1.0, self._vel_scale_x) 97 | self.vy = _interpolate(left_x, 1.0, self._vel_scale_y) 98 | self.wz = _interpolate(right_x, 1.0, self._vel_scale_rot) 99 | 100 | def get_command(self): 101 | return np.array([self.vx, self.vy, self.wz]) 102 | 103 | def stop(self): 104 | self.is_running = False 105 | 106 | 107 | if __name__ == "__main__": 108 | gamepad = Gamepad() 109 | while True: 110 | print(gamepad.get_command()) 111 | time.sleep(0.1) 112 | -------------------------------------------------------------------------------- /mujoco_playground/experimental/sim2sim/onnx/apollo_policy.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/experimental/sim2sim/onnx/apollo_policy.onnx -------------------------------------------------------------------------------- /mujoco_playground/experimental/sim2sim/onnx/bh_policy.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/experimental/sim2sim/onnx/bh_policy.onnx -------------------------------------------------------------------------------- /mujoco_playground/experimental/sim2sim/onnx/g1_policy.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/experimental/sim2sim/onnx/g1_policy.onnx -------------------------------------------------------------------------------- /mujoco_playground/experimental/sim2sim/onnx/go1_policy.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/experimental/sim2sim/onnx/go1_policy.onnx -------------------------------------------------------------------------------- /mujoco_playground/experimental/sim2sim/onnx/leap_reorient_policy.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/experimental/sim2sim/onnx/leap_reorient_policy.onnx -------------------------------------------------------------------------------- /mujoco_playground/experimental/sim2sim/onnx/t1_policy.onnx: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/experimental/sim2sim/onnx/t1_policy.onnx -------------------------------------------------------------------------------- /mujoco_playground/py.typed: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/mujoco_playground/ff0ea5629bd89662f6ffa54464e247653737ea45/mujoco_playground/py.typed -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["hatchling"] 3 | build-backend = "hatchling.build" 4 | 5 | [project] 6 | name = "playground" 7 | version = "0.0.4" 8 | readme = "README.md" 9 | authors = [ 10 | {name = "Google DeepMind", email = "mujoco@deepmind.com"}, 11 | ] 12 | requires-python = ">=3.10" 13 | dynamic = ["description"] 14 | classifiers = [ 15 | "Development Status :: 5 - Production/Stable", 16 | "Framework :: Robot Framework :: Library", 17 | "Intended Audience :: Developers", 18 | "Intended Audience :: Science/Research", 19 | "License :: OSI Approved :: Apache Software License", 20 | "Programming Language :: Python :: 3.10", 21 | "Programming Language :: Python :: 3.11", 22 | "Programming Language :: Python :: 3.12", 23 | "Topic :: Scientific/Engineering", 24 | ] 25 | dependencies = [ 26 | "brax>=0.12.1", 27 | "etils", 28 | "flax", 29 | "jax", 30 | "lxml", 31 | "ml_collections", 32 | "mujoco-mjx>=3.2.7", 33 | "mujoco>=3.2.7", 34 | "tqdm", 35 | ] 36 | keywords = ["mjx", "mujoco", "sim2real", "reinforcement learning"] 37 | 38 | [project.optional-dependencies] 39 | test = [ 40 | "absl-py", 41 | "dm_control", 42 | "pytest", 43 | "pytest-xdist", 44 | ] 45 | notebooks = [ 46 | "matplotlib", 47 | "mediapy", 48 | "jupyter", 49 | "wandb", 50 | ] 51 | dev = [ 52 | "playground[test]", 53 | "mypy", 54 | "pyink", 55 | "pytype", 56 | "ruff", 57 | "pre-commit", 58 | "pylint", 59 | "pytest-xdist", 60 | ] 61 | all = [ 62 | "playground[dev]", 63 | "playground[notebooks]", 64 | ] 65 | 66 | [tool.hatch.build.targets.wheel] 67 | packages = ["mujoco_playground"] 68 | 69 | [tool.ruff] 70 | lint.select = [ 71 | "E", # pycodestyle errors. 72 | "F", # Pyflakes rules. 73 | "PLC", # Pylint convention warnings. 74 | "PLE", # Pylint errors. 75 | "PLR", # Pylint refactor recommendations. 76 | "PLW", # Pylint warnings. 77 | "I" # Import sorting. 78 | ] 79 | lint.ignore = [ 80 | "E741", # Ambiguous variable name. (l, O, or I) 81 | "E501", # Line too long. 82 | "PLR2004", # Magic value used in comparison. 83 | "PLR0915", # Too many statements. 84 | "PLR0913", # Too many arguments. 85 | "PLC0414", # Import alias does not rename variable. (this is used for exporting names) 86 | "PLC1901", # Use falsey strings. 87 | "PLR5501", # Use `elif` instead of `else if`. 88 | "PLR0911", # Too many return statements. 89 | "PLR0912", # Too many branches. 90 | "PLW0603", # Global statement updates are discouraged. 91 | "PLW2901" # For loop variable overwritten. 92 | ] 93 | exclude = [ 94 | "mujoco_menagerie/", 95 | ] 96 | 97 | [tool.mypy] 98 | python_version = "3.12" 99 | ignore_missing_imports = true 100 | warn_unused_configs = true 101 | exclude = ["scripts/", "mujoco_menagerie/"] 102 | 103 | [tool.pytest.ini_options] 104 | testpaths = [ 105 | "mujoco_playground/_src/", 106 | ] 107 | 108 | [tool.isort] 109 | force_single_line = true 110 | force_sort_within_sections = true 111 | lexicographical = true 112 | single_line_exclusions = ["typing"] 113 | order_by_type = false 114 | group_by_package = true 115 | line_length = 120 116 | use_parentheses = true 117 | multi_line_output = 3 118 | include_trailing_comma = true 119 | skip = ["mujoco_menagerie", ".pytype"] 120 | skip_glob = ["**/*.ipynb"] 121 | 122 | [tool.pyink] 123 | line-length = 80 124 | unstable = true 125 | pyink-indentation = 2 126 | pyink-use-majority-quotes = true 127 | extend-exclude = '''( 128 | mujoco_menagerie 129 | | .ipynb$ 130 | )''' 131 | 132 | [pytype] 133 | inputs = "." 134 | exclude = [ 135 | "**/*_test.py", 136 | "**/test_*.py", 137 | "mujoco_menagerie/**", 138 | "mujoco_playground/_src/wrapper_torch.py", 139 | "mujoco_playground/learning/getup_rl_games.py" 140 | ] 141 | output = ".pytype" 142 | report_errors = true 143 | 144 | [tool.pylint] 145 | ignore-paths = 'mujoco_playground/experimental/**$' 146 | 147 | [tool.hatch.build] 148 | include = [ 149 | "mujoco_playground/__init__.py", 150 | "mujoco_playground/_src/**/*", 151 | "mujoco_playground/config/**/*", 152 | ] 153 | --------------------------------------------------------------------------------