Mobile
Building and Testing Robot Policies with MuJoCo
vmodal_ai DEV Community
1 views
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/solimp for 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
Read original: https://dev.to/vmodal_ai/building-and-testing-robot-policies-with-mujoco-2j1k
← Previous
Creating a Robot Sensor Data Recording and Replay System
Next →
Using VR Controllers for Robot Teleoperation
Related
Comments0
No comments yet — be the first