From e789838b608888e7c7d6524f7a3972d8b0a36d29 Mon Sep 17 00:00:00 2001 From: mvanhorn Date: Sun, 28 Jun 2026 15:54:57 -0700 Subject: [PATCH] fix: make terrain_levels_vel promotion threshold track commanded velocity --- src/mjlab/tasks/velocity/mdp/curriculums.py | 37 +++++- tests/test_envs_curriculums.py | 135 +++++++++++++++++++- 2 files changed, 164 insertions(+), 8 deletions(-) diff --git a/src/mjlab/tasks/velocity/mdp/curriculums.py b/src/mjlab/tasks/velocity/mdp/curriculums.py index b2487ebfad..13b5cc14b0 100644 --- a/src/mjlab/tasks/velocity/mdp/curriculums.py +++ b/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) diff --git a/tests/test_envs_curriculums.py b/tests/test_envs_curriculums.py index 477b1317b2..53b9d2be8e 100644 --- a/tests/test_envs_curriculums.py +++ b/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