google-deepmind / google-deepmind/mujoco_playground
State explosion to NaN in cube_reorient.ipynb during test rollout
- Dominant language
- Python
- Stars
- 2.2k
- Forks
- 359
- Avg merge
- 7d 3h
- Merged PRs (30d)
- 1
Description
### Description
When running the test rollout code block in mujoco_playground/experimental/learning/cube_reorient.ipynb, the simulation becomes unstable. The state values (qpos and qvel) explode to very large numbers in the first step and then to NaN in the second step. This causes the episode to terminate prematurely.
The issue is present in the original code. I was able to observe the numerical explosion by adding two print statements inside the for loop for debugging purposes.
### To Reproduce
Steps to reproduce the behavior:
1. Open the notebook mujoco_playground/experimental/learning/cube_reorient.ipynb.
2. Run all cells sequentially up to the test rollout block (the cell that begins with rng = jax.random.PRNGKey(1234)).
3. The original code in that cell will cause the simulation to terminate almost immediately. To clearly see the state explosion, you can add the debug lines as shown below.
The following is the original code from the cell, with only the two print lines added for debugging.
```
rng = jax.random.PRNGKey(1234)
rollout = [state := jit_reset(rng)]
actions = []
rewards = []
cube_angvel = []
cube_angacc = []
torques = []
for i in range(env_cfg.episode_length):
act_rng, rng = jax.random.split(rng)
ctrl, _ = jit_inference_fn(state.obs, act_rng)
state = jit_step(state, ctrl)
# --------------------------for debug--------------------------
# These two lines were added to observe the state values.
print(f" max num of qpos: {jp.max(jp.abs(state.data.qpos))} in step {i}")
print(f" max num of qvel: {jp.max(jp.abs(state.data.qvel))} in step {i}")
# -------------------------------------------------------------
rollout.append(state)
rewards.append({k[7:]: v for k, v in state.metrics.items() if k.startswith("reward/")})
actions.append({
"policy_output": ctrl,
"motor_targets": state.info["motor_targets"],
})
torques.append(jp.linalg.norm(state.data.actuator_force))
cube_angvel.append(env.get_cube_angvel(state.data))
cube_angacc.append(env.get_cube_angacc(state.data))
if state.done:
print("Done detected, stopping rollout.")
break
print(rollout[-1].info["success_count"])
```
### Actual Behavior
The simulation state becomes unstable immediately after the first step. The qpos and qvel values explode and then turn into NaN, triggering the state.done condition and stopping the rollout.
Output (from the code with debug lines):
```
max num of qpos: 1444.453369140625 in step 0
max num of qvel: 159579.703125 in step 0
max num of qpos: nan in step 1
max num of qvel: nan in step 1
Done detected, stopping rollout.
0
```
Expected Behavior
The test rollout should run for a significant number of steps (up to env_cfg.episode_length) without the physics state becoming unstable or exploding to NaN.
Environment:
OS: Ubuntu 22.04
Python version: 3.11.12
JAX version: 0.7.1
MuJoCo version: 3.3.6.dev802089588
Contributor guide
Assessment
This issue has not been assessed yet.