fix: make terrain_levels_vel promotion threshold track commanded velocity

src/mjlab/tasks/velocity/mdp/curriculums.py

@@ -27,6 +27,8 @@ def terrain_levels_vel(
   env_ids: torch.Tensor,
   command_name: str,
   asset_cfg: SceneEntityCfg = _DEFAULT_SCENE_CFG,
+  promotion_frac: float = 0.8,
+  min_expected_distance: float = 0.1,
 ) -> dict[str, torch.Tensor]:
   asset: Entity = env.scene[asset_cfg.name]
 
@@ -44,15 +46,38 @@ def terrain_levels_vel(
     dim=1,
   )
 
-  # Robots that walked far enough progress to harder terrains.
-  move_up = distance > terrain_generator.size[0] / 2
+  # Estimate expected straight-line displacement from the commanded velocity.
+  command_speed = torch.norm(command[env_ids, :2], dim=1)
+  expected_distance = command_speed * env.max_episode_length_s
+
+  # If the robot is commanded to turn while moving, its net displacement is the
+  # chord of the commanded arc rather than the full path length.
+  if command.shape[1] >= 3:
+    yaw_delta = torch.abs(command[env_ids, 2]) * env.max_episode_length_s
+    arc_ratio = torch.ones_like(yaw_delta)
+    is_turning = yaw_delta > 1e-6
+    arc_ratio[is_turning] = torch.abs(
+      torch.sin(yaw_delta[is_turning] / 2.0) / (yaw_delta[is_turning] / 2.0)
+    )
+    expected_distance *= arc_ratio
+
+  is_commanded_to_move = expected_distance >= min_expected_distance
+  max_promotion_distance = torch.full_like(
+    expected_distance,
+    terrain_generator.size[0] / 2,
+  )
+  promotion_distance = torch.minimum(
+    expected_distance * promotion_frac,
+    max_promotion_distance,
+  )
+
+  # Robots that walked far enough for their command progress to harder terrains.
+  move_up = is_commanded_to_move & (distance > promotion_distance)
 
   # Robots that walked less than half of their required distance go to
   # simpler terrains.
-  move_down = (
-    distance < torch.norm(command[env_ids, :2], dim=1) * env.max_episode_length_s * 0.5
-  )
-  move_down *= ~move_up
+  move_down = is_commanded_to_move & (distance < expected_distance * 0.5)
+  move_down &= ~move_up
 
   # Update terrain levels.
   terrain.update_env_origins(env_ids, move_up, move_down)

tests/test_envs_curriculums.py

@@ -1,6 +1,7 @@
-"""Tests for reward_curriculum and termination_curriculum."""
+"""Tests for curriculum terms."""
 
-from unittest.mock import Mock
+from types import SimpleNamespace
+from unittest.mock import MagicMock, Mock
 
 import pytest
 import torch
@@ -9,6 +10,7 @@
 from mjlab.managers.curriculum_manager import CurriculumTermCfg
 from mjlab.managers.reward_manager import RewardTermCfg
 from mjlab.managers.termination_manager import TerminationTermCfg
+from mjlab.tasks.velocity.mdp.curriculums import terrain_levels_vel
 
 
 def _reward_func(env):
@@ -67,6 +69,135 @@ def _make_termination_env(step_counter, term_cfg):
   return env
 
 
