google-deepmind / google-deepmind/mujoco_playground

State explosion to NaN in cube_reorient.ipynb during test rollout

Open
#203 2 comments 0 reactions 1 assignee Claimed by @sen-code-lost View on GitHub
question
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.