Height scanner returns Inf values on rough terrain configuration

[sunpihai-up]

Description
When using the height scanner in the rough terrain (g1_rough) configuration, observation data may contain Inf values.

Steps to Reproduce

1. Add an assertion in compute_observations inside base_env.py:

assert torch.isfinite(height_scan).all(), "Height scan contains NaN or Inf values"

2. Run the training with:

python legged_lab/scripts/train.py --task=g1_rough --headless --logger=tensorboard --num_envs=4096

3. Around iteration ~327/50000, the following error appears:

Traceback (most recent call last):
  File "/home/eai/code/locomotion/LeggedLab/legged_lab/scripts/train.py", line 101, in <module>
    train()
  File "/home/eai/code/locomotion/LeggedLab/legged_lab/scripts/train.py", line 97, in train
    runner.learn(num_learning_iterations=agent_cfg.max_iterations, init_at_random_ep_len=True)
  File "/home/eai/anaconda3/envs/leggedlabsim/lib/python3.11/site-packages/rsl_rl/runners/on_policy_runner.py", line 206, in learn
    obs, rewards, dones, infos = self.env.step(actions.to(self.env.device))
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/eai/code/locomotion/LeggedLab/legged_lab/envs/base/base_env.py", line 246, in step
    actor_obs, critic_obs = self.compute_observations()
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/eai/code/locomotion/LeggedLab/legged_lab/envs/base/base_env.py", line 176, in compute_observations
    assert torch.isfinite(height_scan).all(), "Height scan contains NaN or Inf values"
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: Height scan contains NaN or Inf values

Relevant Code

if self.cfg.scene.height_scanner.enable_height_scan:
    height_scan = (
        self.height_scanner.data.pos_w[:, 2].unsqueeze(1)
        - self.height_scanner.data.ray_hits_w[..., 2]
        - self.cfg.normalization.height_scan_offset
    ) * self.obs_scales.height_scan
    assert torch.isfinite(height_scan).all(), "Height scan contains NaN or Inf values"

I want to confirm whether this behavior is what you expect. The ray_caster returns inf because it didn't hit any meshes, but that doesn't seem right.
While this can be easily fixed with nan2num, I'd like to know the cause of this behavior so I can fix it.

[sunpihai-up]

Additional Information

• I also observed this issue when simply playing the step-0 policy network weights for some time. This suggests that the problem is not directly related to training, but rather something that can occur probabilistically.

• I tried increasing the height scanner offset to 100m, and also expanding the terrain boundaries, but the issue still appeared.

• In each case, only a few points in the entire height map become Inf. So ignoring the issue or simply replacing Inf with 0 might not have a major negative impact.

However, I would really like to understand the root cause, and I hope someone can help clarify this.

[Hellod035]

I guess the robot is at the edge of the terrain, which causes the height_scaner ray to not touch the groud, resulting in an invalid value. The good thing is that we have cliped the observations, which does not affect the training and the policy will adapt to it.

[sunpihai-up]

I personally don't think the problem is the perception range exceeding the terrain. This is because the same issue occurs when I run the policy network trained for one iteration. At this point, the robot should be constantly falling, unable to move beyond the terrain. I've also tried increasing the terrain boundaries to ensure the robot doesn't move beyond the terrain, but the issue still occurs. However, as you mentioned, after clipping, this no longer affects training.