+def _make_terrain_levels_env(
+  command: torch.Tensor,
+  distance: torch.Tensor,
+  *,
+  max_episode_length_s: float = 10.0,
+  terrain_size: float = 8.0,
+):
+  num_envs = command.shape[0]
+  env_origins = torch.zeros((num_envs, 3))
+  root_link_pos_w = env_origins.clone()
+  root_link_pos_w[:, 0] = distance
+  asset = SimpleNamespace(
+    data=SimpleNamespace(root_link_pos_w=root_link_pos_w),
+  )
+
+  terrain_generator = SimpleNamespace(
+    size=(terrain_size, terrain_size),
+    sub_terrains={"flat": object(), "rough": object()},
+  )
+  terrain = SimpleNamespace(
+    cfg=SimpleNamespace(terrain_generator=terrain_generator),
+    update_env_origins=Mock(),
+    terrain_levels=torch.tensor([1, 3], dtype=torch.int64)[:num_envs],
+    terrain_types=torch.tensor([0, 1], dtype=torch.int64)[:num_envs],
+    terrain_origins=torch.zeros((4, 2, 3)),
+  )
+
+  scene = MagicMock()
+  scene.__getitem__.return_value = asset
+  scene.env_origins = env_origins
+  scene.terrain = terrain
+
+  env = SimpleNamespace(
+    scene=scene,
+    command_manager=Mock(),
+    max_episode_length_s=max_episode_length_s,
+  )
+  env.command_manager.get_command.return_value = command
+
+  return env, torch.arange(num_envs), terrain
+
+
+def _terrain_level_masks(terrain) -> tuple[torch.Tensor, torch.Tensor]:
+  _, move_up, move_down = terrain.update_env_origins.call_args.args
+  return move_up, move_down
+
+
+# Terrain levels: velocity task
+
+
+def test_terrain_levels_vel_promotes_full_expected_distance():
+  command = torch.tensor([[0.5, 0.0, 0.0]])
+  env, env_ids, terrain = _make_terrain_levels_env(
+    command,
+    torch.tensor([5.0]),
+    max_episode_length_s=10.0,
+  )
+
+  terrain_levels_vel(env, env_ids, command_name="twist")
+
+  move_up, move_down = _terrain_level_masks(terrain)
+  assert move_up.tolist() == [True]
+  assert move_down.tolist() == [False]
+
+
+def test_terrain_levels_vel_promotes_low_speed_tracking_on_short_tiles():
+  command = torch.tensor([[0.2, 0.0, 0.0]])
+  env, env_ids, terrain = _make_terrain_levels_env(
+    command,
+    torch.tensor([1.0]),
+    max_episode_length_s=5.0,
+    terrain_size=3.0,
+  )
+
+  terrain_levels_vel(env, env_ids, command_name="twist")
+
+  move_up, move_down = _terrain_level_masks(terrain)
+  assert move_up.tolist() == [True]
+  assert move_down.tolist() == [False]
+
+
+def test_terrain_levels_vel_demotes_under_travel_and_masks_are_exclusive():
+  command = torch.tensor([[0.5, 0.0, 0.0]])
+  env, env_ids, terrain = _make_terrain_levels_env(
+    command,
+    torch.tensor([2.0]),
+    max_episode_length_s=10.0,
+  )
+
+  terrain_levels_vel(env, env_ids, command_name="twist")
+
+  move_up, move_down = _terrain_level_masks(terrain)
+  assert move_up.tolist() == [False]
+  assert move_down.tolist() == [True]
+  assert not torch.any(move_up & move_down)
+
+
+def test_terrain_levels_vel_ignores_near_zero_commands():
+  command = torch.tensor([[0.01, 0.0, 0.0]])
+  env, env_ids, terrain = _make_terrain_levels_env(
+    command,
+    torch.tensor([0.02]),
+    max_episode_length_s=1.0,
+  )
+
+  terrain_levels_vel(env, env_ids, command_name="twist")
+
+  move_up, move_down = _terrain_level_masks(terrain)
+  assert move_up.tolist() == [False]
+  assert move_down.tolist() == [False]
+
+
+def test_terrain_levels_vel_original_arguments_and_result_keys():
+  command = torch.tensor(
+    [
+      [0.5, 0.0, 0.0],
+      [0.5, 0.0, 0.0],
+    ]
+  )
+  env, env_ids, _terrain = _make_terrain_levels_env(
+    command,
+    torch.tensor([5.0, 0.0]),
+  )
+
+  result = terrain_levels_vel(env, env_ids, command_name="twist")
+
+  assert set(result) == {"mean", "max", "flat", "rough"}
+
+
 # Reward: weight