32a241c28f
Some checks failed
nightly / Test against latest dependencies (py3.10) (push) Has been cancelled
nightly / Test against latest dependencies (py3.13) (push) Has been cancelled
tests / tests (3.13, locked) (push) Has been cancelled
tests / tests (3.13, unlocked) (push) Has been cancelled
tests / pyright (3.10) (push) Has been cancelled
tests / lint-format (push) Has been cancelled
tests / tests (3.10, locked) (push) Has been cancelled
tests / tests (3.11, locked) (push) Has been cancelled
tests / tests (3.12, locked) (push) Has been cancelled
tests / pyright (3.11) (push) Has been cancelled
tests / pyright (3.12) (push) Has been cancelled
tests / pyright (3.13) (push) Has been cancelled
tests / ty-check (3.10) (push) Has been cancelled
tests / ty-check (3.11) (push) Has been cancelled
tests / ty-check (3.12) (push) Has been cancelled
tests / ty-check (3.13) (push) Has been cancelled
tests / stubs (push) Has been cancelled
tests / smoke-test (push) Has been cancelled
Docker / check_paths (push) Has been cancelled
docs / build (push) Has been cancelled
Docker / build (push) Has been cancelled
Upstream: https://github.com/michaelgillett/mjlab Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047 Upstream-Branch: main
80 lines
3.0 KiB
Python
80 lines
3.0 KiB
Python
"""Tests for the velocity terrain-level curriculum."""
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import torch
|
|
|
|
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
|
from mjlab.tasks.velocity.mdp.curriculums import terrain_levels_vel
|
|
|
|
|
|
class _FakeTerrain:
|
|
"""Minimal terrain stub implementing the real level-update arithmetic."""
|
|
|
|
def __init__(self, num_envs: int, size: float = 8.0):
|
|
self.terrain_levels = torch.zeros(num_envs, dtype=torch.long)
|
|
self.terrain_types = torch.zeros(num_envs, dtype=torch.long)
|
|
self.max_terrain_level = 10
|
|
# [num_rows, num_cols, 3]. One column, enough rows for promotion.
|
|
self.terrain_origins = torch.zeros(self.max_terrain_level, 1, 3)
|
|
self.env_origins = torch.zeros(num_envs, 3)
|
|
self.cfg = SimpleNamespace(
|
|
terrain_generator=SimpleNamespace(size=(size, size), sub_terrains={})
|
|
)
|
|
|
|
def update_env_origins(self, env_ids, move_up, move_down):
|
|
# Mirror TerrainEntity.update_env_origins exactly, including the cap-wrap
|
|
# branch and the env_origins write, so the fake can't pass vacuously.
|
|
self.terrain_levels[env_ids] += 1 * move_up - 1 * move_down
|
|
self.terrain_levels[env_ids] = torch.where(
|
|
self.terrain_levels[env_ids] >= self.max_terrain_level,
|
|
torch.randint_like(self.terrain_levels[env_ids], self.max_terrain_level),
|
|
torch.clip(self.terrain_levels[env_ids], 0),
|
|
)
|
|
self.env_origins[env_ids] = self.terrain_origins[
|
|
self.terrain_levels[env_ids], self.terrain_types[env_ids]
|
|
]
|
|
|
|
|
|
def _make_env(terrain: _FakeTerrain, common_step_counter: int, walked: float):
|
|
num_envs = terrain.terrain_levels.shape[0]
|
|
# Robot has "walked" `walked` meters in x from its origin.
|
|
asset = Mock()
|
|
asset.data.root_link_pos_w = torch.tensor([[walked, 0.0, 0.0]] * num_envs)
|
|
|
|
env = Mock()
|
|
env.common_step_counter = common_step_counter
|
|
env.max_episode_length_s = 20.0
|
|
env.scene.__getitem__ = Mock(return_value=asset)
|
|
env.scene.terrain = terrain
|
|
env.scene.env_origins = torch.zeros(num_envs, 3)
|
|
# Zero command so move_down stays inactive; isolates move_up behavior.
|
|
env.command_manager.get_command = Mock(return_value=torch.zeros(num_envs, 2))
|
|
return env
|
|
|
|
|
|
def test_first_reset_does_not_promote_levels():
|
|
"""On the initial reset the far spawn/origin gap must not bump levels."""
|
|
terrain = _FakeTerrain(num_envs=4, size=8.0)
|
|
# walked >> size/2 (=4.0), which would normally trigger move_up.
|
|
env = _make_env(terrain, common_step_counter=0, walked=100.0)
|
|
|
|
terrain_levels_vel(
|
|
env, torch.arange(4), command_name="twist", asset_cfg=SceneEntityCfg("robot")
|
|
)
|
|
|
|
assert torch.all(terrain.terrain_levels == 0)
|
|
|
|
|
|
def test_subsequent_reset_promotes_on_long_walk():
|
|
"""After stepping, a long walk still promotes the level as before."""
|
|
terrain = _FakeTerrain(num_envs=4, size=8.0)
|
|
env = _make_env(terrain, common_step_counter=1, walked=100.0)
|
|
|
|
terrain_levels_vel(
|
|
env, torch.arange(4), command_name="twist", asset_cfg=SceneEntityCfg("robot")
|
|
)
|
|
|
|
assert torch.all(terrain.terrain_levels == 1)
|