Training of environment Mjlab-Velocity-Rough-Unitree-Go1 crashes because of abnormal reward values

[chengruiz]

Hello,

Training in the built-in environment Mjlab-Velocity-Rough-Unitree-Go1 crashes with

Traceback (most recent call last):
  File "/home/chengrui/Workspaces/mjlab/src/mjlab/scripts/train.py", line 255, in <module>
    main()
    ~~~~^^
  File "/home/chengrui/Workspaces/mjlab/src/mjlab/scripts/train.py", line 251, in main
    launch_training(task_id=chosen_task, args=args)
    ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chengrui/Workspaces/mjlab/src/mjlab/scripts/train.py", line 202, in launch_training
    run_train(task_id, args, log_dir)
    ~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chengrui/Workspaces/mjlab/src/mjlab/scripts/train.py", line 172, in run_train
    runner.learn(
    ~~~~~~~~~~~~^
      num_learning_iterations=cfg.agent.max_iterations, init_at_random_ep_len=True
      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/home/chengrui/Libraries/miniforge3/envs/mjlab/lib/python3.13/site-packages/rsl_rl/runners/on_policy_runner.py", line 102, in learn
    loss_dict = self.alg.update()
  File "/home/chengrui/Libraries/miniforge3/envs/mjlab/lib/python3.13/site-packages/rsl_rl/algorithms/ppo.py", line 259, in update
    self.actor(obs_batch, masks=masks_batch, hidden_state=hidden_states_batch[0], stochastic_output=True)
    ~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chengrui/Libraries/miniforge3/envs/mjlab/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/home/chengrui/Libraries/miniforge3/envs/mjlab/lib/python3.13/site-packages/torch/nn/modules/module.py", line 1787, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/chengrui/Libraries/miniforge3/envs/mjlab/lib/python3.13/site-packages/rsl_rl/models/mlp_model.py", line 122, in forward
    return self.distribution.sample()
           ~~~~~~~~~~~~~~~~~~~~~~~~^^
  File "/home/chengrui/Libraries/miniforge3/envs/mjlab/lib/python3.13/site-packages/torch/distributions/normal.py", line 81, in sample
    return torch.normal(self.loc.expand(shape), self.scale.expand(shape))
           ~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: normal expects all elements of std >= 0.0

This might be because abnormal reward values make training unstable.

[kevinzakka]

We have a bunch of fixes coming out for rough terrain training, including terrain-aware feet shaping terms (current ones assume flat terrain and might encourage the policy to ram the foot into the ground). Should be out before the end of the month.

[kevinzakka]

Hi @chengruiz and @NaCl-1374, can you check if the latest mjlab commit fixes this issue?

[chengruiz]

Hi, I'm currently running the latest version of mjlab and mujoco_warp. It keeps warning Warning: opt.ccd_iterations, currently set to 500, needs to be increased, like many issues mentioned. Even if I increase it to 2000, when it costs 43GB of memory, the warning still exists. I'll ignore this warning and check if the reward values are reasonable.

[NaCl-1374]

Hi, I'm currently running the latest version of mjlab and mujoco_warp. It keeps warning Warning: opt.ccd_iterations, currently set to 500, needs to be increased, like many issues mentioned. Even if I increase it to 2000, when it costs 43GB of memory, the warning still exists. I'll ignore this warning and check if the reward values are reasonable.您好，我目前正在运行最新版本的 mjlab 和 mujoco_warp。它一直出现警告 Warning: opt.ccd_iterations, currently set to 500, needs to be increased ，就像很多问题中提到的那样。即使我将其增加到 2000，占用 43GB 内存，警告仍然存在。我会忽略此警告，并检查奖励值是否合理。

@kevinzakka @chengruiz I'm experiencing a similar issue; the high memory usage is due to the massive demands of collision calculations. I'm using the same environment and reward settings in Legged Gym, and the training results are vastly different.

[chengruiz]

The training exceeded 10,000 iterations, and no abnormal rewards were yielded. I think the current rough terrain setting is much more stable.

[kevinzakka]

Cool I'm gonna mark this as resolved then. Please re-open if you see more instabilities!

[NaCl-1374]

@kevinzakka The latest version still generates a lot ofWarning: opt.ccd_iterations, currently set to 500, needs to be increased.

[kevinzakka]

It's a harmless warning, see #783. This issue is strictly about NaNs / instabilities.

[NaCl-1374]

@kevinzakka There are too many printouts, I can't see the RSL log anymore.