Building and Testing Robot Policies with MuJoCo
MuJoCo has become one of the go-to physics simulators for robot learning — it's fast, accurate for contact-rich manipulation, and has first-class Python bindings via mujoco and dm_control. This tutorial walks through setting up a robot environment in MuJoCo, wiring it up to a policy, and building a testing loop you can trust before ever touching real hardware.
Why MuJoCo for Policy Testing?
- Speed — you can run thousands of steps per second on CPU, which matters when you're iterating on policies or running large-scale evaluation
- Contact accuracy — MuJoCo's solver is well suited to grasping and manipulation, where contact dynamics dominate behavior
- Deterministic, scriptable resets — essential for reproducible evaluation and for domain randomization (covered in the next tutorial)
Step 1: Setting Up a Robot MJCF Model
MuJoCo scenes are described in MJCF (an XML format). Most common robot arms already have community or vendor-provided MJCF models (e.g., via MuJoCo Menagerie). A minimal scene combining a robot arm and a manipulable object looks like:
<mujoco model="arm_scene">
<include file="robot_arm.xml"/>
<worldbody>
<body name="target_object" pos="0.5 0 0.1">
<freejoint/>
<geom type="box" size="0.02 0.02 0.02" rgba="0.8 0.2 0.2 1" mass="0.1"/>
</body>
<geom name="table" type="plane" size="1 1 0.1" pos="0 0 0"/>
</worldbody>
</mujoco>
Keep the object and robot definitions in separate files and <include> them — this makes it trivial to swap objects for domain randomization or multi-task training later.
Step 2: Wrapping the Sim in a Gym-Style Environment
Wrapping MuJoCo in a standard reset()/step() interface keeps your policy code simulator-agnostic:
import mujoco
import numpy as np
class MuJoCoManipulationEnv:
def __init__(self, xml_path, max_steps=200):
self.model = mujoco.MjModel.from_xml_path(xml_path)
self.data = mujoco.MjData(self.model)
self.max_steps = max_steps
self.step_count = 0
def reset(self):
mujoco.mj_resetData(self.model, self.data)
# Randomize object starting position slightly
obj_qpos_addr = self.model.jnt_qposadr[
self.model.joint("target_object_freejoint").id
]
self.data.qpos[obj_qpos_addr:obj_qpos_addr + 2] += np.random.uniform(-0.05, 0.05, 2)
mujoco.mj_forward(self.model, self.data)
self.step_count = 0
return self._get_obs()
def step(self, action):
self.data.ctrl[:] = action
mujoco.mj_step(self.model, self.data)
self.step_count += 1
obs = self._get_obs()
done = self.step_count >= self.max_steps
success = self._check_success()
return obs, 0.0, done, {"success": success}
def _get_obs(self):
return {
"joint_positions": self.data.qpos[: self.model.nu].copy(),
"joint_velocities": self.data.qvel[: self.model.nu].copy(),
"ee_pose": self._compute_ee_pose(),
}
def _compute_ee_pose(self):
ee_id = self.model.body("end_effector").id
return np.concatenate([self.data.xpos[ee_id], self.data.xquat[ee_id]])
def _check_success(self):
obj_id = self.model.body("target_object").id
target_pos = np.array([0.5, 0.3, 0.1])
return np.linalg.norm(self.data.xpos[obj_id] - target_pos) < 0.03
Step 3: Rendering for Vision-Based Policies
If your policy consumes images (as most imitation learning policies from earlier tutorials in this series do), use MuJoCo's offscreen renderer:
class RenderedMuJoCoEnv(MuJoCoManipulationEnv):
def __init__(self, xml_path, camera_name="front_cam", width=224, height=224, **kwargs):
super().__init__(xml_path, **kwargs)
self.renderer = mujoco.Renderer(self.model, height=height, width=width)
self.camera_name = camera_name
def render(self):
self.renderer.update_scene(self.data, camera=self.camera_name)
return self.renderer.render()
def _get_obs(self):
obs = super()._get_obs()
obs["image"] = self.render()
return obs
Match the render resolution and camera placement to your real robot's camera setup as closely as possible — this is the first step toward sim-to-real consistency, which the next tutorial covers in depth.
Step 4: Running Policy Evaluation
With the environment wrapped, plug in any policy (from the imitation learning pipeline, or a scripted baseline) and run batched evaluation:
def evaluate_policy(env, policy, n_episodes=50):
successes = 0
for ep in range(n_episodes):
obs = env.reset()
done = False
while not done:
action = policy.predict(obs)
obs, _, done, info = env.step(action)
successes += int(info["success"])
success_rate = successes / n_episodes
print(f"Success rate: {success_rate:.1%} over {n_episodes} episodes")
return success_rate
Step 5: Debugging with the Interactive Viewer
Before running headless evaluation at scale, use MuJoCo's built-in interactive viewer to sanity-check your scene, joint limits, and control ranges:
import mujoco.viewer
def debug_viewer(env):
with mujoco.viewer.launch_passive(env.model, env.data) as viewer:
while viewer.is_running():
action = np.zeros(env.model.nu) # replace with policy or teleop input
env.step(action)
viewer.sync()
This catches issues — like a gripper that can't physically reach the object, or a control range that's too narrow — far faster than debugging through headless logs.
Common Pitfalls
-
Control range mismatch. MuJoCo actuators have their own
ctrlrange; make sure your policy's action space is normalized/scaled consistently with what the XML defines. -
Timestep mismatch with real robot control rate. If your real robot runs at 20Hz but your MuJoCo simulation timestep is much finer, use
frame_skip(apply the same action for N physics steps) to match effective control frequency. -
Unrealistic contact stiffness. Default MuJoCo contact parameters can make objects behave unrealistically bouncy or sticky — tune
solref/solimpfor materials that matter to your task.
What's Next
Once your policy performs well in MuJoCo, the real question becomes: will it work on the physical robot? That's the subject of the next tutorial — sim-to-real transfer.
Useful Links
Website: www.v-modal.com
SDK Flutter: v-modal/vmodal_sdk_flutter
SDK Android: v-modal/vmodal_sdk_android
Discord: https://discord.gg/K72z28KUx
Top comments (0)