DOCSEARCHbyAlgolia
Submit Docs

MuJoCo Playground

repository·main·Indexed 24 days ago

https://github.com/google-deepmind/mujoco_playground

A suite of GPU-accelerated environments for robot learning research and sim-to-real transfer, leveraging MuJoCo MJX and MuJoCo Warp. It includes re-implementations of the DeepMind Control Suite, locomotion and manipulation environments, and support for training RL agents using Brax PPO and RSL-RL. The package provides tools for vision-based training, high-fidelity tendon-driven manipulation models like the Tetheria Aero Hand, and utilities for exporting policies to ONNX for Sim2Sim deployment.

Tokens
29.4K
Snippets
87
Records
106
Agent score
84%

What's inside MuJoCo Playground

  1. Overview of the Tetheria Aero Hand Open implementation

    main

    The Tetheria Aero Hand Open implementation provides a high-fidelity tendon-driven manipulation model. It is designed for sim-to-real transfer by modeling the hand's physical actuation system (tendons, springs, and pulleys) directly from real-world specifications.

    Core Components

    • Tendon Drives: Modeled as spatial tendons in MuJoCo following real cable routing paths.
    • Springs: Modeled as tendon components to provide restoring forces for finger extension.
    • Pulleys: Modeled as cylinders to ensure accurate tendon and spring routing.

    File Structure

    • Implementation:
      • aero_hand_constants.py: Constants and configuration.
      • rotate_z.py: Implementation of the cube rotation task.
    • XML Models:
      • xmls/right_hand.xml: Main hand model with the tendon system.
      • xmls/scene_mjx_cube.xml: Manipulation scene.
      • xmls/reorientation_cube.xml: Cube reorientation task.

  2. Overview of DeepMind Control Suite implementation

    main

    The DeepMind Control Suite in this repository is a re-implementation of the original suite using MJX. This allows for high-performance simulation using MuJoCo's JAX-based engine.

    Below is the current implementation status of the environments:

    | Env                       | Implemented   |
| ------------------------- | ------------- |
| acrobot-swingup           | ✅            |
| acrobot-swingup_sparse    | ✅            |
| ball_in_cup-catch         | ✅            |
| cartpole-balance          | ✅            |
| cartpole-balance_sparse   | ✅            |
| cartpole-swingup          | ✅            |
| cartpole-swingup_sparse   | ✅            |
| cheetah-run               | ✅            |
| finger-spin               | ✅            |
| finger_turn_easy          | ✅            |
| finger_turn_hard          | ✅            |
| fish-upright              | ✅            |
| fish-swim                 | ✅            |
| hopper-stand              | ✅            |
| hopper-hop                | ✅            |
| humanoid-stand            | ✅            |
| humanoid-walk             | ✅            |
| humanoid-run              | ✅            |
| pendulum-swingup          | ✅            |
| point_mass-easy           | ✅            |
| reacher-easy              | ✅            |
| reacher-hard              | ✅            |
| swimmer-swimmer6          | ✅            |
| swimmer-swimmer15        | ✅            |
| walker-stand              | ✅            |
| walker-walk               | ✅            |
| walker-run                | ✅            |
| manipulator-bring_ball    | ❌            |
| manipulator-bring_peg     | ❌            |
| manipulator-insert_ball   | ❌            |
| manipulator-insert_peg   | ❌            |
| dog-stand                 | ❌            |
| dog-walk                  | ❌            |
| dog-trot                  | ❌            |
| dog-run                   | ❌            |
| dog-fetch                 | ❌            |

  3. Control Gait and Commands in Locomotion Environments

    main

    In joystick-based gait tracking environments, you can manipulate the agent's behavior by modifying the state.info dictionary during a rollout.

    Supported keys in state.info:

    • command: A jax.numpy.array representing [linear_velocity_x, linear_velocity_y, yaw_rate].
    • gait: An integer index selecting the gait type from a predefined set (e.g., trot, walk, pace, bound, pronk).
    • phase: The phase timing array corresponding to the selected gait.
    • gait_freq: The frequency of the gait.
    • phase_dt: The calculated phase delta time based on gait_freq and env.dt.
    • foot_height: The target height for foot clearance.

  4. XML Changes Compared to DM Control

    main

    The MuJoCo Playground environments in the dm_control_suite directory include specific XML modifications compared to the original DeepMind Control (DM Control) environments. These changes are primarily implemented to optimize performance on XPU hardware while minimizing deviations from the original physics.

    Key global changes applied to all environments include:

    • Disabling eulerdamp.
    • Reducing solver iterations.
    • Reducing ls_iterations.

    Specific environments may also have modified timestep values, contact culling settings (max_contact_points and max_geom_pairs), or disabled contacts entirely.

  6. Collect performance data for Madrona MJX benchmarks

    main

    To collect performance data for the CartpoleBalance and PandaPickCubeCartesian environments using Madrona MJX, you can use the provided scripts:

    1. Individual Benchmarking: Use benchmark.py to instantiate a specific environment and configuration. It creates and benchmarks an unroll based on your CLI arguments. Each execution appends a new row of results to data/madrona_mjx.csv.
    2. Automated Batch Collection: Use get_data.sh to run multiple trials of benchmark.py with varying arguments. This script collects all necessary data for analysis, except for PPO training curves, which are expected to be loaded from ../data.

  7. Configure MuJoCo for GPU EGL rendering

    main

    When running in environments like Google Colab where the NVIDIA ICD (Installable Client Driver) might be missing, you may need to manually create the ICD config and set the MUJOCO_GL environment variable to egl to enable GPU-accelerated rendering.

    # Create NVIDIA ICD config if missing
