NAN from the simulation will ruin the training

https://github.com/google-deepmind/mujoco_playground/blob/bfac26213883acd08a104fd614a23021b251966d/mujoco_playground/_src/dm_control_suite/humanoid.py#L119
When training with the DM_Control environment, it will face nan problem because of the simulation instability. The previous code show that we will set the environment to done if nan is detected. However we didn't do this operation to the reward and I think reward will pass to the loss computation and thus pollute the critic, policy and other value relevant to reward. As observed in my training, after the nan appear then all the reward will be nan and training loss will also be nan.
The solution of this problem is to add the same detection as done. If the nan is detected then set the reward to zero to cut the spread of nan.
```
def step(self, state: mjx_env.State, action: jax.Array) -> mjx_env.State:
data = mjx_env.step(self.mjx_model, state.data, action, self.n_substeps)
reward = self._get_reward(data, action, state.info, state.metrics) # pylint: disable=redefined-outer-name
obs = self._get_obs(data, state.info)
done = jp.isnan(data.qpos).any() | jp.isnan(data.qvel).any()
done = done.astype(float)
**reward = (1-done) * reward**
return mjx_env.State(data, obs, reward, done, state.metrics, state.info)
```
Training instability is a common problem which I think not only in JAX based MJX but also in CPU based MuJoCo. Is this modification rational?