Import upstream snapshot d424a0c899f6b33cbd3daeb279913134349c0b63
Upstream: https://github.com/pollen-robotics/microduck_rl Upstream-Commit: d424a0c899f6b33cbd3daeb279913134349c0b63 Upstream-Branch: develop
This commit is contained in:
Upstream Snapshot
@@ -0,0 +1,139 @@
|
||||
"""On linux-aarch64 (DGX Spark / GB10) PyPI's torch wheel is CPU-ONLY:
|
||||
torch.version.cuda is None -> torch.cuda.device_count() == 0 -> mjlab's
|
||||
select_gpus() indexes an empty list and dies with
|
||||
`IndexError: list index out of range` BEFORE the first training step
|
||||
(mjlab/utils/gpu.py:70).
|
||||
|
||||
The fix (pyproject.toml) routes torch to PyTorch's CUDA index, on aarch64
|
||||
only. It has two SILENT break points, locked in by these tests — in both
|
||||
cases `uv sync` succeeds and you only find out when you launch a run:
|
||||
|
||||
1. `torch` must stay a DIRECT dependency: uv applies [tool.uv.sources] to
|
||||
direct dependencies only, so deleting the `torch==...` line (which looks
|
||||
redundant, since torch already comes in via mjlab/rsl_rl) makes the
|
||||
source binding a no-op without any warning.
|
||||
2. The x86_64 resolution must stay on PyPI, otherwise HF Jobs silently
|
||||
switch wheels.
|
||||
"""
|
||||
|
||||
import platform
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tomllib
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
_ROOT = Path(__file__).resolve().parents[1]
|
||||
_CUDA_INDEX = "https://download.pytorch.org/whl/cu"
|
||||
|
||||
|
||||
def _packages(name):
|
||||
lock = tomllib.loads((_ROOT / "uv.lock").read_text())
|
||||
return [p for p in lock["package"] if p["name"] == name]
|
||||
|
||||
|
||||
def _registry(pkg):
|
||||
return pkg.get("source", {}).get("registry", "")
|
||||
|
||||
|
||||
def _markers(pkg):
|
||||
return " ".join(pkg.get("resolution-markers", []))
|
||||
|
||||
|
||||
def _aarch64_entry(pkgs):
|
||||
"""The entry whose resolution-markers SELECT linux-aarch64."""
|
||||
hits = [
|
||||
p
|
||||
for p in pkgs
|
||||
if "platform_machine == 'aarch64'" in _markers(p)
|
||||
and "sys_platform == 'linux'" in _markers(p)
|
||||
]
|
||||
assert len(hits) == 1, f"expected 1 aarch64 entry, found {len(hits)}"
|
||||
return hits[0]
|
||||
|
||||
|
||||
def test_torch_is_a_direct_dependency():
|
||||
"""Without this, [tool.uv.sources] for torch is a silent no-op."""
|
||||
pyproject = tomllib.loads((_ROOT / "pyproject.toml").read_text())
|
||||
deps = pyproject["project"]["dependencies"]
|
||||
assert any(d.split("=")[0].split("[")[0].strip() == "torch" for d in deps), (
|
||||
"torch must stay in [project.dependencies]: uv applies "
|
||||
"[tool.uv.sources] to DIRECT dependencies only. Removing it silently "
|
||||
"drops aarch64 back onto PyPI's CPU-only wheel."
|
||||
)
|
||||
|
||||
|
||||
def test_torch_source_is_pinned_to_a_cuda_index_on_aarch64():
|
||||
pyproject = tomllib.loads((_ROOT / "pyproject.toml").read_text())
|
||||
uv_cfg = pyproject["tool"]["uv"]
|
||||
assert "torch" in uv_cfg.get("sources", {}), (
|
||||
"[tool.uv.sources] no longer has a torch entry -> aarch64 falls back "
|
||||
"to PyPI's CPU wheel and `train` dies with IndexError in select_gpus()."
|
||||
)
|
||||
sources = uv_cfg["sources"]["torch"]
|
||||
indexes = {p["name"]: p["url"] for p in uv_cfg.get("index", [])}
|
||||
for src in sources:
|
||||
assert "aarch64" in src["marker"], "the torch source must stay aarch64-scoped"
|
||||
assert indexes[src["index"]].startswith(_CUDA_INDEX), (
|
||||
f"index {src['index']} is not a PyTorch CUDA index"
|
||||
)
|
||||
|
||||
|
||||
def test_lockfile_routes_aarch64_torch_to_cuda_wheels():
|
||||
torch_pkgs = _packages("torch")
|
||||
aarch64 = _aarch64_entry(torch_pkgs)
|
||||
assert _registry(aarch64).startswith(_CUDA_INDEX), (
|
||||
f"torch on aarch64 comes from {_registry(aarch64)!r} — a CPU wheel. "
|
||||
"Re-run `uv lock` after checking [tool.uv.sources]."
|
||||
)
|
||||
wheels = " ".join(w["url"] for w in aarch64["wheels"])
|
||||
assert "aarch64" in wheels, "no aarch64 wheel in the aarch64 torch entry"
|
||||
assert "%2Bcu" in wheels or "+cu" in wheels, (
|
||||
"the aarch64 wheel has no +cuXXX local version -> CPU build"
|
||||
)
|
||||
|
||||
|
||||
def test_x86_64_resolution_stays_on_pypi():
|
||||
"""HF Jobs run on x86_64: their resolution must not move."""
|
||||
others = [
|
||||
p
|
||||
for p in _packages("torch")
|
||||
if "platform_machine == 'aarch64'" not in _markers(p)
|
||||
]
|
||||
assert others, "no non-aarch64 torch entry found"
|
||||
for pkg in others:
|
||||
assert _registry(pkg) == "https://pypi.org/simple", (
|
||||
f"x86_64 torch moved to {_registry(pkg)!r} — HF Jobs would switch "
|
||||
"wheels."
|
||||
)
|
||||
assert "+cu" not in pkg["version"], "x86_64 torch must not be CUDA-pinned"
|
||||
|
||||
|
||||
def test_torch_version_identical_across_platforms():
|
||||
"""The fix changes only the wheel's SOURCE, not its version: the CUDA
|
||||
index carries newer builds than the PyPI pin, so a `>=` drags torch
|
||||
2.9.1 -> 2.13.0 with nothing having validated that bump."""
|
||||
versions = {p["version"].split("+")[0] for p in _packages("torch")}
|
||||
assert len(versions) == 1, f"torch versions diverge across platforms: {versions}"
|
||||
|
||||
|
||||
def _on_spark():
|
||||
return (
|
||||
sys.platform == "linux"
|
||||
and platform.machine() == "aarch64"
|
||||
and shutil.which("nvidia-smi") is not None
|
||||
and subprocess.run(["nvidia-smi"], capture_output=True).returncode == 0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not _on_spark(), reason="not a linux-aarch64 machine with a GPU")
|
||||
def test_installed_torch_actually_sees_the_gpu():
|
||||
"""Direct reproduction of the crash: this is exactly what select_gpus() reads."""
|
||||
import torch
|
||||
|
||||
assert torch.cuda.device_count() > 0, (
|
||||
f"torch {torch.__version__} (cuda={torch.version.cuda}) sees no GPU "
|
||||
"although nvidia-smi reports one -> select_gpus() will raise IndexError."
|
||||
)
|
||||
@@ -0,0 +1,90 @@
|
||||
import math
|
||||
import torch
|
||||
from mjlab_microduck.tasks import mdp
|
||||
|
||||
|
||||
def test_crouch_height_target_endpoints_are_high():
|
||||
# phase 0 (début) et phase ~1 (fin) → hauteur haute (debout)
|
||||
phase = torch.tensor([0.0, 0.999])
|
||||
t = mdp.crouch_height_target(phase, height_low=0.075, height_high=0.11)
|
||||
assert torch.allclose(t, torch.tensor([0.11, 0.11]), atol=2e-3)
|
||||
|
||||
|
||||
def test_crouch_height_target_plateau_is_low():
|
||||
# tout le palier [0.375, 0.625] → hauteur basse constante
|
||||
phase = torch.tensor([0.375, 0.5, 0.624])
|
||||
t = mdp.crouch_height_target(phase, height_low=0.075, height_high=0.11)
|
||||
assert torch.allclose(t, torch.full((3,), 0.075), atol=1e-6)
|
||||
|
||||
|
||||
def test_crouch_height_target_descent_midpoint():
|
||||
# milieu de la descente (phase = hold_lo/2 = 0.1875) → milieu des deux hauteurs
|
||||
phase = torch.tensor([0.1875])
|
||||
t = mdp.crouch_height_target(phase, height_low=0.075, height_high=0.11)
|
||||
assert torch.allclose(t, torch.tensor([(0.11 + 0.075) / 2]), atol=1e-6)
|
||||
|
||||
|
||||
def test_crouch_height_target_rise_midpoint():
|
||||
# milieu de la remontée (phase = 0.8125) → milieu des deux hauteurs
|
||||
phase = torch.tensor([0.8125])
|
||||
t = mdp.crouch_height_target(phase, height_low=0.075, height_high=0.11)
|
||||
assert torch.allclose(t, torch.tensor([(0.11 + 0.075) / 2]), atol=1e-6)
|
||||
|
||||
|
||||
# ── crouch_pose_blend : 4 segments (descente / bas / montée / debout) ─────────
|
||||
# breakpoints de test : descente [0,0.1), bas [0.1,0.5), montée [0.5,0.6),
|
||||
# debout [0.6,1.0).
|
||||
_BLEND = dict(descent_end=0.10, hold_end=0.50, rise_end=0.60)
|
||||
|
||||
|
||||
def test_blend_zero_standing_at_start_and_top_hold():
|
||||
phase = torch.tensor([0.0, 0.6, 0.8, 0.999]) # début + palier haut
|
||||
b = mdp.crouch_pose_blend(phase, **_BLEND)
|
||||
assert torch.allclose(b, torch.zeros(4), atol=1e-6)
|
||||
|
||||
|
||||
def test_blend_one_on_low_hold():
|
||||
phase = torch.tensor([0.10, 0.3, 0.499]) # palier bas
|
||||
b = mdp.crouch_pose_blend(phase, **_BLEND)
|
||||
assert torch.allclose(b, torch.ones(3), atol=1e-6)
|
||||
|
||||
|
||||
def test_blend_descent_and_rise_midpoints():
|
||||
# milieu descente (0.05 sur [0,0.1)) → 0.5 ; milieu montée (0.55 sur [0.5,0.6)) → 0.5
|
||||
phase = torch.tensor([0.05, 0.55])
|
||||
b = mdp.crouch_pose_blend(phase, **_BLEND)
|
||||
assert torch.allclose(b, torch.tensor([0.5, 0.5]), atol=1e-6)
|
||||
|
||||
|
||||
def test_reward_is_one_when_height_matches_target():
|
||||
# phase 0.5 (plein palier) → cible = height_low ; si com_height == height_low → reward 1
|
||||
cmd_cos = torch.tensor([math.cos(2 * math.pi * 0.5)]) # -1
|
||||
cmd_sin = torch.tensor([math.sin(2 * math.pi * 0.5)]) # ~0
|
||||
com_height = torch.tensor([0.075])
|
||||
r = mdp.crouch_glide_reward_from_values(
|
||||
com_height, cmd_cos, cmd_sin, height_low=0.075, height_high=0.11, std=0.02
|
||||
)
|
||||
assert torch.allclose(r, torch.tensor([1.0]), atol=1e-3)
|
||||
|
||||
|
||||
def test_reward_decays_when_off_by_one_std():
|
||||
# à height_low + std de la cible → exp(-1) ≈ 0.368
|
||||
cmd_cos = torch.tensor([math.cos(2 * math.pi * 0.5)])
|
||||
cmd_sin = torch.tensor([math.sin(2 * math.pi * 0.5)])
|
||||
com_height = torch.tensor([0.075 + 0.02])
|
||||
r = mdp.crouch_glide_reward_from_values(
|
||||
com_height, cmd_cos, cmd_sin, height_low=0.075, height_high=0.11, std=0.02
|
||||
)
|
||||
assert torch.allclose(r, torch.tensor([math.exp(-1.0)]), atol=1e-3)
|
||||
|
||||
|
||||
def test_reward_at_phase_zero_expects_high_stance():
|
||||
# phase 0 → cible = height_high ; rester debout est récompensé, être accroupi non
|
||||
cmd_cos = torch.tensor([1.0, 1.0]) # cos(0)
|
||||
cmd_sin = torch.tensor([0.0, 0.0]) # sin(0)
|
||||
com_height = torch.tensor([0.11, 0.075]) # debout vs accroupi
|
||||
r = mdp.crouch_glide_reward_from_values(
|
||||
com_height, cmd_cos, cmd_sin, height_low=0.075, height_high=0.11, std=0.02
|
||||
)
|
||||
assert r[0] > 0.99 # debout à phase 0 → ~1
|
||||
assert r[1] < 0.2 # accroupi à phase 0 → faible
|
||||
@@ -0,0 +1,48 @@
|
||||
"""descent_speed_reward : récompense la vitesse d'avance vers le bas de la pente
|
||||
(monde +x), plafonnée à `cap`, nulle si le robot recule/remonte, NaN-safe.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from mjlab_microduck.tasks.mdp import descent_speed_reward
|
||||
|
||||
|
||||
class _Data:
|
||||
def __init__(self, vx):
|
||||
self.root_link_lin_vel_w = torch.tensor(vx, dtype=torch.float32).reshape(-1, 1).repeat(1, 3)
|
||||
# seule la colonne 0 (x) est lue ; on met vx en x
|
||||
self.root_link_lin_vel_w[:, 0] = torch.tensor(vx, dtype=torch.float32)
|
||||
|
||||
|
||||
class _Asset:
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
|
||||
|
||||
class _Env:
|
||||
def __init__(self, vx):
|
||||
self._a = _Asset(_Data(vx))
|
||||
self.scene = self
|
||||
|
||||
def __getitem__(self, _k):
|
||||
return self._a
|
||||
|
||||
|
||||
def test_rewards_forward_speed_up_to_cap():
|
||||
out = descent_speed_reward(_Env([0.5]), cap=0.8)
|
||||
assert abs(float(out[0]) - 0.5) < 1e-6
|
||||
|
||||
|
||||
def test_caps_high_speed():
|
||||
out = descent_speed_reward(_Env([1.5]), cap=0.8)
|
||||
assert abs(float(out[0]) - 0.8) < 1e-6
|
||||
|
||||
|
||||
def test_zero_for_backward_or_uphill():
|
||||
out = descent_speed_reward(_Env([-0.4]), cap=0.8)
|
||||
assert float(out[0]) == 0.0
|
||||
|
||||
|
||||
def test_nan_safe():
|
||||
out = descent_speed_reward(_Env([float("nan")]), cap=0.8)
|
||||
assert float(out[0]) == 0.0
|
||||
@@ -0,0 +1,57 @@
|
||||
from mjlab_microduck.tasks.microduck_ground_pick_env_cfg import (
|
||||
make_microduck_ground_pick_env_cfg,
|
||||
)
|
||||
from mjlab_microduck.tasks.mdp import GroundPickPhaseCommand
|
||||
|
||||
|
||||
def test_ground_pick_cfg_task_space_rewards():
|
||||
"""Objectif espace-tâche : bouche près du sol (sans toucher) + orientée."""
|
||||
cfg = make_microduck_ground_pick_env_cfg()
|
||||
r = cfg.rewards
|
||||
# proximité bouche->sol (tire vers le bas)
|
||||
assert "mouth_ground_proximity" in r
|
||||
assert r["mouth_ground_proximity"].weight == 3.0
|
||||
assert r["mouth_ground_proximity"].params["target_height"] == 0.0
|
||||
# orientation bouche vers le bas
|
||||
assert "mouth_perpendicular_to_ground" in r
|
||||
assert r["mouth_perpendicular_to_ground"].weight == 2.0
|
||||
# no-touch : pénalité de contact forte + seuil bas
|
||||
assert "head_impact_penalty" in r
|
||||
assert r["head_impact_penalty"].weight == -2.0
|
||||
assert r["head_impact_penalty"].params["threshold"] == 1.0
|
||||
# pieds au sol ET à plat (anti-bascule sur la cheville)
|
||||
assert "feet_grounded" in r and r["feet_grounded"].weight == 3.0
|
||||
assert "feet_flat" in r and r["feet_flat"].weight == -2.0
|
||||
# retour debout + aide au relever (upright gaté sur la remontée)
|
||||
assert "ground_pick_return_pose_legs" in r
|
||||
assert "ground_pick_return_pose_neck" in r
|
||||
assert "return_upright" in r and r["return_upright"].weight == 4.0
|
||||
# plus d'approche par pose interpolée
|
||||
assert "phase_pose_track_head" not in r
|
||||
assert "phase_pose_track_legs" not in r
|
||||
|
||||
|
||||
def test_ground_pick_mouth_payload_wired():
|
||||
cfg = make_microduck_ground_pick_env_cfg()
|
||||
# hook d'application (poids 0) + event de tirage du payload
|
||||
assert "mouth_payload_force" in cfg.rewards
|
||||
assert cfg.rewards["mouth_payload_force"].weight == 0.0
|
||||
assert "sample_mouth_payload" in cfg.events
|
||||
assert cfg.events["sample_mouth_payload"].params["min_kg"] == 0.01
|
||||
assert cfg.events["sample_mouth_payload"].params["max_kg"] == 0.04
|
||||
|
||||
|
||||
def test_ground_pick_cfg_command_is_phase():
|
||||
cfg = make_microduck_ground_pick_env_cfg()
|
||||
cmd = cfg.commands["twist"]
|
||||
assert cmd.class_type is GroundPickPhaseCommand
|
||||
|
||||
|
||||
def test_ground_pick_rough_variant_builds():
|
||||
cfg = make_microduck_ground_pick_env_cfg(rough=True)
|
||||
assert "mouth_ground_proximity" in cfg.rewards
|
||||
|
||||
|
||||
def test_ground_pick_play_variant_builds():
|
||||
cfg = make_microduck_ground_pick_env_cfg(play=True)
|
||||
assert "mouth_ground_proximity" in cfg.rewards
|
||||
@@ -0,0 +1,123 @@
|
||||
import torch
|
||||
from mjlab_microduck.tasks.mdp import phase_pose_blend
|
||||
|
||||
DESCENT_END, HOLD_END, RISE_END = 0.15, 0.50, 0.65
|
||||
|
||||
|
||||
def test_phase_pose_blend_keypoints():
|
||||
phase = torch.tensor([0.0, 0.075, 0.15, 0.30, 0.50, 0.575, 0.65, 0.80])
|
||||
b = phase_pose_blend(phase, DESCENT_END, HOLD_END, RISE_END)
|
||||
expected = torch.tensor([0.0, 0.5, 1.0, 1.0, 1.0, 0.5, 0.0, 0.0])
|
||||
assert torch.allclose(b, expected, atol=1e-6), b
|
||||
|
||||
|
||||
def test_phase_pose_blend_range():
|
||||
phase = torch.linspace(0.0, 1.0, 101)
|
||||
b = phase_pose_blend(phase, DESCENT_END, HOLD_END, RISE_END)
|
||||
assert b.min() >= 0.0 and b.max() <= 1.0
|
||||
|
||||
|
||||
from mjlab_microduck.tasks.mdp import phase_pose_track, phase_pose_track_l1
|
||||
|
||||
|
||||
class _FakeData:
|
||||
def __init__(self, joint_pos, default_pos):
|
||||
self.joint_pos = joint_pos
|
||||
self.default_joint_pos = default_pos
|
||||
|
||||
|
||||
class _FakeAsset:
|
||||
def __init__(self, names, joint_pos, default_pos):
|
||||
self._ids = {n: i for i, n in enumerate(names)}
|
||||
self.data = _FakeData(joint_pos, default_pos)
|
||||
|
||||
def find_joints(self, query):
|
||||
# mjlab renvoie (ids, names) ; on ne gère que la requête [name]
|
||||
(name,) = query
|
||||
return ([self._ids[name]], [name])
|
||||
|
||||
|
||||
class _FakeCmdMgr:
|
||||
def __init__(self, cmd):
|
||||
self._cmd = cmd
|
||||
|
||||
def get_command(self, _name):
|
||||
return self._cmd
|
||||
|
||||
|
||||
class _FakeEnv:
|
||||
def __init__(self, names, joint_pos, default_pos, phase):
|
||||
import math
|
||||
self.device = "cpu"
|
||||
self.scene = {"robot": _FakeAsset(names, joint_pos, default_pos)}
|
||||
ang = 2 * math.pi * phase
|
||||
cmd = torch.tensor([[math.cos(ang), math.sin(ang), 0.0]])
|
||||
self.command_manager = _FakeCmdMgr(cmd)
|
||||
|
||||
|
||||
NAMES = ["j0", "j1"]
|
||||
DOWN = {"j0": 1.0, "j1": -1.0}
|
||||
# HOME (STAND source) = 0 pour les deux joints
|
||||
HOME = torch.tensor([[0.0, 0.0]])
|
||||
|
||||
|
||||
def _env(cur, phase):
|
||||
return _FakeEnv(NAMES, torch.tensor([cur]), HOME.clone(), phase)
|
||||
|
||||
|
||||
def test_phase_pose_track_perfect_at_down():
|
||||
# phase 0.30 -> blend 1 -> cible = DOWN ; cur == DOWN -> gaussienne 1, l1 0
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
cfg = SceneEntityCfg("robot")
|
||||
env = _env([1.0, -1.0], phase=0.30)
|
||||
r = phase_pose_track(env, target_pose=DOWN, asset_cfg=cfg)
|
||||
assert torch.allclose(r, torch.tensor([1.0]), atol=1e-6), r
|
||||
env2 = _env([1.0, -1.0], phase=0.30)
|
||||
l1 = phase_pose_track_l1(env2, target_pose=DOWN, asset_cfg=cfg)
|
||||
assert torch.allclose(l1, torch.tensor([0.0]), atol=1e-6), l1
|
||||
|
||||
|
||||
def test_phase_pose_track_l1_at_home_when_down_target():
|
||||
# phase 0.30 -> cible DOWN=[1,-1] ; cur=HOME=[0,0] -> l1 = -mean(|1|,|1|) = -1
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
cfg = SceneEntityCfg("robot")
|
||||
env = _env([0.0, 0.0], phase=0.30)
|
||||
l1 = phase_pose_track_l1(env, target_pose=DOWN, asset_cfg=cfg)
|
||||
assert torch.allclose(l1, torch.tensor([-1.0]), atol=1e-6), l1
|
||||
|
||||
|
||||
def test_phase_pose_track_returns_to_stand():
|
||||
# phase 0.80 -> blend 0 -> cible = HOME ; cur=HOME -> gaussienne 1
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
cfg = SceneEntityCfg("robot")
|
||||
env = _env([0.0, 0.0], phase=0.80)
|
||||
r = phase_pose_track(env, target_pose=DOWN, asset_cfg=cfg)
|
||||
assert torch.allclose(r, torch.tensor([1.0]), atol=1e-6), r
|
||||
|
||||
|
||||
def test_phase_pose_track_affine_interpolation_nonzero_home():
|
||||
# HOME (source) nonzero, blend 0.5 at phase 0.075:
|
||||
# target = home + 0.5*(down-home) = [0.4,-0.4] + 0.5*([1,-1]-[0.4,-0.4]) = [0.7,-0.7]
|
||||
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
||||
cfg = SceneEntityCfg("robot")
|
||||
home = torch.tensor([[0.4, -0.4]])
|
||||
env = _FakeEnv(NAMES, torch.tensor([[0.7, -0.7]]), home.clone(), phase=0.075)
|
||||
r = phase_pose_track(env, target_pose=DOWN, asset_cfg=cfg)
|
||||
assert torch.allclose(r, torch.tensor([1.0]), atol=1e-6), r
|
||||
|
||||
|
||||
def test_ground_pick_cmd_cfg_has_randomize_phase_default_true():
|
||||
from mjlab_microduck.tasks.mdp import GroundPickPhaseCommandCfg
|
||||
from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
|
||||
# construit une cfg minimale en copiant une cfg velocity par défaut
|
||||
# NOTE: adapté de `asset_name` (brief) -> `entity_name` (API locale de
|
||||
# UniformVelocityCommandCfg, qui n'a pas de champ `asset_name`).
|
||||
base = UniformVelocityCommandCfg(
|
||||
entity_name="robot", resampling_time_range=(10.0, 10.0),
|
||||
ranges=UniformVelocityCommandCfg.Ranges(
|
||||
lin_vel_x=(0.0, 0.0), lin_vel_y=(0.0, 0.0), ang_vel_z=(0.0, 0.0),
|
||||
),
|
||||
)
|
||||
cfg = GroundPickPhaseCommandCfg(**{**vars(base)})
|
||||
assert cfg.randomize_phase is True
|
||||
assert cfg.period == 4.0
|
||||
@@ -0,0 +1,176 @@
|
||||
"""head_pose_bias_penalty: prices sustained standing droop, never the recovery.
|
||||
|
||||
The velocity-env lesson (run 5yay13u4): instantaneous posture precision is an
|
||||
unescapable tax on motion. The standup lesson (retired head_impact_penalty):
|
||||
any head cost active during the ground phase blocks the head-pivot flip. So
|
||||
this term must (a) charge only the DC bias, (b) accumulate NOTHING while
|
||||
fallen, and (c) start the clock from ~zero on arrival upright.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
from mjlab_microduck.tasks import mdp as microduck_mdp
|
||||
|
||||
|
||||
class _Data:
|
||||
def __init__(self, n):
|
||||
self.joint_pos = torch.zeros(n, 6)
|
||||
self.default_joint_pos = torch.zeros(n, 6)
|
||||
self.root_link_pos_w = torch.zeros(n, 3)
|
||||
self.root_link_quat_w = torch.zeros(n, 4)
|
||||
self.root_link_quat_w[:, 0] = 1.0 # upright
|
||||
|
||||
|
||||
class _Asset:
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
|
||||
|
||||
class _Terrain:
|
||||
def __init__(self, n):
|
||||
self.env_origins = torch.zeros(n, 3)
|
||||
|
||||
|
||||
class _Scene:
|
||||
def __init__(self, asset, n):
|
||||
self._asset = asset
|
||||
self.terrain = _Terrain(n)
|
||||
|
||||
def __getitem__(self, _):
|
||||
return self._asset
|
||||
|
||||
|
||||
class _Cmd:
|
||||
def __init__(self, n, dim=4):
|
||||
self.cmd = torch.zeros(n, dim)
|
||||
|
||||
def get_command(self, _):
|
||||
return self.cmd
|
||||
|
||||
|
||||
class _Env:
|
||||
def __init__(self, n):
|
||||
self.num_envs = n
|
||||
self.device = "cpu"
|
||||
self.step_dt = 0.02
|
||||
self.episode_length_buf = torch.full((n,), 10, dtype=torch.long)
|
||||
self.scene = _Scene(_Asset(_Data(n)), n)
|
||||
self.command_manager = _Cmd(n)
|
||||
# Pre-seed the neck-id cache (normally built from the real asset).
|
||||
self._head_pose_neck_ids = torch.tensor([0, 1, 2, 3])
|
||||
self._head_pose_bl_ids = torch.tensor([0, 0, 0, 0])
|
||||
self._head_pose_bl_mask = torch.zeros(4)
|
||||
|
||||
|
||||
GATE = dict(
|
||||
gate_height_low=0.09, gate_height_high=0.11,
|
||||
gate_tilt_full_deg=20.0, gate_tilt_zero_deg=45.0,
|
||||
)
|
||||
|
||||
|
||||
def _set_pose(env, z, pitch_deg):
|
||||
env.scene._asset.data.root_link_pos_w[:, 2] = z
|
||||
half = math.radians(pitch_deg) / 2
|
||||
q = torch.tensor([math.cos(half), 0.0, math.sin(half), 0.0])
|
||||
env.scene._asset.data.root_link_quat_w[:] = q
|
||||
|
||||
|
||||
def _run(env, steps, **kw):
|
||||
out = None
|
||||
for _ in range(steps):
|
||||
out = microduck_mdp.head_pose_bias_penalty(env, tau_s=1.0, **kw)
|
||||
return out
|
||||
|
||||
|
||||
def test_prone_accumulates_nothing():
|
||||
env = _Env(2)
|
||||
_set_pose(env, z=0.05, pitch_deg=90.0) # face-down on the ground
|
||||
env.scene._asset.data.joint_pos[:, :4] = 0.5 # huge head "error" (28°)
|
||||
out = _run(env, 200, **GATE) # 4 s of prone thrash
|
||||
assert torch.allclose(out, torch.zeros(2), atol=1e-9), out
|
||||
|
||||
|
||||
def test_arrival_starts_from_zero_then_charges_true_bias():
|
||||
env = _Env(2)
|
||||
_set_pose(env, z=0.05, pitch_deg=90.0)
|
||||
env.scene._asset.data.joint_pos[:, :4] = 0.5
|
||||
_run(env, 200, **GATE) # ground phase: EMA stays 0
|
||||
_set_pose(env, z=0.117, pitch_deg=0.0) # recovery completes
|
||||
env.scene._asset.data.joint_pos[:, :4] = math.radians(15) # 15° droop
|
||||
first = _run(env, 1, **GATE)
|
||||
assert first.abs().max() < 0.01, f"finish-line wall: {first}" # no arrival spike
|
||||
settled = _run(env, 300, **GATE) # 6 s standing
|
||||
assert abs(-settled[0].item() - math.radians(15)) < 0.01 # charges the droop
|
||||
|
||||
|
||||
def test_fall_stops_the_charge_immediately():
|
||||
env = _Env(2)
|
||||
_set_pose(env, z=0.117, pitch_deg=0.0)
|
||||
env.scene._asset.data.joint_pos[:, :4] = math.radians(15)
|
||||
_run(env, 300, **GATE)
|
||||
_set_pose(env, z=0.05, pitch_deg=90.0) # falls
|
||||
out = _run(env, 1, **GATE)
|
||||
assert torch.allclose(out, torch.zeros(2), atol=1e-9), out
|
||||
|
||||
|
||||
def test_ungated_matches_velocity_env_behavior():
|
||||
# No gate params -> plain EMA of the raw error (the velocity-env term).
|
||||
env = _Env(2)
|
||||
_set_pose(env, z=0.05, pitch_deg=90.0) # pose must be irrelevant
|
||||
env.scene._asset.data.joint_pos[:, :4] = math.radians(15)
|
||||
out = _run(env, 300)
|
||||
assert abs(-out[0].item() - math.radians(15)) < 0.01
|
||||
|
||||
|
||||
def test_reset_clears_the_ema():
|
||||
env = _Env(2)
|
||||
_set_pose(env, z=0.117, pitch_deg=0.0)
|
||||
env.scene._asset.data.joint_pos[:, :4] = math.radians(15)
|
||||
_run(env, 300, **GATE)
|
||||
env.episode_length_buf[0] = 1 # env 0 just reset
|
||||
out = _run(env, 1, **GATE)
|
||||
# One step after reset the EMA holds exactly alpha*err (~0.005), while the
|
||||
# non-reset env still carries the full settled bias (~0.26).
|
||||
assert out[0].abs() < 0.01 and out[1].abs() > 0.2
|
||||
|
||||
|
||||
def test_standup_cfg_wiring():
|
||||
from mjlab_microduck.tasks.microduck_standup_env_cfg import (
|
||||
make_microduck_standup_env_cfg,
|
||||
)
|
||||
|
||||
cfg = make_microduck_standup_env_cfg()
|
||||
term = cfg.rewards["head_pose_bias"]
|
||||
assert term.weight == 0.0 # discovery phase untouched
|
||||
assert term.params["gate_height_low"] is not None
|
||||
# Gate values identical to arrival_damping so "standing" means one thing.
|
||||
ad = cfg.rewards["arrival_damping"].params
|
||||
assert term.params["gate_height_low"] == ad["height_low"]
|
||||
assert term.params["gate_tilt_full_deg"] == ad["tilt_full_deg"]
|
||||
stages = cfg.curriculum["head_pose_bias_weight"].params["weight_stages"]
|
||||
assert stages[0]["weight"] == 0.0
|
||||
assert min(s["step"] for s in stages if s["weight"] > 0) >= 3000 * 24
|
||||
|
||||
|
||||
def test_velocity_cfg_unchanged_no_gate():
|
||||
from mjlab_microduck.tasks.microduck_velocity_env_cfg import (
|
||||
make_microduck_velocity_env_cfg,
|
||||
)
|
||||
|
||||
cfg = make_microduck_velocity_env_cfg()
|
||||
assert "gate_height_low" not in cfg.rewards["head_pose_bias"].params
|
||||
|
||||
|
||||
def test_velstand_inherited_term_is_gated():
|
||||
# Velstand episodes survive falls — the inherited velocity-env term must not
|
||||
# charge the ground phase.
|
||||
from mjlab_microduck.tasks.microduck_velstand_env_cfg import (
|
||||
make_microduck_velstand_env_cfg,
|
||||
)
|
||||
|
||||
cfg = make_microduck_velstand_env_cfg()
|
||||
params = cfg.rewards["head_pose_bias"].params
|
||||
assert params.get("gate_height_low") is not None
|
||||
assert params["gate_tilt_zero_deg"] == 40.0 # REWARD_GATE_TILT_DEG
|
||||
@@ -0,0 +1,66 @@
|
||||
"""robot_state_is_nan doit attraper un état non-fini n'importe où (joints OU base
|
||||
OU roues), pas seulement dans joint_pos — sinon un free-joint qui diverge en NaN
|
||||
échappe au reset et corrompt l'obs critic (base_lin_vel/wheel_vel), ce qui tue
|
||||
l'entraînement via le check_nan global de rsl_rl.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from mjlab_microduck.tasks.mdp import robot_state_is_nan
|
||||
|
||||
|
||||
class _Data:
|
||||
def __init__(self, n):
|
||||
self.joint_pos = torch.zeros(n, 4)
|
||||
self.joint_vel = torch.zeros(n, 4)
|
||||
self.root_link_pos_w = torch.zeros(n, 3)
|
||||
self.root_link_quat_w = torch.zeros(n, 4)
|
||||
self.root_link_lin_vel_w = torch.zeros(n, 3)
|
||||
self.root_link_ang_vel_w = torch.zeros(n, 3)
|
||||
|
||||
|
||||
class _Asset:
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
|
||||
|
||||
class _Scene:
|
||||
def __init__(self, asset):
|
||||
self._a = asset
|
||||
|
||||
def __getitem__(self, _key):
|
||||
return self._a
|
||||
|
||||
|
||||
class _Env:
|
||||
def __init__(self, data):
|
||||
self.scene = _Scene(_Asset(data))
|
||||
|
||||
|
||||
def test_catches_base_linear_velocity_nan():
|
||||
# env 1 : vitesse de base NaN (free-joint divergé) — joint_pos reste fini.
|
||||
d = _Data(3)
|
||||
d.root_link_lin_vel_w[1, 0] = float("nan")
|
||||
out = robot_state_is_nan(_Env(d))
|
||||
assert out.tolist() == [False, True, False]
|
||||
|
||||
|
||||
def test_catches_base_velocity_inf():
|
||||
# inf dans la vitesse angulaire de base (avant qu'il ne devienne NaN).
|
||||
d = _Data(2)
|
||||
d.root_link_ang_vel_w[0, 2] = float("inf")
|
||||
out = robot_state_is_nan(_Env(d))
|
||||
assert out.tolist() == [True, False]
|
||||
|
||||
|
||||
def test_still_catches_joint_pos_nan():
|
||||
# comportement historique préservé.
|
||||
d = _Data(2)
|
||||
d.joint_pos[0, 1] = float("nan")
|
||||
out = robot_state_is_nan(_Env(d))
|
||||
assert out.tolist() == [True, False]
|
||||
|
||||
|
||||
def test_clean_state_is_not_flagged():
|
||||
out = robot_state_is_nan(_Env(_Data(4)))
|
||||
assert out.tolist() == [False, False, False, False]
|
||||
@@ -0,0 +1,137 @@
|
||||
"""The critic obs must survive a non-finite sensor reading.
|
||||
|
||||
Regression for the 2026-08-21 crash: rsl_rl's check_nan killed a
|
||||
Velocity2-Rough-Backlash run with "observation group 'critic' contains NaN".
|
||||
`nan_state` (robot_state_is_nan) only covered joint + root state, but the
|
||||
critic also carries three SENSOR-derived terms (raycast heights, contact
|
||||
air-time, contact forces). MuJoCo can return a non-finite contact force while
|
||||
the integrated robot state is still clean, so the env was never reset and the
|
||||
NaN reached the runner.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from mjlab_microduck.tasks import mdp as microduck_mdp
|
||||
|
||||
|
||||
class _SensorData:
|
||||
def __init__(self, force=None, heights=None):
|
||||
self.force = force
|
||||
self.heights = heights
|
||||
|
||||
|
||||
class _Sensor:
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
|
||||
|
||||
class _Scene:
|
||||
def __init__(self, sensors, asset):
|
||||
self.sensors = sensors
|
||||
self._asset = asset
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self.sensors[key] if key in self.sensors else self._asset
|
||||
|
||||
|
||||
class _AssetData:
|
||||
def __init__(self, n):
|
||||
self.joint_pos = torch.zeros(n, 4)
|
||||
self.joint_vel = torch.zeros(n, 4)
|
||||
self.root_link_pos_w = torch.zeros(n, 3)
|
||||
self.root_link_quat_w = torch.zeros(n, 4)
|
||||
self.root_link_lin_vel_w = torch.zeros(n, 3)
|
||||
self.root_link_ang_vel_w = torch.zeros(n, 3)
|
||||
|
||||
|
||||
class _Asset:
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
|
||||
|
||||
class _Env:
|
||||
def __init__(self, n, force):
|
||||
self.num_envs = n
|
||||
self.device = "cpu"
|
||||
asset = _Asset(_AssetData(n))
|
||||
self.scene = _Scene({"feet": _Sensor(_SensorData(force=force))}, asset)
|
||||
|
||||
|
||||
def _force(n, bad_env=None, value=float("nan")):
|
||||
f = torch.ones(n, 2, 3)
|
||||
if bad_env is not None:
|
||||
f[bad_env, 0, 0] = value
|
||||
return f
|
||||
|
||||
|
||||
def test_state_only_check_misses_bad_contact_force():
|
||||
# This is the gap that killed the run: robot state is clean, force is not.
|
||||
env = _Env(3, _force(3, bad_env=1))
|
||||
assert not microduck_mdp.robot_state_is_nan(env).any()
|
||||
|
||||
|
||||
def test_termination_catches_nan_contact_force():
|
||||
env = _Env(3, _force(3, bad_env=1))
|
||||
out = microduck_mdp.robot_state_is_nan(env, sensor_names=("feet",))
|
||||
assert out.tolist() == [False, True, False]
|
||||
|
||||
|
||||
def test_termination_catches_inf_contact_force():
|
||||
env = _Env(3, _force(3, bad_env=2, value=float("inf")))
|
||||
out = microduck_mdp.robot_state_is_nan(env, sensor_names=("feet",))
|
||||
assert out.tolist() == [False, False, True]
|
||||
|
||||
|
||||
def test_termination_ignores_missing_sensor():
|
||||
env = _Env(2, _force(2))
|
||||
assert not microduck_mdp.robot_state_is_nan(env, sensor_names=("nope",)).any()
|
||||
|
||||
|
||||
def test_finite_helper_sanitizes_nan_and_inf():
|
||||
x = torch.tensor([[1.0, float("nan"), float("inf"), float("-inf")]])
|
||||
out = microduck_mdp._finite(x)
|
||||
assert torch.isfinite(out).all()
|
||||
assert out[0, 0] == 1.0
|
||||
|
||||
|
||||
def test_safe_obs_wrappers_are_wired_into_the_critic():
|
||||
# Guards must actually be installed on the env cfg, not just exist.
|
||||
from mjlab_microduck.tasks.microduck_velocity_env_cfg import (
|
||||
make_microduck_velocity_env_cfg,
|
||||
)
|
||||
|
||||
cfg = make_microduck_velocity_env_cfg(rough=True)
|
||||
terms = cfg.observations["critic"].terms
|
||||
for name in ("foot_contact_forces", "foot_height", "foot_air_time"):
|
||||
assert terms[name].func.__name__.endswith("_safe"), (
|
||||
f"critic/{name} lost its NaN guard"
|
||||
)
|
||||
|
||||
|
||||
def test_nan_state_termination_watches_the_contact_sensor():
|
||||
from mjlab_microduck.tasks.microduck_velocity_env_cfg import (
|
||||
make_microduck_velocity_env_cfg,
|
||||
)
|
||||
|
||||
cfg = make_microduck_velocity_env_cfg(rough=True)
|
||||
params = cfg.terminations["nan_state"].params
|
||||
assert params.get("sensor_names"), "nan_state no longer watches contact forces"
|
||||
|
||||
|
||||
def test_standup_env_is_also_guarded():
|
||||
# The deployed standing policy trains on StandUp, which builds on mjlab's
|
||||
# base env (NOT the microduck velocity env) and therefore does not inherit
|
||||
# the guards wired there.
|
||||
from mjlab_microduck.tasks.microduck_standup_env_cfg import (
|
||||
make_microduck_standup_env_cfg,
|
||||
)
|
||||
|
||||
cfg = make_microduck_standup_env_cfg()
|
||||
terms = cfg.observations["critic"].terms
|
||||
for name in ("foot_contact_forces", "foot_air_time"):
|
||||
assert terms[name].func.__name__.endswith("_safe"), (
|
||||
f"standup critic/{name} lost its NaN guard"
|
||||
)
|
||||
assert cfg.terminations["nan_state"].params.get("sensor_names"), (
|
||||
"standup nan_state no longer watches contact forces"
|
||||
)
|
||||
@@ -0,0 +1,48 @@
|
||||
from mjlab_microduck.tasks.microduck_roller_crouch_env_cfg import (
|
||||
make_microduck_roller_crouch_env_cfg,
|
||||
)
|
||||
from mjlab_microduck.tasks import mdp as microduck_mdp
|
||||
|
||||
|
||||
def test_cfg_uses_phase_command():
|
||||
cfg = make_microduck_roller_crouch_env_cfg()
|
||||
cmd = cfg.commands["twist"]
|
||||
assert isinstance(cmd, microduck_mdp.GroundPickPhaseCommandCfg)
|
||||
# period must match --ground-pick-period at deploy
|
||||
# (0.5s down + 2s low + 0.5s up + 2s standing = 5s)
|
||||
assert cmd.period == 5.0
|
||||
# each episode starts standing (phase 0), matching the runtime trigger
|
||||
assert cmd.randomize_phase is False
|
||||
|
||||
|
||||
def test_cfg_has_crouch_and_forward_rewards():
|
||||
cfg = make_microduck_roller_crouch_env_cfg()
|
||||
# pose-based objective (standing<->crouch) + L1 bootstrap
|
||||
assert "crouch_glide_pose" in cfg.rewards
|
||||
assert "crouch_glide_pose_l1" in cfg.rewards
|
||||
assert "forward_speed" in cfg.rewards
|
||||
# léger penché avant pendant l'accroupi (cible positive = vers l'avant)
|
||||
assert "crouch_forward_lean" in cfg.rewards
|
||||
assert cfg.rewards["crouch_forward_lean"].params["target_pitch"] > 0.0
|
||||
# the crouch pose is carried by-name and includes the leg fold
|
||||
cp = cfg.rewards["crouch_glide_pose"].params["crouch_pose"]
|
||||
assert "left_knee" in cp and "right_knee" in cp
|
||||
# rewards de patinage actif retirées (pas de stride pendant le trick)
|
||||
for gone in ("braking", "skating_air_time", "single_support", "glide", "wheel_speed"):
|
||||
assert gone not in cfg.rewards
|
||||
|
||||
|
||||
def test_entry_velocity_applied_safely_via_reset_base():
|
||||
# Regression: entry momentum must be injected through reset_base's
|
||||
# velocity_range (reset_root_state_uniform sets it from the clean default
|
||||
# state), NOT via a mode="reset" push_by_setting_velocity event, which adds
|
||||
# to the current (possibly divergent) root velocity and blows the base
|
||||
# free-joint up to NaN. See the env cfg comment on ENTRY_VELOCITY_X.
|
||||
cfg = make_microduck_roller_crouch_env_cfg()
|
||||
# the buggy reset-push event must NOT exist
|
||||
assert "entry_velocity" not in cfg.events
|
||||
# forward entry velocity must be carried by reset_base, with a positive range
|
||||
vr = cfg.events["reset_base"].params.get("velocity_range")
|
||||
assert vr and "x" in vr
|
||||
lo, hi = vr["x"]
|
||||
assert lo > 0.0 and hi >= lo
|
||||
@@ -0,0 +1,105 @@
|
||||
from mjlab_microduck.tasks.microduck_roller_slope_env_cfg import (
|
||||
make_microduck_roller_slope_env_cfg,
|
||||
)
|
||||
from mjlab_microduck.tasks.slope_terrain import FlatRampTerrainCfg
|
||||
|
||||
|
||||
def test_terrain_is_flat_ramp_generator():
|
||||
cfg = make_microduck_roller_slope_env_cfg()
|
||||
assert cfg.scene.terrain.terrain_type == "generator"
|
||||
gen = cfg.scene.terrain.terrain_generator
|
||||
assert gen is not None and gen.curriculum is True
|
||||
assert any(isinstance(st, FlatRampTerrainCfg) for st in gen.sub_terrains.values())
|
||||
|
||||
|
||||
def test_command_is_neutralised():
|
||||
cfg = make_microduck_roller_slope_env_cfg()
|
||||
cmd = cfg.commands["twist"]
|
||||
assert cmd.rel_standing_envs == 1.0
|
||||
assert cmd.rel_heading_envs == 0.0
|
||||
assert cmd.ranges.lin_vel_x == (0.0, 0.0)
|
||||
assert cmd.ranges.lin_vel_y == (0.0, 0.0)
|
||||
if getattr(cmd.ranges, "ang_vel_z", None) is not None:
|
||||
assert cmd.ranges.ang_vel_z == (0.0, 0.0)
|
||||
|
||||
|
||||
def test_rolling_entry_no_base_push():
|
||||
# élan donné en ROULEMENT (reset_rolling_entry), pas en poussée de base
|
||||
# (base seule + roues immobiles = à-coup de patinage). Donc reset_base ne
|
||||
# met aucune vitesse de base.
|
||||
cfg = make_microduck_roller_slope_env_cfg()
|
||||
assert cfg.events["reset_base"].params["velocity_range"] == {}
|
||||
assert "reset_rolling_entry" in cfg.events
|
||||
lo, hi = cfg.events["reset_rolling_entry"].params["speed_range"]
|
||||
assert 0.0 < lo <= hi <= 0.6
|
||||
|
||||
|
||||
def test_has_heading_hold_reward():
|
||||
# aller droit : maintien du yaw de spawn
|
||||
cfg = make_microduck_roller_slope_env_cfg()
|
||||
assert "heading_hold" in cfg.rewards
|
||||
assert cfg.rewards["heading_hold"].weight > 0.0
|
||||
|
||||
|
||||
def test_balance_rewards_no_fixed_pose():
|
||||
# équilibre libre : upright/alive/glisse présents, mais PAS de pose fixe
|
||||
# imposée (il doit pouvoir bouger son centre de gravité pour tenir la pente).
|
||||
cfg = make_microduck_roller_slope_env_cfg()
|
||||
for name in ("upright", "alive", "feet_flat", "wheel_glide", "neck_joint_pos_l2"):
|
||||
assert name in cfg.rewards
|
||||
assert "standing_pose" not in cfg.rewards
|
||||
assert "standing_pose_l1" not in cfg.rewards
|
||||
|
||||
|
||||
def test_has_wheel_glide_reward_not_base_speed():
|
||||
# "se laisser glisser" = rouler (roues), pas récompenser la vitesse de base
|
||||
# (qu'il atteignait en courant). wheel_glide présent, descent_speed absent.
|
||||
cfg = make_microduck_roller_slope_env_cfg()
|
||||
assert "wheel_glide" in cfg.rewards
|
||||
assert cfg.rewards["wheel_glide"].weight > 0.0
|
||||
assert "descent_speed" not in cfg.rewards
|
||||
|
||||
|
||||
def test_no_roller_skating_rewards_survive():
|
||||
# les rewards de PATINAGE du roller ne doivent pas survivre (heading_hold est
|
||||
# ré-ajouté volontairement pour aller droit, donc pas dans cette liste).
|
||||
cfg = make_microduck_roller_slope_env_cfg()
|
||||
for name in ("wheel_speed", "braking", "skating_air_time", "glide", "forward_lean"):
|
||||
assert name not in cfg.rewards
|
||||
|
||||
|
||||
def test_spawn_yaw_faces_downhill():
|
||||
# yaw fixe à 0 : toujours face au bas de la pente (+x), pas le -pi/+pi hérité
|
||||
cfg = make_microduck_roller_slope_env_cfg()
|
||||
assert cfg.events["reset_base"].params["pose_range"]["yaw"] == (0.0, 0.0)
|
||||
|
||||
|
||||
def test_void_termination_present_no_edge_termination():
|
||||
cfg = make_microduck_roller_slope_env_cfg()
|
||||
assert "fell_into_void" in cfg.terminations
|
||||
assert "fell_over" in cfg.terminations
|
||||
# plus de terminaison « bord de terrain » (remplacée par le plat de sortie)
|
||||
assert "reached_bottom" not in cfg.terminations
|
||||
assert "out_of_terrain_bounds" not in cfg.terminations
|
||||
|
||||
|
||||
def test_obs_nan_policy_sanitize():
|
||||
# obs assainies : un NaN de contact rare ne doit pas tuer l'entraînement
|
||||
cfg = make_microduck_roller_slope_env_cfg()
|
||||
assert cfg.observations["actor"].nan_policy == "sanitize"
|
||||
assert cfg.observations["critic"].nan_policy == "sanitize"
|
||||
|
||||
|
||||
def test_curriculum_present_and_starts_gentle():
|
||||
# curriculum doux->raide : démarre sur la rampe la plus douce, promotion active
|
||||
cfg = make_microduck_roller_slope_env_cfg() # play=False (entraînement)
|
||||
assert "terrain_levels" in cfg.curriculum
|
||||
assert cfg.scene.terrain.max_init_terrain_level == 0
|
||||
|
||||
|
||||
def test_terrain_tile_fits_geometry():
|
||||
# la tuile doit contenir plat + rampe_max + sortie
|
||||
cfg = make_microduck_roller_slope_env_cfg()
|
||||
gen = cfg.scene.terrain.terrain_generator
|
||||
st = next(iter(gen.sub_terrains.values()))
|
||||
assert st.flat_length + st.ramp_length_range[1] + st.runout_length <= gen.size[0]
|
||||
@@ -0,0 +1,539 @@
|
||||
import pytest
|
||||
|
||||
from mjlab_microduck.tasks.microduck_roller_standup_env_cfg import (
|
||||
EPISODE_LENGTH_S,
|
||||
make_microduck_roller_standup_env_cfg,
|
||||
)
|
||||
from mjlab_microduck.tasks.microduck_velocity_rollers_env_cfg import (
|
||||
make_microduck_velocity_rollers_env_cfg,
|
||||
)
|
||||
|
||||
# Récompenses de PATINAGE : elles ne doivent pas survivre dans un env de relevé.
|
||||
SKATING_REWARDS = (
|
||||
"wheel_speed",
|
||||
"braking",
|
||||
"skating_air_time",
|
||||
"glide",
|
||||
"single_support",
|
||||
"gait_symmetry",
|
||||
"forward_lean",
|
||||
"heading_hold",
|
||||
"feet_flat",
|
||||
"hip_roll_neutral",
|
||||
"pose",
|
||||
"com_height_target",
|
||||
"upright",
|
||||
)
|
||||
|
||||
|
||||
def test_env_builds_train_and_play():
|
||||
assert make_microduck_roller_standup_env_cfg() is not None
|
||||
assert make_microduck_roller_standup_env_cfg(play=True) is not None
|
||||
|
||||
|
||||
def test_episode_is_short():
|
||||
# Épisode court : monter puis stabiliser, comme standup (6 s).
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
assert cfg.episode_length_s == EPISODE_LENGTH_S == 6.0
|
||||
|
||||
|
||||
def test_no_skating_rewards_survive():
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
for name in SKATING_REWARDS:
|
||||
assert name not in cfg.rewards, f"reward de patinage survivante : {name}"
|
||||
|
||||
|
||||
def test_smoothness_regularisers_kept():
|
||||
# Gardées de l'héritage roller : le relevé a besoin de douceur sim2real, mais
|
||||
# body_ang_vel doit rester LÉGER (standup documente qu'à -0.15 il gelait).
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
for name in (
|
||||
"action_over_limit",
|
||||
"self_collisions",
|
||||
"body_ang_vel",
|
||||
"angular_momentum",
|
||||
"action_rate_l2",
|
||||
"neck_action_rate_l2",
|
||||
"neck_joint_pos_l2",
|
||||
"joint_torques_l2",
|
||||
):
|
||||
assert name in cfg.rewards, f"régularisateur perdu : {name}"
|
||||
assert cfg.rewards["body_ang_vel"].weight == -0.05
|
||||
|
||||
|
||||
def test_twist_command_is_neutralised():
|
||||
# Pas de pilotage : la policy se déploie en --standing, où le runtime laisse
|
||||
# le slot twist à zéro (cf. infer_policy.py:239).
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
cmd = cfg.commands["twist"]
|
||||
assert cmd.ranges.lin_vel_x == (-0.01, 0.01)
|
||||
assert cmd.ranges.lin_vel_y == (-0.01, 0.01)
|
||||
assert cmd.ranges.ang_vel_z == (-0.05, 0.05)
|
||||
assert cmd.heading_command is False
|
||||
assert cmd.ranges.heading is None
|
||||
assert cmd.rel_standing_envs == 0.0
|
||||
|
||||
|
||||
def test_twist_command_is_not_heading_relative():
|
||||
# L'env roller installe un RelativeHeadingVelocityCommandCfg (cmd[2] = erreur
|
||||
# de cap, calculée en interne). Ici cmd[2] doit être un vrai zéro bruité.
|
||||
from mjlab_microduck.tasks import mdp as microduck_mdp
|
||||
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
cmd = cfg.commands["twist"]
|
||||
assert isinstance(cmd, microduck_mdp.VelocityCommandCommandOnlyCfg)
|
||||
assert not isinstance(cmd, microduck_mdp.RelativeHeadingVelocityCommandCfg)
|
||||
|
||||
|
||||
def test_obs_nan_policy_sanitize():
|
||||
# Un contact rare fait diverger le free-joint en NaN : on assainit l'obs
|
||||
# plutôt que de tuer l'entraînement (même choix que roller_slope).
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
assert cfg.observations["actor"].nan_policy == "sanitize"
|
||||
assert cfg.observations["critic"].nan_policy == "sanitize"
|
||||
|
||||
|
||||
def test_obs_parity_with_roller_env():
|
||||
# Parité 61D obligatoire : sinon l'ONNX ne se charge pas dans un slot runtime.
|
||||
standup = make_microduck_roller_standup_env_cfg()
|
||||
roller = make_microduck_velocity_rollers_env_cfg()
|
||||
for grp in ("actor", "critic"):
|
||||
assert list(standup.observations[grp].terms.keys()) == list(
|
||||
roller.observations[grp].terms.keys()
|
||||
), f"layout d'observation divergent sur le groupe {grp}"
|
||||
|
||||
|
||||
def test_terrain_is_plain_plane():
|
||||
# Hérité de l'env roller : sol plat, pas de générateur. Pas de variante rough
|
||||
# pour cette v1.
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
assert cfg.scene.terrain.terrain_type == "plane"
|
||||
assert cfg.scene.terrain.terrain_generator is None
|
||||
|
||||
|
||||
def test_task_is_registered():
|
||||
from mjlab.tasks.registry import list_tasks
|
||||
|
||||
import mjlab_microduck.tasks # noqa: F401 (l'import déclenche l'enregistrement)
|
||||
|
||||
assert "Mjlab-RollerStandUp-Flat-MicroDuck" in list_tasks()
|
||||
|
||||
|
||||
def test_joint_indices_match_actual_roller_model():
|
||||
"""Verrou : les roues passives sont intercalées dans l'ordre des joints.
|
||||
|
||||
Réutiliser les indices du standup ([0-4, 9-13]) donnerait des récompenses
|
||||
qui pointent sur des roues. Ce test compile le vrai MjSpec du robot rollers
|
||||
et vérifie les noms aux indices utilisés. Pur CPU, pas de sim.
|
||||
"""
|
||||
import mujoco
|
||||
|
||||
from mjlab_microduck.robot.microduck_constants import get_walk_rollers_spec
|
||||
from mjlab_microduck.tasks.microduck_roller_standup_env_cfg import (
|
||||
_LEG_JOINTS,
|
||||
_NECK_JOINTS,
|
||||
_WHEEL_JOINTS,
|
||||
)
|
||||
|
||||
model = get_walk_rollers_spec().compile()
|
||||
articulated = [
|
||||
mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, j)
|
||||
for j in range(model.njnt)
|
||||
if model.jnt_type[j] != mujoco.mjtJoint.mjJNT_FREE
|
||||
]
|
||||
|
||||
assert [articulated[i] for i in _LEG_JOINTS] == [
|
||||
"left_hip_yaw", "left_hip_roll", "left_hip_pitch", "left_knee", "left_ankle",
|
||||
"right_hip_yaw", "right_hip_roll", "right_hip_pitch", "right_knee", "right_ankle",
|
||||
]
|
||||
assert [articulated[i] for i in _NECK_JOINTS] == [
|
||||
"neck_pitch", "head_pitch", "head_yaw", "head_roll",
|
||||
]
|
||||
assert [articulated[i] for i in _WHEEL_JOINTS] == [
|
||||
"passive_LF_wheel", "passive_LR_wheel", "passive_RF_wheel", "passive_RR_wheel",
|
||||
]
|
||||
# Aucun recouvrement, et les trois listes couvrent tous les joints.
|
||||
assert len(set(_LEG_JOINTS) | set(_NECK_JOINTS) | set(_WHEEL_JOINTS)) == len(articulated)
|
||||
|
||||
|
||||
def test_recovery_rewards_present_with_expected_weights():
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
expected = {
|
||||
"pose_stand_legs": 8.0,
|
||||
"pose_stand_l1": 5.0,
|
||||
"height_stand": 4.0,
|
||||
"height_stand_sharp": 4.0,
|
||||
"height_stand_l1": 30.0,
|
||||
"com_upward_velocity": 3.0,
|
||||
# gentle_rise : poids POSITIF. trunk_vertical_accel_penalty renvoie déjà
|
||||
# -|a_z|, donc un poids négatif en faisait une RÉCOMPENSE de la violence
|
||||
# (bug mesuré : Episode_Reward/gentle_rise loggée à +0.0118).
|
||||
"gentle_rise": +0.02,
|
||||
"upright_linear": 6.0,
|
||||
"upright_sharp": 6.0,
|
||||
"standing_composite": 15.0,
|
||||
# -2e-3 ne contribuait que -0.0002/pas face à +41.6 de tâche : nul.
|
||||
# -2.0 a mesuré -0.255/pas (run d8rnko6p) — pas le gel, mais on redescend
|
||||
# à -0.2 pour dégager le budget d'amortissement pendant qu'on isole.
|
||||
"joint_torque_rate_l2": -0.2,
|
||||
}
|
||||
for name, weight in expected.items():
|
||||
assert name in cfg.rewards, f"récompense de relevé manquante : {name}"
|
||||
assert cfg.rewards[name].weight == weight, f"poids inattendu sur {name}"
|
||||
|
||||
|
||||
def test_recovery_rewards_use_roller_heights_not_walker_heights():
|
||||
from mjlab_microduck.tasks.microduck_roller_standup_env_cfg import (
|
||||
ROLLER_PRONE_Z,
|
||||
ROLLER_STAND_Z,
|
||||
)
|
||||
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
assert ROLLER_STAND_Z == 0.138 # PAS le 0.115 du modèle sans roues
|
||||
for name in ("height_stand", "height_stand_sharp", "height_stand_l1"):
|
||||
assert cfg.rewards[name].params["target_height"] == ROLLER_STAND_Z
|
||||
assert cfg.rewards["standing_composite"].params["target_height"] == ROLLER_STAND_Z
|
||||
# com_upward_velocity se coupe juste AU-DESSUS de la cible (10 mm de marge),
|
||||
# sinon la policy se gare à l'altitude de coupure sans finir la montée.
|
||||
assert cfg.rewards["com_upward_velocity"].params["max_height"] == ROLLER_STAND_Z + 0.010
|
||||
# upright_sharp est gatée entre le repos au sol et la station debout.
|
||||
assert cfg.rewards["upright_sharp"].params["height_low"] == ROLLER_PRONE_Z
|
||||
assert cfg.rewards["upright_sharp"].params["height_high"] == ROLLER_STAND_Z
|
||||
|
||||
|
||||
def test_pose_rewards_target_legs_only_at_roller_indices():
|
||||
from mjlab_microduck.tasks.microduck_roller_standup_env_cfg import _LEG_JOINTS
|
||||
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
for name in ("pose_stand_legs", "pose_stand_l1", "standing_composite"):
|
||||
assert cfg.rewards[name].params["joint_indices"] == _LEG_JOINTS
|
||||
# target_overrides=None → la cible est HOME (default_joint_pos).
|
||||
assert cfg.rewards[name].params["target_overrides"] is None
|
||||
|
||||
|
||||
def test_trunk_asset_cfgs_are_distinct_objects():
|
||||
"""mjlab résout et MUTE les SceneEntityCfg en place : un objet partagé entre
|
||||
plusieurs termes provoque des indices périmés. Chaque terme doit avoir le sien.
|
||||
"""
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
names = (
|
||||
"height_stand", "height_stand_sharp", "height_stand_l1",
|
||||
"com_upward_velocity", "gentle_rise", "upright_linear",
|
||||
"upright_sharp", "standing_composite",
|
||||
)
|
||||
seen = [id(cfg.rewards[n].params["asset_cfg"]) for n in names]
|
||||
assert len(set(seen)) == len(seen), "asset_cfg partagé entre plusieurs termes"
|
||||
|
||||
|
||||
def test_starts_from_ground_states():
|
||||
# Ventre + dos + debout. Pas de bucket "assis" : il n'existait dans standup
|
||||
# que pour le hand-off depuis la policy sit, dont il n'y a pas d'équivalent
|
||||
# roller — et ses sitting_joint_overrides sont des indices du modèle SANS roues.
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
assert "set_ground_state" in cfg.events
|
||||
params = cfg.events["set_ground_state"].params
|
||||
assert params["sitting_prob"] == 0.0
|
||||
assert params["sitting_joint_overrides"] is None
|
||||
assert params["face_down_prob"] > 0.0
|
||||
assert params["standing_prob"] > 0.0
|
||||
# face_up (le dos) démarre à 0 : introduit tard par le curriculum.
|
||||
assert params["face_up_prob"] == 0.0
|
||||
|
||||
|
||||
def test_ground_state_heights_are_roller_specific():
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
params = cfg.events["set_ground_state"].params
|
||||
# Ventre et dos partagent une seule plage de z, mais leurs contacts diffèrent :
|
||||
# le ventre ne décolle du sol qu'à partir de 0.0752, le dos repose à 0.0475.
|
||||
# prone_z_min = 0.076 pour éliminer toute interpénétration côté ventre.
|
||||
assert (params["prone_z_min"], params["prone_z_max"]) == (0.076, 0.09)
|
||||
# Sous 0.0752 (contact mesuré, pose HOME), le départ ventre commence DANS le
|
||||
# sol — un pushout de contact que la policy paierait via gentle_rise /
|
||||
# joint_torque_rate_l2. prone_z_min doit rester au-dessus.
|
||||
assert params["prone_z_min"] >= 0.0752
|
||||
# Debout : hauteur ROLLER (+23 mm vs le modèle sans roues, qui est à 0.11–0.12).
|
||||
assert params["standing_z_min"] == 0.134
|
||||
assert params["standing_z_max"] == 0.144
|
||||
assert params["standing_z_min"] < 0.138 < params["standing_z_max"]
|
||||
|
||||
|
||||
def test_ground_state_event_runs_after_base_reset():
|
||||
# set_ground_state écrase la pose posée par reset_base / reset_robot_joints :
|
||||
# l'ordre des événements suit l'ordre d'insertion, il doit donc venir APRÈS.
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
order = list(cfg.events.keys())
|
||||
assert order.index("set_ground_state") > order.index("reset_base")
|
||||
assert order.index("set_ground_state") > order.index("reset_robot_joints")
|
||||
|
||||
|
||||
def test_no_fall_termination():
|
||||
# Le robot DÉMARRE tombé : une terminaison sur inclinaison tuerait l'épisode
|
||||
# au premier pas. nan_state (hérité) reste, lui.
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
assert "fell_over" not in cfg.terminations
|
||||
assert "nan_state" in cfg.terminations
|
||||
|
||||
|
||||
def test_ground_state_curriculum_ramps_easy_to_hard():
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
assert "ground_state_mix" in cfg.curriculum
|
||||
stages = cfg.curriculum["ground_state_mix"].params["param_stages"]
|
||||
assert cfg.curriculum["ground_state_mix"].params["event_name"] == "set_ground_state"
|
||||
# Les steps sont croissants et démarrent à 0.
|
||||
steps = [s["step"] for s in stages]
|
||||
assert steps[0] == 0 and steps == sorted(steps) and len(set(steps)) == len(steps)
|
||||
# Le dos (face_up) est introduit tard puis croît de façon monotone.
|
||||
face_up = [s["params"]["face_up_prob"] for s in stages]
|
||||
assert face_up[0] == 0.0
|
||||
assert face_up == sorted(face_up)
|
||||
assert face_up[-1] >= 0.35
|
||||
# Chaque palier est une distribution valide, et le "déjà debout" ne disparaît
|
||||
# jamais (sinon la policy se relève puis retombe faute d'apprendre à tenir).
|
||||
for stage in stages:
|
||||
p = stage["params"]
|
||||
total = (
|
||||
p["standing_prob"] + p["sitting_prob"]
|
||||
+ p["face_down_prob"] + p["face_up_prob"]
|
||||
)
|
||||
assert abs(total - 1.0) < 1e-9
|
||||
assert p["sitting_prob"] == 0.0
|
||||
assert p["standing_prob"] > 0.0
|
||||
|
||||
|
||||
def test_wheel_friction_curriculum_is_decreasing():
|
||||
"""La pièce nouvelle : roues FREINÉES → LIBRES.
|
||||
|
||||
Les roues roulent, donc il n'y a aucune adhérence longitudinale pour pousser
|
||||
sur le sol. On bootstrappe avec des roulements quasi bloqués (le relevé se
|
||||
fait comme avec des pieds) puis on rampe vers la vraie valeur. L'env roller,
|
||||
lui, fait MONTER cette friction (0 → 0.0015) : le sens est bien inversé ici.
|
||||
"""
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
stages = cfg.curriculum["wheel_friction"].params["ranges_stages"]
|
||||
assert cfg.curriculum["wheel_friction"].params["event_name"] == "randomize_wheel_friction"
|
||||
|
||||
steps = [s["step"] for s in stages]
|
||||
assert steps[0] == 0 and steps == sorted(steps) and len(set(steps)) == len(steps)
|
||||
|
||||
lows = [s["ranges"][0] for s in stages]
|
||||
assert lows == sorted(lows, reverse=True), "la friction doit DÉCROÎTRE"
|
||||
assert lows[0] >= 0.02, "départ franchement freiné pour bootstrapper le geste"
|
||||
# Arrivée sur la vraie valeur du roulement (celle de l'env roller).
|
||||
assert stages[-1]["ranges"] == (0.0015, 0.0015)
|
||||
for stage in stages:
|
||||
assert stage["ranges"][0] == stage["ranges"][1]
|
||||
|
||||
|
||||
def test_wheel_friction_event_default_matches_stage_zero():
|
||||
# Le curriculum manager tourne AVANT les événements de reset à chaque reset
|
||||
# (y compris le tout premier), et wheel_friction_curriculum défaut lui-même
|
||||
# sur le palier 0 : cette valeur par défaut de l'événement n'est donc jamais
|
||||
# lue en pratique. On vérifie juste qu'elle reste cohérente avec le palier 0
|
||||
# du curriculum — redondance défensive utile si le curriculum disparaît un
|
||||
# jour en laissant l'événement en place.
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
stage0 = cfg.curriculum["wheel_friction"].params["ranges_stages"][0]["ranges"]
|
||||
assert cfg.events["randomize_wheel_friction"].params["ranges"] == stage0
|
||||
|
||||
|
||||
def test_action_rate_ramp_is_the_standup_one_not_the_roller_one():
|
||||
# L'env roller monte à -2.0 (gait calme) : c'est un bloqueur de mouvement,
|
||||
# il ralentit l'action rapide dont le relevé depuis le dos a besoin. On
|
||||
# reprend la rampe du standup, qui plafonne à -1.0.
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
weights = [
|
||||
s["weight"] for s in cfg.curriculum["action_rate_weight"].params["weight_stages"]
|
||||
]
|
||||
assert weights == [-0.4, -0.8, -1.0]
|
||||
assert cfg.rewards["action_rate_l2"].weight == -0.6
|
||||
|
||||
|
||||
def test_push_curriculum_ramps_from_zero():
|
||||
# Poussées héritées (±0.2 m/s), mais rampées : une bourrade dès le pas 0
|
||||
# parasite le bootstrap du relevé.
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
assert "push_robot" in cfg.events
|
||||
stages = cfg.curriculum["push_magnitude"].params["push_stages"]
|
||||
assert cfg.curriculum["push_magnitude"].params["event_name"] == "push_robot"
|
||||
assert stages[0]["velocity_range"]["x"] == (0.0, 0.0)
|
||||
assert stages[-1]["velocity_range"]["x"] == (-0.2, 0.2)
|
||||
highs = [s["velocity_range"]["x"][1] for s in stages]
|
||||
assert highs == sorted(highs), "la poussée doit CROÎTRE"
|
||||
|
||||
|
||||
def test_inherited_dr_curricula_survive():
|
||||
# La DR héritée de l'env roller ne doit pas avoir été perdue en chemin.
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
for name in ("com_range", "head_com_range"):
|
||||
assert name in cfg.curriculum, f"curriculum de DR perdu : {name}"
|
||||
for name in (
|
||||
"randomize_com",
|
||||
"randomize_head_com",
|
||||
"randomize_armature",
|
||||
"randomize_joint_friction",
|
||||
"randomize_mass_inertia",
|
||||
"randomize_wheel_friction",
|
||||
"encoder_bias",
|
||||
):
|
||||
assert name in cfg.events, f"événement de DR perdu : {name}"
|
||||
|
||||
|
||||
# ── Override de play : forcer les départs sur le dos ──────────────────────────
|
||||
# Sans override, un play ne montre JAMAIS de départ sur le dos : l'env de play est
|
||||
# reconstruit à neuf, donc common_step_counter repart à 0 et le curriculum applique
|
||||
# son palier 0, où face_up_prob = 0. Or c'est justement le cas le plus dur, celui
|
||||
# qu'on veut inspecter à l'œil. STANDUP_PLAY_FACE_UP force le mélange, sur le
|
||||
# modèle de SLOPE_PLAY_DIFFICULTY dans roller_slope.
|
||||
|
||||
|
||||
def test_play_face_up_override_forces_back_starts(monkeypatch):
|
||||
monkeypatch.setenv("STANDUP_PLAY_FACE_UP", "1.0")
|
||||
cfg = make_microduck_roller_standup_env_cfg(play=True)
|
||||
params = cfg.events["set_ground_state"].params
|
||||
assert params["face_up_prob"] == 1.0
|
||||
assert params["face_down_prob"] == 0.0
|
||||
assert params["standing_prob"] == 0.0
|
||||
# Sans ça, le curriculum réécrirait les probabilités dès le premier reset
|
||||
# (event_param_curriculum tourne AVANT les événements de reset).
|
||||
assert "ground_state_mix" not in cfg.curriculum
|
||||
|
||||
|
||||
def test_play_face_up_override_splits_remainder_like_final_stage(monkeypatch):
|
||||
# 0.4 doit reproduire le DERNIER palier du curriculum (0.40 ventre / 0.20
|
||||
# debout / 0.40 dos) : le reste est réparti dans le rapport 2:1 de ce palier.
|
||||
monkeypatch.setenv("STANDUP_PLAY_FACE_UP", "0.4")
|
||||
params = make_microduck_roller_standup_env_cfg(play=True).events["set_ground_state"].params
|
||||
assert params["face_up_prob"] == pytest.approx(0.40)
|
||||
assert params["face_down_prob"] == pytest.approx(0.40)
|
||||
assert params["standing_prob"] == pytest.approx(0.20)
|
||||
total = params["face_up_prob"] + params["face_down_prob"] + params["standing_prob"]
|
||||
assert total == pytest.approx(1.0)
|
||||
|
||||
|
||||
def test_play_face_up_override_is_clamped(monkeypatch):
|
||||
monkeypatch.setenv("STANDUP_PLAY_FACE_UP", "3.0")
|
||||
params = make_microduck_roller_standup_env_cfg(play=True).events["set_ground_state"].params
|
||||
assert params["face_up_prob"] == 1.0
|
||||
|
||||
|
||||
def test_play_face_up_override_ignored_during_training(monkeypatch):
|
||||
# Garde-fou : la variable ne doit JAMAIS toucher l'entraînement, sinon on
|
||||
# casserait le curriculum easy->hard sans s'en apercevoir.
|
||||
monkeypatch.setenv("STANDUP_PLAY_FACE_UP", "1.0")
|
||||
cfg = make_microduck_roller_standup_env_cfg(play=False)
|
||||
assert cfg.events["set_ground_state"].params["face_up_prob"] == 0.00
|
||||
assert "ground_state_mix" in cfg.curriculum
|
||||
|
||||
|
||||
def test_play_without_override_keeps_curriculum_mix(monkeypatch):
|
||||
# Comportement par défaut inchangé : palier 0, pas de départ sur le dos.
|
||||
monkeypatch.delenv("STANDUP_PLAY_FACE_UP", raising=False)
|
||||
cfg = make_microduck_roller_standup_env_cfg(play=True)
|
||||
assert cfg.events["set_ground_state"].params["face_up_prob"] == 0.00
|
||||
assert "ground_state_mix" in cfg.curriculum
|
||||
|
||||
|
||||
def test_play_face_up_override_invalid_value_falls_back(monkeypatch):
|
||||
monkeypatch.setenv("STANDUP_PLAY_FACE_UP", "pouet")
|
||||
cfg = make_microduck_roller_standup_env_cfg(play=True)
|
||||
assert cfg.events["set_ground_state"].params["face_up_prob"] == 0.00
|
||||
assert "ground_state_mix" in cfg.curriculum
|
||||
|
||||
|
||||
def test_play_face_up_override_none_keyword_disables(monkeypatch):
|
||||
monkeypatch.setenv("STANDUP_PLAY_FACE_UP", "none")
|
||||
cfg = make_microduck_roller_standup_env_cfg(play=True)
|
||||
assert cfg.events["set_ground_state"].params["face_up_prob"] == 0.00
|
||||
assert "ground_state_mix" in cfg.curriculum
|
||||
|
||||
|
||||
# ── Anti-violence : corrections après test sur le robot ───────────────────────
|
||||
# Symptômes observés (checkpoint 4000+, EN SIMU AUSSI donc pas du sim2real) :
|
||||
# mouvements très brusques, la tête tape le sol, échec du relevé depuis le dos
|
||||
# sur le vrai robot. Diagnostic mesuré dans wandb (run vweolw91, iter 7500).
|
||||
|
||||
|
||||
def test_already_negative_penalties_use_positive_weights():
|
||||
"""Verrou sur la classe de bug qui rendait la policy violente.
|
||||
|
||||
mdp.py mélange DEUX conventions de signe : certaines fonctions de pénalité
|
||||
renvoient une magnitude positive (à multiplier par un poids négatif), d'autres
|
||||
renvoient déjà une valeur négative (à multiplier par un poids POSITIF).
|
||||
trunk_vertical_accel_penalty renvoie -|a_z| : avec le poids -0.02 hérité du
|
||||
standup, le double négatif RÉCOMPENSAIT l'accélération verticale — mesuré à
|
||||
Episode_Reward/gentle_rise = +0.0118, seul terme de pénalité loggé positif.
|
||||
"""
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
# Ces trois termes appellent des fonctions qui renvoient déjà du négatif
|
||||
# (height_l1_penalty, pose_l1_penalty, trunk_vertical_accel_penalty).
|
||||
for name in ("height_stand_l1", "pose_stand_l1", "gentle_rise"):
|
||||
assert cfg.rewards[name].weight > 0, (
|
||||
f"{name} appelle une fonction qui renvoie déjà du négatif : "
|
||||
f"un poids négatif en ferait une récompense"
|
||||
)
|
||||
# Et ces termes renvoient une magnitude positive → poids négatif.
|
||||
for name in ("joint_torques_l2", "joint_torque_rate_l2", "action_rate_l2"):
|
||||
assert cfg.rewards[name].weight < 0, f"{name} attend un poids négatif"
|
||||
|
||||
|
||||
def test_no_ungated_head_impact_penalty():
|
||||
"""PAS de pénalité d'impact tête non gatée — elle gelait la policy.
|
||||
|
||||
Essayée à -1.0 (valeurs de velstand) : la policy a convergé vers rester
|
||||
couchée, inerte. Mesuré sur le run d8rnko6p : head_impact_penalty -1.01/pas,
|
||||
le plus gros terme négatif, pendant que standing_composite s'effondrait de
|
||||
+14.3 à +3.3.
|
||||
|
||||
L'erreur de raisonnement était de croire qu'une pénalité « ciblée » ne bride
|
||||
pas le mouvement. Faux ici : pour se relever du dos, ce robot PIVOTE sur sa
|
||||
tête et ses épaules. La tête est le point d'appui du retournement, pas un
|
||||
dégât collatéral — la pénaliser, c'est pénaliser le seul mécanisme disponible.
|
||||
|
||||
Si le slam revient une fois le signe de gentle_rise corrigé, la reprise doit
|
||||
être une pénalité GATÉE EN HAUTEUR (comme upright_sharp l'est), qui épargne la
|
||||
phase de retournement au sol. Pas celle-ci.
|
||||
"""
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
assert "head_impact_penalty" not in cfg.rewards
|
||||
assert "head_impact_contact" not in [s.name for s in cfg.scene.sensors]
|
||||
|
||||
|
||||
def test_inherited_sensors_intact():
|
||||
# Les capteurs hérités de l'env roller sont utilisés par des récompenses
|
||||
# gardées (self_collisions) et par les observations.
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
names = [s.name for s in cfg.scene.sensors]
|
||||
assert "feet_ground_contact" in names
|
||||
assert "self_collision" in names
|
||||
|
||||
|
||||
def test_lazy_prone_optimum_is_documented_risk():
|
||||
"""Le gel vient d'un optimum paresseux : couché, jambes à HOME, ça paye.
|
||||
|
||||
pose_stand_legs restait à +7.72 sur 8 alors que le robot était allongé — les
|
||||
jambes sont à HOME en position couchée, donc la récompense de pose est encaissée
|
||||
quasi gratuitement. C'est le contrepoids qui rend « ne rien faire » viable dès
|
||||
qu'on ajoute un coût au mouvement. height_stand_l1 (poids +30) est le terme
|
||||
censé rendre « rester au sol » net négatif : il doit rester fort.
|
||||
"""
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
assert cfg.rewards["height_stand_l1"].weight >= 30.0
|
||||
assert cfg.rewards["com_upward_velocity"].weight > 0.0
|
||||
|
||||
|
||||
def test_damping_terms_are_not_numerically_negligible():
|
||||
"""Les amortisseurs dédiés ne pesaient littéralement rien.
|
||||
|
||||
Mesuré à convergence : joint_torque_rate_l2 -0.0002/pas et joint_torques_l2
|
||||
-0.0001/pas, face à ~+41.6 de récompense de tâche (rapport ~35:1 pour tous
|
||||
les amortisseurs réunis). joint_torque_rate_l2 est le levier SÛR à remonter :
|
||||
il pénalise la VARIATION de couple, pas le mouvement, donc il n'agit pas comme
|
||||
bloqueur de mouvement — le standup documente que body_ang_vel et action_rate,
|
||||
eux, gelaient le relevé depuis le dos.
|
||||
"""
|
||||
cfg = make_microduck_roller_standup_env_cfg()
|
||||
assert abs(cfg.rewards["joint_torque_rate_l2"].weight) >= 0.1
|
||||
# Les bloqueurs de mouvement restent à leurs valeurs « se relève de partout ».
|
||||
assert cfg.rewards["body_ang_vel"].weight == -0.05
|
||||
weights = [s["weight"] for s in cfg.curriculum["action_rate_weight"].params["weight_stages"]]
|
||||
assert min(weights) >= -1.0, "action_rate au-delà de -1.0 gelait le relevé (standup)"
|
||||
@@ -0,0 +1,39 @@
|
||||
import torch
|
||||
from mjlab_microduck.tasks.mdp import slope_move_masks
|
||||
|
||||
|
||||
def test_move_up_when_reached_bottom():
|
||||
# distance > size_x*0.4 (=3.2) → monte en difficulté
|
||||
dist = torch.tensor([5.0, 4.1])
|
||||
up, down = slope_move_masks(dist, size_x=8.0)
|
||||
assert bool(up[0]) and bool(up[1])
|
||||
assert not bool(down[0]) and not bool(down[1])
|
||||
|
||||
|
||||
def test_move_down_when_stuck_early():
|
||||
# distance < size_x*0.2 (=1.6) → descend en difficulté
|
||||
dist = torch.tensor([0.5, 1.0])
|
||||
up, down = slope_move_masks(dist, size_x=8.0)
|
||||
assert not bool(up[0]) and not bool(up[1])
|
||||
assert bool(down[0]) and bool(down[1])
|
||||
|
||||
|
||||
def test_stay_in_middle_band():
|
||||
# entre 1.6 et 3.2 → ni haut ni bas
|
||||
dist = torch.tensor([2.5])
|
||||
up, down = slope_move_masks(dist, size_x=8.0)
|
||||
assert not bool(up[0]) and not bool(down[0])
|
||||
|
||||
|
||||
def test_move_up_boundary_at_04():
|
||||
# promotion dès qu'on a descendu > 0.4*size_x (le robot a parcouru une bonne
|
||||
# partie de la rampe avant d'atteindre le plat de sortie).
|
||||
dist = torch.tensor([3.3])
|
||||
up, down = slope_move_masks(dist, size_x=8.0)
|
||||
assert bool(up[0])
|
||||
assert not bool(down[0])
|
||||
|
||||
# 3.0 reste dans la bande médiane (3.0 < 3.2 et 3.0 > 1.6)
|
||||
dist_mid = torch.tensor([3.0])
|
||||
up_mid, down_mid = slope_move_masks(dist_mid, size_x=8.0)
|
||||
assert not bool(up_mid[0]) and not bool(down_mid[0])
|
||||
@@ -0,0 +1,96 @@
|
||||
import math
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
|
||||
from mjlab_microduck.tasks.slope_terrain import (
|
||||
ramp_angle_by_difficulty,
|
||||
RAMP_DEG_MIN,
|
||||
RAMP_DEG_MAX,
|
||||
FlatRampTerrainCfg,
|
||||
)
|
||||
|
||||
|
||||
def test_ramp_angle_endpoints():
|
||||
assert math.isclose(ramp_angle_by_difficulty(0.0), math.radians(RAMP_DEG_MIN), abs_tol=1e-9)
|
||||
assert math.isclose(ramp_angle_by_difficulty(1.0), math.radians(RAMP_DEG_MAX), abs_tol=1e-9)
|
||||
|
||||
|
||||
def test_ramp_angle_midpoint():
|
||||
mid_deg = (RAMP_DEG_MIN + RAMP_DEG_MAX) / 2.0
|
||||
assert math.isclose(ramp_angle_by_difficulty(0.5), math.radians(mid_deg), abs_tol=1e-9)
|
||||
|
||||
|
||||
def test_ramp_angle_clamps_out_of_range():
|
||||
assert math.isclose(ramp_angle_by_difficulty(-1.0), math.radians(RAMP_DEG_MIN), abs_tol=1e-9)
|
||||
assert math.isclose(ramp_angle_by_difficulty(2.0), math.radians(RAMP_DEG_MAX), abs_tol=1e-9)
|
||||
|
||||
|
||||
def _empty_terrain_spec():
|
||||
spec = mujoco.MjSpec()
|
||||
spec.worldbody.add_body(name="terrain")
|
||||
return spec
|
||||
|
||||
|
||||
def test_flat_ramp_builds_geoms_and_origin_on_flat():
|
||||
cfg = FlatRampTerrainCfg(flat_length=2.0)
|
||||
cfg.size = (15.0, 4.0) # posé normalement par le générateur
|
||||
spec = _empty_terrain_spec()
|
||||
out = cfg.function(difficulty=0.5, spec=spec, rng=np.random.default_rng(0))
|
||||
# trois géométries : plat de départ + rampe + plat de sortie
|
||||
assert len(out.geometries) == 3
|
||||
# origine SUR la rampe (au-delà du plat), donc x > flat_length et z < 0
|
||||
assert out.origin[0] == cfg.flat_length + cfg.spawn_on_ramp
|
||||
assert out.origin[2] < 0.0
|
||||
# z = surface inclinée à spawn_on_ramp du bord (drop = d * tan(angle))
|
||||
angle = ramp_angle_by_difficulty(0.5, cfg.deg_min, cfg.deg_max)
|
||||
assert abs(out.origin[2] - (-cfg.spawn_on_ramp * math.tan(angle))) < 1e-9
|
||||
|
||||
|
||||
def test_flat_ramp_steeper_at_higher_difficulty():
|
||||
# à difficulté plus haute, le bout de rampe descend plus bas
|
||||
cfg = FlatRampTerrainCfg()
|
||||
cfg.size = (15.0, 4.0)
|
||||
easy = cfg.function(0.0, _empty_terrain_spec(), np.random.default_rng(0))
|
||||
hard = cfg.function(1.0, _empty_terrain_spec(), np.random.default_rng(0))
|
||||
# la rampe (2e géométrie) est plus basse (centre z plus négatif) en difficile
|
||||
# (même rng -> même longueur tirée -> seule la pente change)
|
||||
assert hard.geometries[1].geom.pos[2] < easy.geometries[1].geom.pos[2]
|
||||
|
||||
|
||||
def test_ramp_joins_flat_platform_no_gap():
|
||||
# le haut de la rampe doit toucher le bord de la plateforme plate (x=flat_length) :
|
||||
# centre rampe décalé de -(t/2)*sin(angle) en x.
|
||||
cfg = FlatRampTerrainCfg()
|
||||
cfg.size = (15.0, 4.0)
|
||||
out = cfg.function(0.5, _empty_terrain_spec(), np.random.default_rng(0))
|
||||
ramp = out.geometries[1].geom
|
||||
angle = ramp_angle_by_difficulty(0.5, cfg.deg_min, cfg.deg_max)
|
||||
surf_half = ramp.size[0]
|
||||
ramp_len = surf_half * 2.0 * math.cos(angle)
|
||||
expected_cx = cfg.flat_length + ramp_len / 2.0 - (cfg.thickness / 2.0) * math.sin(angle)
|
||||
assert abs(ramp.pos[0] - expected_cx) < 1e-6
|
||||
|
||||
|
||||
def test_flat_ramp_runout_at_ramp_bottom():
|
||||
# le plat de sortie (3e géométrie) est au niveau du bas de la rampe (z<0),
|
||||
# et sa surface est plate (box non tourné : quat identité).
|
||||
cfg = FlatRampTerrainCfg()
|
||||
cfg.size = (15.0, 4.0)
|
||||
out = cfg.function(1.0, _empty_terrain_spec(), np.random.default_rng(0))
|
||||
runout = out.geometries[2].geom
|
||||
assert runout.pos[2] < 0.0 # descendu sous le plat de départ
|
||||
# quaternion identité (plat, pas incliné)
|
||||
assert math.isclose(runout.quat[0], 1.0, abs_tol=1e-9)
|
||||
|
||||
|
||||
def test_ramp_length_within_range():
|
||||
cfg = FlatRampTerrainCfg(ramp_length_range=(3.0, 8.0))
|
||||
cfg.size = (15.0, 4.0)
|
||||
# surface de rampe = ramp_length / cos(angle) ; à difficulté 0, angle=2°,
|
||||
# donc surf_len ~= ramp_length. On vérifie sur plusieurs tirages.
|
||||
for seed in range(20):
|
||||
out = cfg.function(0.0, _empty_terrain_spec(), np.random.default_rng(seed))
|
||||
surf_half = out.geometries[1].geom.size[0]
|
||||
ramp_len = surf_half * 2.0 * math.cos(math.radians(2.0))
|
||||
assert 3.0 - 1e-6 <= ramp_len <= 8.0 + 1e-6
|
||||
@@ -0,0 +1,387 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
|
||||
from mjlab_microduck.tasks import mdp
|
||||
|
||||
# Enveloppe du spec : accel 0.5s / régime 1.6s / freinage 0.5s / repos 1.4s sur 4s.
|
||||
_ENV = dict(rate_max=6.0, accel_end=0.125, hold_end=0.525, brake_end=0.650)
|
||||
|
||||
|
||||
def test_spin_rate_segment_boundaries():
|
||||
# bornes des 4 segments : 0 au départ, plein régime sur [accel_end, hold_end],
|
||||
# encore plein régime au tout début du freinage, 0 dès le segment de repos.
|
||||
phase = torch.tensor([0.0, 0.125, 0.30, 0.525, 0.650, 0.80, 0.999])
|
||||
w = mdp.spin_rate_by_phase(phase, **_ENV)
|
||||
expected = torch.tensor([0.0, 6.0, 6.0, 6.0, 0.0, 0.0, 0.0])
|
||||
assert torch.allclose(w, expected, atol=1e-6)
|
||||
|
||||
|
||||
def test_spin_rate_accel_ramp_is_increasing():
|
||||
phase = torch.linspace(0.0, 0.125, 20)
|
||||
w = mdp.spin_rate_by_phase(phase, **_ENV)
|
||||
assert torch.all(w[1:] >= w[:-1])
|
||||
# milieu de la rampe de lancement -> moitié de la cible
|
||||
mid = mdp.spin_rate_by_phase(torch.tensor([0.0625]), **_ENV)
|
||||
assert torch.allclose(mid, torch.tensor([3.0]), atol=1e-6)
|
||||
|
||||
|
||||
def test_spin_rate_brake_ramp_is_decreasing():
|
||||
phase = torch.linspace(0.525, 0.6499, 20)
|
||||
w = mdp.spin_rate_by_phase(phase, **_ENV)
|
||||
assert torch.all(w[1:] <= w[:-1])
|
||||
# milieu du freinage -> moitié de la cible
|
||||
mid = mdp.spin_rate_by_phase(torch.tensor([0.5875]), **_ENV)
|
||||
assert torch.allclose(mid, torch.tensor([3.0]), atol=1e-6)
|
||||
|
||||
|
||||
def test_spin_rate_integral_matches_trapezoid_shape_at_rate_max_6():
|
||||
# Ce test protège la FORME du trapèze (2.1 * rate_max rad par cycle), pas la
|
||||
# cible réellement expédiée : à rate_max=6.0 (hypothétique, cf. _ENV ci-dessus)
|
||||
# ça vaut ~4*pi rad = 2 tours. Enveloppe exacte = 12.6 rad, 4*pi = 12.566 ->
|
||||
# tolérance 1 %. La cible EN VIGUEUR est couverte par le test suivant.
|
||||
n = 100_000
|
||||
phase = (torch.arange(n, dtype=torch.float64) + 0.5) / n
|
||||
w = mdp.spin_rate_by_phase(phase, **_ENV)
|
||||
integral = float(w.mean()) * 4.0
|
||||
assert abs(integral - 4 * math.pi) / (4 * math.pi) < 0.01
|
||||
|
||||
|
||||
def test_spin_rate_max_integrates_to_2_1_times_itself_per_cycle():
|
||||
# LE test qui protège la cible EXPÉDIÉE (mdp.SPIN_RATE_MAX), par opposition au
|
||||
# test ci-dessus qui ne teste que la forme à rate_max=6.0. L'aire sous
|
||||
# l'enveloppe sur un cycle vaut 2.1 * rate_max rad, quel que soit rate_max
|
||||
# (0.25 + 1.6 + 0.25 = 2.1, cf. le commentaire au-dessus des constantes dans
|
||||
# mdp.py). Avec le réglage actuel (SPIN_RATE_MAX = 3.0) ça donne 6.3 rad,
|
||||
# soit ~1 tour -- pas 2. Ce test échoue bruyamment si quelqu'un change la
|
||||
# cible sans réfléchir au nombre de tours que ça implique.
|
||||
n = 100_000
|
||||
phase = (torch.arange(n, dtype=torch.float64) + 0.5) / n
|
||||
w = mdp.spin_rate_by_phase(
|
||||
phase,
|
||||
rate_max=mdp.SPIN_RATE_MAX,
|
||||
accel_end=mdp.SPIN_ACCEL_END,
|
||||
hold_end=mdp.SPIN_HOLD_END,
|
||||
brake_end=mdp.SPIN_BRAKE_END,
|
||||
)
|
||||
integral = float(w.mean()) * mdp.SPIN_PERIOD
|
||||
expected = 2.1 * mdp.SPIN_RATE_MAX
|
||||
assert abs(integral - expected) / expected < 0.01
|
||||
|
||||
|
||||
def test_spin_gate_is_normalized_rate():
|
||||
phase = torch.tensor([0.0, 0.0625, 0.30, 0.5875, 0.80])
|
||||
gate = mdp.spin_gate_by_phase(phase, **_ENV)
|
||||
rate = mdp.spin_rate_by_phase(phase, **_ENV)
|
||||
assert torch.allclose(gate, rate / 6.0, atol=1e-6)
|
||||
assert torch.all(gate >= 0.0) and torch.all(gate <= 1.0)
|
||||
|
||||
|
||||
def test_spin_gate_is_zero_over_the_whole_rest_segment():
|
||||
# pendant le repos aucune amorce ne doit pousser au ciseau -> porte nulle,
|
||||
# c'est ce qui donne une sortie de trick propre vers la policy roller.
|
||||
phase = torch.linspace(0.650, 0.999, 50)
|
||||
gate = mdp.spin_gate_by_phase(phase, **_ENV)
|
||||
assert torch.allclose(gate, torch.zeros_like(gate), atol=1e-6)
|
||||
|
||||
|
||||
# ── faux env minimal : permet de tester les wrappers de reward sans MuJoCo ────
|
||||
class _FakeData:
|
||||
def __init__(self, ang_vel_b=None, lin_vel_b=None, joint_pos=None, joint_vel=None):
|
||||
self.root_link_ang_vel_b = ang_vel_b
|
||||
self.root_link_lin_vel_b = lin_vel_b
|
||||
self.joint_pos = joint_pos
|
||||
self.joint_vel = joint_vel
|
||||
|
||||
|
||||
class _FakeEntity:
|
||||
"""Entity minimale : find_joints() résout par nom depuis un dict {nom: index}."""
|
||||
|
||||
def __init__(self, data, joint_ids=None):
|
||||
self.data = data
|
||||
self._joint_ids = joint_ids or {}
|
||||
|
||||
def find_joints(self, pattern):
|
||||
import re
|
||||
|
||||
names = list(self._joint_ids.keys())
|
||||
if isinstance(pattern, (list, tuple)):
|
||||
matched = [n for n in names if n in pattern]
|
||||
else:
|
||||
matched = [n for n in names if re.fullmatch(pattern, n)]
|
||||
assert matched, f"aucun joint ne matche {pattern!r} parmi {names}"
|
||||
return [self._joint_ids[n] for n in matched], matched
|
||||
|
||||
|
||||
class _FakeCommandManager:
|
||||
def __init__(self, cmd):
|
||||
self._cmd = cmd
|
||||
|
||||
def get_command(self, name):
|
||||
return self._cmd
|
||||
|
||||
|
||||
class _FakeSensorData:
|
||||
def __init__(self, current_contact_time):
|
||||
self.current_contact_time = current_contact_time
|
||||
|
||||
|
||||
class _FakeSensor:
|
||||
def __init__(self, current_contact_time):
|
||||
self.data = _FakeSensorData(current_contact_time)
|
||||
|
||||
|
||||
class _FakeEnv:
|
||||
def __init__(self, entity, cmd=None, sensors=None):
|
||||
self.scene = {"robot": entity, **(sensors or {})}
|
||||
self.command_manager = _FakeCommandManager(cmd)
|
||||
self.device = "cpu"
|
||||
|
||||
|
||||
def _phase_cmd(phases):
|
||||
"""Commande du slot telle que la voit la policy : [cos(2*pi*phi), sin(...), 0]."""
|
||||
p = torch.as_tensor(phases, dtype=torch.float32)
|
||||
return torch.stack(
|
||||
[torch.cos(2 * math.pi * p), torch.sin(2 * math.pi * p), torch.zeros_like(p)],
|
||||
dim=-1,
|
||||
)
|
||||
|
||||
|
||||
# ── phase recover ────────────────────────────────────────────────────────────
|
||||
def test_spin_phase_from_command_roundtrip():
|
||||
phases = torch.tensor([0.0, 0.125, 0.4, 0.65, 0.9])
|
||||
got = mdp.spin_phase_from_command(_phase_cmd(phases))
|
||||
assert torch.allclose(got, phases, atol=1e-5)
|
||||
|
||||
|
||||
# ── spin_rate_track ──────────────────────────────────────────────────────────
|
||||
def test_spin_rate_reward_peaks_on_exact_match():
|
||||
w = torch.tensor([6.0, 6.0])
|
||||
target = torch.tensor([6.0, 4.5])
|
||||
r = mdp.spin_rate_reward_from_values(w, target, std=1.5)
|
||||
# erreur nulle -> 1.0 ; erreur = 1 std -> exp(-1)
|
||||
assert torch.allclose(r, torch.tensor([1.0, math.exp(-1.0)]), atol=1e-6)
|
||||
|
||||
|
||||
def test_spin_rate_track_uses_yaw_and_phase():
|
||||
# phase 0.30 = plein régime -> cible SPIN_RATE_MAX (3.0 rad/s, défaut appelé
|
||||
# ici implicitement). Un robot qui tourne à la cible doit toucher 1.0 ; un
|
||||
# robot immobile doit être largement en dessous (exp(-(3/1.5)^2) = 0.018 au
|
||||
# réglage courant : std=1.5 reste bien calibré à cette cible, cf. mdp.py).
|
||||
ang = torch.tensor([[0.0, 0.0, mdp.SPIN_RATE_MAX], [0.0, 0.0, 0.0]])
|
||||
env = _FakeEnv(
|
||||
_FakeEntity(_FakeData(ang_vel_b=ang)), cmd=_phase_cmd([0.30, 0.30])
|
||||
)
|
||||
r = mdp.spin_rate_track(env, std=1.5)
|
||||
assert r[0] > 0.99
|
||||
assert r[1] < 0.05
|
||||
|
||||
|
||||
def test_spin_rate_track_wants_stillness_during_rest():
|
||||
# phase 0.80 = repos -> cible 0 : tourner encore est puni, être immobile payé.
|
||||
ang = torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 6.0]])
|
||||
env = _FakeEnv(
|
||||
_FakeEntity(_FakeData(ang_vel_b=ang)), cmd=_phase_cmd([0.80, 0.80])
|
||||
)
|
||||
r = mdp.spin_rate_track(env, std=1.5)
|
||||
assert r[0] > 0.99
|
||||
assert r[1] < 0.01
|
||||
|
||||
|
||||
def test_spin_rate_track_penalizes_wrong_direction():
|
||||
# tourner à -SPIN_RATE_MAX (horaire) quand on demande +SPIN_RATE_MAX doit
|
||||
# être pire qu'immobile.
|
||||
ang = torch.tensor([[0.0, 0.0, -mdp.SPIN_RATE_MAX], [0.0, 0.0, 0.0]])
|
||||
env = _FakeEnv(
|
||||
_FakeEntity(_FakeData(ang_vel_b=ang)), cmd=_phase_cmd([0.30, 0.30])
|
||||
)
|
||||
r = mdp.spin_rate_track(env, std=1.5)
|
||||
assert r[0] < r[1]
|
||||
|
||||
|
||||
# ── spin_rate_l1 ─────────────────────────────────────────────────────────────
|
||||
def test_spin_rate_l1_is_negative_absolute_error():
|
||||
# phase 0.30 = plein régime -> cible SPIN_RATE_MAX (3.0 rad/s, défaut).
|
||||
ang = torch.tensor([[0.0, 0.0, mdp.SPIN_RATE_MAX], [0.0, 0.0, 1.0]])
|
||||
env = _FakeEnv(
|
||||
_FakeEntity(_FakeData(ang_vel_b=ang)), cmd=_phase_cmd([0.30, 0.30])
|
||||
)
|
||||
r = mdp.spin_rate_l1(env)
|
||||
expected = torch.tensor([0.0, -(mdp.SPIN_RATE_MAX - 1.0)])
|
||||
assert torch.allclose(r, expected, atol=1e-5)
|
||||
|
||||
|
||||
# ── spin_stay_in_place ───────────────────────────────────────────────────────
|
||||
def test_spin_stay_in_place_is_squared_planar_speed():
|
||||
# phase 0.30 = plein régime -> coût plein tarif
|
||||
lin = torch.tensor([[0.0, 0.0, 0.0], [0.3, 0.4, 9.0]])
|
||||
env = _FakeEnv(
|
||||
_FakeEntity(_FakeData(lin_vel_b=lin)), cmd=_phase_cmd([0.30, 0.30])
|
||||
)
|
||||
c = mdp.spin_stay_in_place(env)
|
||||
# 0.3^2 + 0.4^2 = 0.25 ; la composante z est ignorée
|
||||
assert torch.allclose(c, torch.tensor([0.0, 0.25]), atol=1e-6)
|
||||
|
||||
|
||||
def test_spin_stay_in_place_is_attenuated_during_the_launch_ramp():
|
||||
# Même vitesse, deux phases : dans la rampe de lancement (0.05 < accel_end) le
|
||||
# coût est multiplié par launch_scale, en régime (0.30) il est plein tarif.
|
||||
# C'est ce qui empêche ce terme de s'opposer à l'injection de moment angulaire.
|
||||
lin = torch.tensor([[0.3, 0.4, 0.0], [0.3, 0.4, 0.0]])
|
||||
env = _FakeEnv(
|
||||
_FakeEntity(_FakeData(lin_vel_b=lin)), cmd=_phase_cmd([0.05, 0.30])
|
||||
)
|
||||
c = mdp.spin_stay_in_place(env, launch_scale=0.2, accel_end=0.125)
|
||||
# 0.25 * 0.2 = 0.05
|
||||
assert torch.allclose(c, torch.tensor([0.05, 0.25]), atol=1e-6)
|
||||
assert c[0] < c[1]
|
||||
|
||||
|
||||
def test_spin_stay_in_place_is_full_price_during_rest():
|
||||
# Pendant le repos on veut le robot IMMOBILE : ce terme ne doit PAS être éteint,
|
||||
# contrairement aux amorces (spin_wheel_differential, spin_grounded, ciseau).
|
||||
lin = torch.tensor([[0.3, 0.4, 0.0]])
|
||||
env = _FakeEnv(_FakeEntity(_FakeData(lin_vel_b=lin)), cmd=_phase_cmd([0.80]))
|
||||
c = mdp.spin_stay_in_place(env)
|
||||
assert torch.allclose(c, torch.tensor([0.25]), atol=1e-6)
|
||||
|
||||
|
||||
# ── spin_wheel_differential ──────────────────────────────────────────────────
|
||||
_WHEEL_IDS = {
|
||||
"passive_LF_wheel": 0,
|
||||
"passive_LR_wheel": 1,
|
||||
"passive_RF_wheel": 2,
|
||||
"passive_RR_wheel": 3,
|
||||
}
|
||||
|
||||
|
||||
def _wheel_env(vel_rows, phases):
|
||||
vel = torch.tensor(vel_rows, dtype=torch.float32)
|
||||
entity = _FakeEntity(_FakeData(joint_vel=vel), joint_ids=_WHEEL_IDS)
|
||||
return _FakeEnv(entity, cmd=_phase_cmd(phases))
|
||||
|
||||
|
||||
def test_wheel_differential_rewards_counter_rolling_wheels():
|
||||
# anti-horaire : roues GAUCHE négatives (patin part en arrière), DROITE
|
||||
# positives -> omega_D - omega_G > 0 -> récompensé.
|
||||
env = _wheel_env(
|
||||
[
|
||||
[-10.0, -10.0, 10.0, 10.0], # bon différentiel
|
||||
[10.0, 10.0, 10.0, 10.0], # tout droit : différentiel nul
|
||||
[10.0, 10.0, -10.0, -10.0], # différentiel inversé (horaire)
|
||||
],
|
||||
[0.30, 0.30, 0.30],
|
||||
)
|
||||
r = mdp.spin_wheel_differential(env, omega_scale=20.0)
|
||||
assert r[0] > 0.5
|
||||
assert torch.allclose(r[1], torch.tensor(0.0), atol=1e-6)
|
||||
assert torch.allclose(r[2], torch.tensor(0.0), atol=1e-6)
|
||||
|
||||
|
||||
def test_wheel_differential_is_gated_off_during_rest():
|
||||
# même bon différentiel, mais en phase de repos -> porte nulle -> pas payé.
|
||||
env = _wheel_env([[-10.0, -10.0, 10.0, 10.0]], [0.80])
|
||||
r = mdp.spin_wheel_differential(env, omega_scale=20.0)
|
||||
assert torch.allclose(r, torch.zeros(1), atol=1e-6)
|
||||
|
||||
|
||||
def test_wheel_differential_saturates():
|
||||
# tanh : au-delà de omega_scale la reward sature, pas de course à la vitesse.
|
||||
env = _wheel_env(
|
||||
[[-10.0, -10.0, 10.0, 10.0], [-100.0, -100.0, 100.0, 100.0]], [0.30, 0.30]
|
||||
)
|
||||
r = mdp.spin_wheel_differential(env, omega_scale=20.0)
|
||||
assert r[1] > r[0]
|
||||
assert r[1] <= 1.0
|
||||
|
||||
|
||||
def test_wheel_differential_from_values_is_pure():
|
||||
diff = torch.tensor([20.0, 0.0, -20.0])
|
||||
gate = torch.ones(3)
|
||||
r = mdp.spin_wheel_differential_from_values(diff, gate, omega_scale=20.0)
|
||||
expected = torch.tensor([math.tanh(1.0), 0.0, 0.0])
|
||||
assert torch.allclose(r, expected, atol=1e-6)
|
||||
|
||||
|
||||
# ── spin_grounded ────────────────────────────────────────────────────────────
|
||||
def test_spin_grounded_rewards_both_blades_down_and_is_gated():
|
||||
contact = torch.tensor([[0.2, 0.3], [0.2, 0.0], [0.0, 0.0], [0.2, 0.3]])
|
||||
entity = _FakeEntity(_FakeData())
|
||||
env = _FakeEnv(
|
||||
entity,
|
||||
cmd=_phase_cmd([0.30, 0.30, 0.30, 0.80]),
|
||||
sensors={"feet_ground_contact": _FakeSensor(contact)},
|
||||
)
|
||||
r = mdp.spin_grounded(env, sensor_name="feet_ground_contact")
|
||||
# deux lames au sol en régime -> porte 1.0 ; une seule ou zéro -> 0 ;
|
||||
# deux lames au sol mais en repos -> porte 0.
|
||||
assert torch.allclose(r, torch.tensor([1.0, 0.0, 0.0, 0.0]), atol=1e-6)
|
||||
|
||||
|
||||
# ── leg_antisymmetry ─────────────────────────────────────────────────────────
|
||||
_LEG_IDS = {
|
||||
"left_hip_pitch": 0,
|
||||
"left_knee": 1,
|
||||
"right_hip_pitch": 2,
|
||||
"right_knee": 3,
|
||||
}
|
||||
|
||||
|
||||
def _leg_env(pos_rows, phases):
|
||||
pos = torch.tensor(pos_rows, dtype=torch.float32)
|
||||
entity = _FakeEntity(_FakeData(joint_pos=pos), joint_ids=_LEG_IDS)
|
||||
return _FakeEnv(entity, cmd=_phase_cmd(phases))
|
||||
|
||||
|
||||
def test_leg_antisymmetry_prefers_scissor_over_mirror():
|
||||
# convention miroir : q_G = -q_D est une pose SYMÉTRIQUE (mauvais ici),
|
||||
# q_G = q_D est le CISEAU (bon ici). Valeur = -mean|q_G - q_D|, donc <= 0.
|
||||
env = _leg_env(
|
||||
[
|
||||
[0.4, 0.3, 0.4, 0.3], # ciseau parfait : q_G == q_D -> 0.0
|
||||
[0.4, 0.3, -0.4, -0.3], # miroir : écart 0.8 et 0.6 -> -0.7
|
||||
],
|
||||
[0.30, 0.30],
|
||||
)
|
||||
r = mdp.leg_antisymmetry(env)
|
||||
assert torch.allclose(r, torch.tensor([0.0, -0.7]), atol=1e-6)
|
||||
assert r[0] > r[1]
|
||||
|
||||
|
||||
def test_leg_antisymmetry_is_gated_off_during_rest():
|
||||
# en repos la porte est nulle : rien ne pousse au ciseau, station neutre libre.
|
||||
env = _leg_env([[0.4, 0.3, -0.4, -0.3]], [0.80])
|
||||
r = mdp.leg_antisymmetry(env)
|
||||
assert torch.allclose(r, torch.zeros(1), atol=1e-6)
|
||||
|
||||
|
||||
# ── neck_joint_pos_l2 : paramètre pattern ────────────────────────────────────
|
||||
_NECK_IDS = {
|
||||
"neck_pitch": 0,
|
||||
"head_pitch": 1,
|
||||
"head_roll": 2,
|
||||
"head_yaw": 3,
|
||||
}
|
||||
|
||||
|
||||
def test_neck_joint_pos_l2_pattern_can_exclude_head_yaw():
|
||||
class _NeckData(_FakeData):
|
||||
def __init__(self, joint_pos, default_joint_pos):
|
||||
super().__init__(joint_pos=joint_pos)
|
||||
self.default_joint_pos = default_joint_pos
|
||||
|
||||
pos = torch.tensor([[0.0, 0.0, 0.0, 1.0]]) # seul head_yaw dévie, de 1 rad
|
||||
default = torch.zeros(1, 4)
|
||||
entity = _FakeEntity(_NeckData(pos, default), joint_ids=_NECK_IDS)
|
||||
env = _FakeEnv(entity)
|
||||
|
||||
# motif par défaut : head_yaw compté -> coût 1.0
|
||||
assert torch.allclose(
|
||||
mdp.neck_joint_pos_l2(env), torch.tensor([1.0]), atol=1e-6
|
||||
)
|
||||
# motif du spin : head_yaw exclu -> coût 0.0 (tête libre en lacet)
|
||||
assert torch.allclose(
|
||||
mdp.neck_joint_pos_l2(env, pattern=r"^(neck_pitch|head_pitch|head_roll)$"),
|
||||
torch.tensor([0.0]),
|
||||
atol=1e-6,
|
||||
)
|
||||
@@ -0,0 +1,110 @@
|
||||
from mjlab_microduck.tasks import mdp as microduck_mdp
|
||||
from mjlab_microduck.tasks.microduck_spin_env_cfg import (
|
||||
make_microduck_spin_env_cfg,
|
||||
MicroduckSpinRlCfg,
|
||||
)
|
||||
|
||||
|
||||
def test_cfg_uses_phase_command_with_runtime_default_period():
|
||||
cfg = make_microduck_spin_env_cfg()
|
||||
cmd = cfg.commands["twist"]
|
||||
assert isinstance(cmd, microduck_mdp.GroundPickPhaseCommandCfg)
|
||||
# 4.0 s = le défaut de --ground-pick-period : rien à passer au runtime
|
||||
assert cmd.period == 4.0
|
||||
# chaque épisode démarre à phase 0 (debout), comme le bouton au déploiement
|
||||
assert cmd.randomize_phase is False
|
||||
|
||||
|
||||
def test_cfg_has_the_spin_rewards():
|
||||
cfg = make_microduck_spin_env_cfg()
|
||||
for name in (
|
||||
"spin_rate_track",
|
||||
"spin_rate_l1",
|
||||
"spin_stay_in_place",
|
||||
"spin_wheel_differential",
|
||||
"spin_grounded",
|
||||
"leg_antisymmetry",
|
||||
):
|
||||
assert name in cfg.rewards, name
|
||||
# objectif principal avec un poids dominant
|
||||
assert cfg.rewards["spin_rate_track"].weight == 6.0
|
||||
# sur-place est un COÛT
|
||||
assert cfg.rewards["spin_stay_in_place"].weight < 0.0
|
||||
|
||||
|
||||
def test_stay_in_place_is_attenuated_during_the_launch_ramp():
|
||||
# Renforcé à -3.0, ce terme s'opposerait à l'injection de moment angulaire s'il
|
||||
# était plein tarif pendant la rampe de lancement : il doit y être atténué.
|
||||
cfg = make_microduck_spin_env_cfg()
|
||||
params = cfg.rewards["spin_stay_in_place"].params
|
||||
assert 0.0 < params["launch_scale"] < 1.0
|
||||
assert params["accel_end"] == microduck_mdp.SPIN_ACCEL_END
|
||||
# cible positive = anti-horaire (le sens est porté par l'enveloppe)
|
||||
assert microduck_mdp.SPIN_RATE_MAX > 0.0
|
||||
|
||||
|
||||
def test_angular_momentum_reward_is_removed():
|
||||
# Régression : angular_momentum_penalty pénalise la NORME 3D du moment
|
||||
# angulaire, elle combattrait directement le spin. Elle doit être absente.
|
||||
cfg = make_microduck_spin_env_cfg()
|
||||
assert "angular_momentum" not in cfg.rewards
|
||||
# body_ang_vel ne pénalise que x/y -> elle reste, elle mate le ballant
|
||||
assert "body_ang_vel" in cfg.rewards
|
||||
|
||||
|
||||
def test_head_yaw_is_free_to_act_as_a_flywheel():
|
||||
cfg = make_microduck_spin_env_cfg()
|
||||
pattern = cfg.rewards["neck_joint_pos_l2"].params["pattern"]
|
||||
assert "head_yaw" not in pattern
|
||||
|
||||
|
||||
def test_entry_velocity_allows_standstill_and_slow_roll():
|
||||
cfg = make_microduck_spin_env_cfg()
|
||||
# jamais via un push en mode reset (régression NaN du crouch)
|
||||
assert "entry_velocity" not in cfg.events
|
||||
lo, hi = cfg.events["reset_base"].params["velocity_range"]["x"]
|
||||
assert lo == 0.0 and hi > 0.0
|
||||
|
||||
|
||||
def test_symmetry_augmentation_is_disabled():
|
||||
# la symétrie G/D transformerait un spin à gauche en spin à droite
|
||||
assert MicroduckSpinRlCfg.algorithm.symmetry_cfg is None
|
||||
|
||||
|
||||
def test_leg_antisymmetry_shaping_decays():
|
||||
cfg = make_microduck_spin_env_cfg()
|
||||
stages = cfg.curriculum["leg_antisym_weight"].params["weight_stages"]
|
||||
weights = [s["weight"] for s in stages]
|
||||
assert weights[0] == cfg.rewards["leg_antisymmetry"].weight
|
||||
assert weights == sorted(weights, reverse=True)
|
||||
assert weights[-1] < weights[0]
|
||||
|
||||
|
||||
def test_actor_observation_keeps_the_61d_slot_layout():
|
||||
# condition pour que l'ONNX charge dans le slot du runtime. L'égalité exacte
|
||||
# des dimensions avec le crouch est vérifiée par test_obs_parity_with_roller_crouch
|
||||
# ci-dessous ; ici on vérifie la structure.
|
||||
cfg = make_microduck_spin_env_cfg()
|
||||
terms = cfg.observations["actor"].terms
|
||||
assert "base_lin_vel" not in terms
|
||||
assert "height_scan" not in terms
|
||||
for padded in ("head_command", "body_command"):
|
||||
assert padded in terms
|
||||
assert terms["head_command"].params["dim"] == 4
|
||||
assert terms["body_command"].params["dim"] == 6
|
||||
|
||||
|
||||
def test_obs_parity_with_roller_crouch():
|
||||
# Parité de layout obligatoire : sinon l'ONNX exporté ne charge pas dans le
|
||||
# slot du runtime. Contrairement au test de structure ci-dessus, celui-ci
|
||||
# compare l'ordre EXACT des termes, groupe par groupe.
|
||||
from mjlab_microduck.tasks.microduck_roller_crouch_env_cfg import (
|
||||
make_microduck_roller_crouch_env_cfg,
|
||||
)
|
||||
|
||||
spin = make_microduck_spin_env_cfg()
|
||||
crouch = make_microduck_roller_crouch_env_cfg()
|
||||
for grp in ("actor", "critic"):
|
||||
assert list(spin.observations[grp].terms.keys()) == list(
|
||||
crouch.observations[grp].terms.keys()
|
||||
), f"layout d'observation divergent sur le groupe {grp}"
|
||||
@@ -0,0 +1,42 @@
|
||||
from mjlab.tasks.velocity import mdp
|
||||
from mjlab_microduck.tasks.microduck_velocity_swizzle_env_cfg import (
|
||||
make_microduck_velocity_swizzle_env_cfg,
|
||||
)
|
||||
from mjlab_microduck.tasks.microduck_velocity_rollers_env_cfg import (
|
||||
make_microduck_velocity_rollers_env_cfg,
|
||||
)
|
||||
|
||||
|
||||
def test_swizzle_head_control_wired():
|
||||
cfg = make_microduck_velocity_swizzle_env_cfg()
|
||||
roller_cfg = make_microduck_velocity_rollers_env_cfg()
|
||||
|
||||
# Head-pose command term exists.
|
||||
assert "head_pose" in cfg.commands
|
||||
|
||||
# head_command obs is the REAL command (not zero-padded) on both groups.
|
||||
for group in ("actor", "critic"):
|
||||
term = cfg.observations[group].terms["head_command"]
|
||||
assert term.func is mdp.generated_commands
|
||||
assert term.params["command_name"] == "head_pose"
|
||||
|
||||
# head_pose_tracking reward exists.
|
||||
assert "head_pose_tracking" in cfg.rewards
|
||||
|
||||
# The two HOME-pullers that would fight the head command are handled:
|
||||
# - neck_joint_pos_l2 removed
|
||||
assert "neck_joint_pos_l2" not in cfg.rewards
|
||||
# - pose reward scoped to leg joints via a negative-lookahead regex that
|
||||
# excludes neck/head (and passive wheels)
|
||||
pose_joints = cfg.rewards["pose"].params["asset_cfg"].joint_names
|
||||
assert any(
|
||||
"(?!" in j and "neck" in j and "head" in j for j in pose_joints
|
||||
), f"pose reward not scoped away from neck/head: {pose_joints}"
|
||||
|
||||
# Pose reward function unchanged (not swapped to a different function)
|
||||
assert cfg.rewards["pose"].func is roller_cfg.rewards["pose"].func, \
|
||||
"pose reward function was swapped (should only scope asset_cfg)"
|
||||
|
||||
# Late head curricula exist.
|
||||
assert "head_pose_tracking_weight" in cfg.curriculum
|
||||
assert "head_pose_range" in cfg.curriculum
|
||||
@@ -0,0 +1,65 @@
|
||||
"""wheel_glide_reward : récompense le ROULEMENT des roues vers l'avant (glisse
|
||||
par gravité), plafonné à cap_speed, nul si les roues reculent, NaN-safe.
|
||||
Indépendant de toute commande (la tâche pente a une commande nulle).
|
||||
"""
|
||||
|
||||
import re
|
||||
|
||||
import torch
|
||||
|
||||
from mjlab_microduck.tasks.mdp import wheel_glide_reward
|
||||
|
||||
# Current model joint names (post 2026-07 re-export: underscore spelling).
|
||||
_WHEELS = {"passive_LF_wheel": 0, "passive_LR_wheel": 1, "passive_RF_wheel": 2, "passive_RR_wheel": 3}
|
||||
|
||||
|
||||
class _Data:
|
||||
def __init__(self, omegas):
|
||||
# 4 roues, colonnes 0..3 dans l'ordre LF,LR,RF,RR
|
||||
self.joint_vel = torch.tensor([omegas], dtype=torch.float32)
|
||||
|
||||
|
||||
class _Asset:
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
|
||||
def find_joints(self, pattern):
|
||||
# Regex resolution like the real Entity.find_joints (mdp queries use
|
||||
# spelling-tolerant patterns such as "passive_LF_?wheel").
|
||||
ids = [i for name, i in _WHEELS.items() if re.fullmatch(pattern, name)]
|
||||
assert ids, pattern
|
||||
return ids, None
|
||||
|
||||
|
||||
class _Env:
|
||||
def __init__(self, omegas):
|
||||
self._a = _Asset(_Data(omegas))
|
||||
|
||||
def __getitem__(self, _k):
|
||||
return self._a
|
||||
|
||||
@property
|
||||
def scene(self):
|
||||
return self
|
||||
|
||||
|
||||
def test_rewards_forward_roll_below_cap():
|
||||
# omega=10 rad/s sur les 4 -> vitesse = 10*0.0175 = 0.175 m/s (< cap 0.35)
|
||||
out = wheel_glide_reward(_Env([10.0, 10.0, 10.0, 10.0]), cap_speed=0.35)
|
||||
assert abs(float(out[0]) - 0.175) < 1e-6
|
||||
|
||||
|
||||
def test_caps_fast_roll():
|
||||
# omega=40 -> 0.7 m/s -> plafonné à 0.35
|
||||
out = wheel_glide_reward(_Env([40.0, 40.0, 40.0, 40.0]), cap_speed=0.35)
|
||||
assert abs(float(out[0]) - 0.35) < 1e-6
|
||||
|
||||
|
||||
def test_zero_when_wheels_roll_backward():
|
||||
out = wheel_glide_reward(_Env([-10.0, -10.0, -10.0, -10.0]), cap_speed=0.35)
|
||||
assert float(out[0]) == 0.0
|
||||
|
||||
|
||||
def test_nan_safe():
|
||||
out = wheel_glide_reward(_Env([float("nan"), 10.0, 10.0, 10.0]), cap_speed=0.35)
|
||||
assert float(out[0]) == 0.0
|
||||
Reference in New Issue
Block a user