NVIDIA_ICD_CONFIG_PATH = '/usr/share/glvnd/egl_vendor.d/10_nvidia.json'
if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):
  with open(NVIDIA_ICD_CONFIG_PATH, 'w') as f:
    f.write('''{
    "file_format_version" : "1.0.0",
    "ICD" : {
        "library_path" : "libEGL_nvidia.so.0"
    }
}''')

# Set MuJoCo to use EGL
%env MUJOCO_GL=egl
    #!pip install mujoco
!pip install mujoco_mjx
!pip install brax
!pip install playground

  8. Configure MuJoCo EGL rendering

    main

    To use EGL for hardware-accelerated offscreen rendering in MuJoCo (often required for headless environments or specific GPU setups), set the MUJOCO_GL environment variable before importing MuJoCo modules.

    import os
os.environ["MUJOCO_GL"] = "egl"

  9. Render Locomotion Rollouts with Visual Overlays

    main

    To visualize a trajectory with visual indicators (like a joystick command overlay), use the env.render method. You can pass modify_scene_fns to inject custom drawing logic into the MuJoCo scene.

    Example workflow:

    1. Collect rollouts by stepping through the environment.
    2. Create modify_scene_fns using draw_joystick_command to visualize the command vector.
    3. Call env.render with the collected trajectories and scene functions.
    from mujoco_playground._src.gait import draw_joystick_command

# ... after collecting 'rollout' and 'modify_scene_fns' ...

frames = eval_env.render(
    traj,
    camera="track",
    scene_option=scene_option,
    modify_scene_fns=modify_scene_fns,
    height=480,
    width=640,
)
media.show_video(frames, fps=fps)

  11. Train a Locomotion Policy with PPO

    main

    Training is typically performed using the brax.training.agents.ppo.train function. You can wrap a MuJoCo environment using BraxEnvWrapper to make it compatible with Brax training.

    Key parameters for ppo.train include:

    • num_timesteps: Total training steps.
    • network_factory: A factory function (often created via ppo_networks.make_ppo_networks) to define the policy architecture.
    • policy_params_fn: A callback for saving checkpoints during training.
    • randomization_fn: The domain randomization function for the environment.
    from mujoco_playground import BraxEnvWrapper
from brax.training.agents.ppo import train as ppo

# ... setup env, env_cfg, randomizer, and train_fn ...

env = locomotion.load(env_name, config=env_cfg)
eval_env = locomotion.load(env_name, config=env_cfg)

make_inference_fn, params, _ = train_fn(
    environment=BraxEnvWrapper(env),
    progress_fn=progress,
    eval_env=BraxEnvWrapper(eval_env),
)

  12. Implement Inverse Kinematics (IK) for locomotion using mink

    main

    To drive a robot (like the Go1) through a gait, use mink to define tasks and solve for velocities via Inverse Kinematics.

    Workflow

    1. Initialize Configuration: Use mink.Configuration(model) to wrap the MuJoCo model.
    2. Define Tasks:
      • mink.FrameTask: For controlling specific frames (e.g., trunk for body pose or site for feet).
      • mink.PostureTask: To maintain a preferred posture (e.g., from a keyframe).
    3. Set Targets: Update task targets (e.g., feet_tasks[i].set_target(...)) based on your gait trajectory.
    4. Solve: Use mink.solve_ik to compute the velocity command.
    5. Integrate: Apply the velocity to the configuration using configuration.integrate_inplace(vel, ctrl_dt).

    Example: Driving feet to target heights

    # ... setup model, configuration, and tasks ...

# Loop through gait steps
for r in rs:
  for i, foot in enumerate(feet):
    foot_pos = feet_positions[i].copy()
    foot_pos[-1] = r[i]  # Update Z-height from gait trajectory
    feet_tasks[i].set_target(mink.SE3.from_translation(foot_pos))

  # Solve IK and integrate
  vel = mink.solve_ik(configuration, tasks, ctrl_dt, solver, 1e-5)
  configuration.integrate_inplace(vel, ctrl_dt)
  mujoco.mj_forward(model, data)
    import mink
import mujoco

# Setup
configuration = mink.Configuration(model)
base_task = mink.FrameTask(frame_name="trunk", frame_type="body", position_cost=1.0, orientation_cost=1.0)
posture_task = mink.PostureTask(model, cost=1e-5)

# Foot tasks
feet_tasks = []
for foot in feet:
  task = mink.FrameTask(frame_name=foot, frame_type="site", position_cost=1.0, orientation_cost=0.0)
  feet_tasks.append(task)

tasks = [base_task, posture_task, *feet_tasks]

# Solve loop snippet
vel = mink.solve_ik(configuration, tasks, ctrl_dt, solver, 1e-5)
configuration.integrate_inplace(vel, ctrl_dt)
mujoco.mj_forward(model, data)