microduck/microduck_rl
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
microduck_rl/docs/superpowers/plans/2026-08-04-spin-env.md
Upstream Snapshot 47372443ff Import upstream snapshot d424a0c899f6b33cbd3daeb279913134349c0b63 
Upstream: https://github.com/pollen-robotics/microduck_rl
Upstream-Commit: d424a0c899f6b33cbd3daeb279913134349c0b63
Upstream-Branch: develop
2026-08-28 15:41:56 +08:00

1565 lines
64 KiB
Markdown
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# Spin Env Implementation Plan
> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.
**Goal:** Ajouter une tâche RL `Mjlab-Spin-Flat-MicroDuck` qui apprend au microduck sur rollers à faire ~2 tours anti-horaire sur place à ~6 rad/s puis à s'arrêter, geste cyclique piloté par la phase du slot bouton du runtime.
**Architecture :** Nouvel env cfg qui clone la structure de `microduck_roller_crouch_env_cfg.py` (robot rollers, obs 61D, DR complète, `GroundPickPhaseCommand`) mais remplace les rewards de pose par des rewards de **résultat** (suivi d'une vitesse de lacet cible en trapèze sur la phase), plus deux amorces de *shaping* décroissantes qui poussent vers le roulement différentiel. Toutes les nouvelles fonctions de reward vont dans `src/mjlab_microduck/tasks/mdp.py`, chacune découpée en une **fonction pure sur valeurs** (testable sans simulateur) + un **wrapper env** — c'est l'idiome déjà présent dans le repo (`crouch_glide_reward_from_values`).
**Tech Stack :** Python 3.12, mjlab 1.3.0, MuJoCo / mujoco-warp, torch, rsl_rl, pytest, uv.
**Spec de référence :** `docs/superpowers/specs/2026-08-04-spin-env-design.md`
## Global Constraints
- Toutes les valeurs numériques de l'enveloppe sont fixées par le spec : `SPIN_PERIOD = 4.0` s, `SPIN_RATE_MAX = 6.0` rad/s, `SPIN_ACCEL_END = 0.125`, `SPIN_HOLD_END = 0.525`, `SPIN_BRAKE_END = 0.650`. Ne pas les changer sans changer le spec.
- Sens de rotation : **anti-horaire uniquement**, ω_z cible **positive**.
- `ENABLE_SYMMETRY = False` et `symmetry_cfg=None` : l'augmentation de symétrie G/D détruirait un spin à sens unique.
- La reward `angular_momentum` (norme 3D du moment angulaire) doit être **absente** de l'env : elle combattrait le spin.
- L'obs actor doit rester à **61 dimensions** (layout identique à roller / ground_pick / crouch), sinon l'ONNX ne charge pas dans le slot du runtime.
- Les joints sont **toujours** résolus par nom / regex via `asset.find_joints(...)`, jamais par index en dur : les 4 roues passives sont intercalées dans l'ordre des joints du modèle rollers.
- La vitesse d'entrée est injectée via `cfg.events["reset_base"].params["velocity_range"]`, **jamais** via un `push_by_setting_velocity` en `mode="reset"` (régression NaN connue).
- Style du repo : commentaires en français dans les env cfg, docstrings des fonctions `mdp.py` en anglais ou français selon le voisinage, pas de `Co-Authored-By` dans les commits.
- Lancer les tests avec `uv run --with pytest pytest tests/ -q`.
---
### Task 1 : Enveloppe de phase (fonctions pures)
Les deux fonctions purement mathématiques qui portent toute la définition du geste : le profil de vitesse de lacet cible et la porte de shaping qui s'en déduit.
**Files:**
- Modify: `src/mjlab_microduck/tasks/mdp.py` (ajouter à la fin du fichier, après `RelativeHeadingVelocityCommandCfg` et les rewards de patinage)
- Test: `tests/test_spin.py` (créer)
**Interfaces:**
- Consumes: rien (fonctions pures sur `torch.Tensor`)
- Produces:
- `SPIN_PERIOD: float = 4.0`, `SPIN_RATE_MAX: float = 6.0`, `SPIN_ACCEL_END: float = 0.125`, `SPIN_HOLD_END: float = 0.525`, `SPIN_BRAKE_END: float = 0.650`
- `spin_rate_by_phase(phase: torch.Tensor, rate_max: float = 6.0, accel_end: float = 0.125, hold_end: float = 0.525, brake_end: float = 0.650) -> torch.Tensor`
- `spin_gate_by_phase(phase: torch.Tensor, rate_max: float = 6.0, accel_end: float = 0.125, hold_end: float = 0.525, brake_end: float = 0.650) -> torch.Tensor`
- [ ] **Step 1 : Écrire les tests qui échouent**
Créer `tests/test_spin.py` :
```python
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_is_two_turns():
# LE test qui protège la cible du spec : l'aire sous l'enveloppe sur un cycle
# de 4 s doit valoir ~4*pi rad = 2 tours. Enveloppe exacte = 12.6 rad,
# 4*pi = 12.566 -> tolérance 1 %.
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_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)
```
- [ ] **Step 2 : Lancer les tests pour vérifier qu'ils échouent**
Run: `uv run --with pytest pytest tests/test_spin.py -q`
Expected: FAIL avec `AttributeError: module 'mjlab_microduck.tasks.mdp' has no attribute 'spin_rate_by_phase'`
- [ ] **Step 3 : Implémenter les deux fonctions**
Ajouter à la fin de `src/mjlab_microduck/tasks/mdp.py` :
```python
# --------------------------------------------------------------------------- #
# Tâche SPIN — rotation rapide sur place sur rollers #
# --------------------------------------------------------------------------- #
# Enveloppe de phase : la commande du slot bouton porte une phase, qui pilote
# une VITESSE DE LACET cible en trapèze (et non une pose comme le crouch).
# [0, accel_end) 0.5 s 0 -> rate_max (lancement)
# [accel_end, hold_end) 1.6 s rate_max (régime)
# [hold_end, brake_end) 0.5 s rate_max -> 0 (freinage)
# [brake_end, 1.0) 1.4 s 0 (repos debout)
# Aire sous l'enveloppe sur un cycle de 4 s = 12.6 rad ~ 2 tours.
SPIN_PERIOD = 4.0
SPIN_RATE_MAX = 6.0
SPIN_ACCEL_END = 0.125
SPIN_HOLD_END = 0.525
SPIN_BRAKE_END = 0.650
def spin_rate_by_phase(
phase: torch.Tensor,
rate_max: float = SPIN_RATE_MAX,
accel_end: float = SPIN_ACCEL_END,
hold_end: float = SPIN_HOLD_END,
brake_end: float = SPIN_BRAKE_END,
) -> torch.Tensor:
"""Vitesse de lacet cible (rad/s, positive = anti-horaire) le long de la phase."""
w = torch.zeros_like(phase)
accel = phase < accel_end
w = torch.where(accel, rate_max * phase / accel_end, w)
hold = (phase >= accel_end) & (phase < hold_end)
w = torch.where(hold, torch.full_like(phase, rate_max), w)
brake = (phase >= hold_end) & (phase < brake_end)
w = torch.where(
brake, rate_max * (1.0 - (phase - hold_end) / (brake_end - hold_end)), w
)
return w
def spin_gate_by_phase(
phase: torch.Tensor,
rate_max: float = SPIN_RATE_MAX,
accel_end: float = SPIN_ACCEL_END,
hold_end: float = SPIN_HOLD_END,
brake_end: float = SPIN_BRAKE_END,
) -> torch.Tensor:
"""Porte de shaping dans [0,1] = enveloppe normalisée.
Vaut 0 sur tout le segment de repos : les amorces (ciseau des jambes,
différentiel des roues) ne s'appliquent que pendant lancement + régime, donc
le robot revient en station neutre avant de rendre la main à la policy roller.
"""
return spin_rate_by_phase(phase, rate_max, accel_end, hold_end, brake_end) / rate_max
```
- [ ] **Step 4 : Lancer les tests pour vérifier qu'ils passent**
Run: `uv run --with pytest pytest tests/test_spin.py -q`
Expected: PASS (6 tests)
- [ ] **Step 5 : Commit**
```bash
git add src/mjlab_microduck/tasks/mdp.py tests/test_spin.py
git commit -m "spin: enveloppe de phase (vitesse de lacet cible + porte de shaping)"
```
---
### Task 2 : Rewards de suivi de lacet et « sur place »
L'objectif principal (`spin_rate_track`), son bootstrap L1, et la pénalité qui garde le robot sur place. Cette tâche introduit aussi le **faux env** partagé par les tests des tâches 2 à 4, ce qui permet de tester les wrappers env sans lancer MuJoCo.
**Files:**
- Modify: `src/mjlab_microduck/tasks/mdp.py` (à la suite de la Task 1)
- Modify: `tests/test_spin.py`
**Interfaces:**
- Consumes: `spin_rate_by_phase`, `spin_gate_by_phase`, `SPIN_*` (Task 1)
- Produces:
- `spin_phase_from_command(cmd: torch.Tensor) -> torch.Tensor`
- `spin_rate_reward_from_values(omega_z: torch.Tensor, omega_target: torch.Tensor, std: float) -> torch.Tensor`
- `spin_rate_track(env, command_name="twist", std=1.5, rate_max=..., accel_end=..., hold_end=..., brake_end=..., asset_cfg=_DEFAULT_ASSET_CFG) -> torch.Tensor`
- `spin_rate_l1(env, command_name="twist", rate_max=..., accel_end=..., hold_end=..., brake_end=..., asset_cfg=_DEFAULT_ASSET_CFG) -> torch.Tensor`
- `spin_stay_in_place(env, asset_cfg=_DEFAULT_ASSET_CFG) -> torch.Tensor`
- [ ] **Step 1 : Écrire les tests qui échouent**
Ajouter à `tests/test_spin.py` (le faux env sert aussi aux tâches 3 et 4) :
```python
# ── 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 6 rad/s. Un robot qui tourne à 6 rad/s
# doit toucher 1.0 ; un robot immobile doit être largement en dessous.
ang = torch.tensor([[0.0, 0.0, 6.0], [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.01
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 à -6 rad/s (horaire) quand on demande +6 doit être pire qu'immobile
ang = torch.tensor([[0.0, 0.0, -6.0], [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():
ang = torch.tensor([[0.0, 0.0, 6.0], [0.0, 0.0, 2.0]])
env = _FakeEnv(
_FakeEntity(_FakeData(ang_vel_b=ang)), cmd=_phase_cmd([0.30, 0.30])
)
r = mdp.spin_rate_l1(env)
assert torch.allclose(r, torch.tensor([0.0, -4.0]), atol=1e-5)
# ── spin_stay_in_place ───────────────────────────────────────────────────────
def test_spin_stay_in_place_is_squared_planar_speed():
lin = torch.tensor([[0.0, 0.0, 0.0], [0.3, 0.4, 9.0]])
env = _FakeEnv(_FakeEntity(_FakeData(lin_vel_b=lin)))
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)
```
- [ ] **Step 2 : Lancer les tests pour vérifier qu'ils échouent**
Run: `uv run --with pytest pytest tests/test_spin.py -q`
Expected: FAIL — `AttributeError: ... has no attribute 'spin_phase_from_command'` (les 6 tests de la Task 1 continuent de passer)
- [ ] **Step 3 : Implémenter les rewards**
Ajouter à la suite dans `src/mjlab_microduck/tasks/mdp.py` :
```python
def spin_phase_from_command(cmd: torch.Tensor) -> torch.Tensor:
"""Récupère la phase [0,1) depuis la commande [cos(2πφ), sin(2πφ), 0] du slot."""
return (torch.atan2(cmd[:, 1], cmd[:, 0]) / (2 * torch.pi)) % 1.0
def _spin_target_rate(
env: ManagerBasedRlEnv,
command_name: str,
rate_max: float,
accel_end: float,
hold_end: float,
brake_end: float,
) -> torch.Tensor:
phase = spin_phase_from_command(env.command_manager.get_command(command_name))
return spin_rate_by_phase(phase, rate_max, accel_end, hold_end, brake_end)
def _spin_gate(
env: ManagerBasedRlEnv,
command_name: str,
rate_max: float,
accel_end: float,
hold_end: float,
brake_end: float,
) -> torch.Tensor:
phase = spin_phase_from_command(env.command_manager.get_command(command_name))
return spin_gate_by_phase(phase, rate_max, accel_end, hold_end, brake_end)
def spin_rate_reward_from_values(
omega_z: torch.Tensor, omega_target: torch.Tensor, std: float
) -> torch.Tensor:
"""Gaussienne sur l'erreur de vitesse de lacet (fonction pure, testable)."""
return torch.exp(-(((omega_z - omega_target) / std) ** 2))
def spin_rate_track(
env: ManagerBasedRlEnv,
command_name: str = "twist",
std: float = 1.5,
rate_max: float = SPIN_RATE_MAX,
accel_end: float = SPIN_ACCEL_END,
hold_end: float = SPIN_HOLD_END,
brake_end: float = SPIN_BRAKE_END,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Objectif principal du spin : suivre la vitesse de lacet cible ω*(φ).
ω_z est pris en repère corps (c'est ce que voit le gyro de l'IMU, donc ce que
la policy observe). Une rotation dans le mauvais sens est plus punie que
l'immobilité, la gaussienne étant centrée sur une cible positive.
"""
asset: Entity = env.scene[asset_cfg.name]
omega_z = asset.data.root_link_ang_vel_b[:, 2]
target = _spin_target_rate(env, command_name, rate_max, accel_end, hold_end, brake_end)
return spin_rate_reward_from_values(omega_z, target, std)
def spin_rate_l1(
env: ManagerBasedRlEnv,
command_name: str = "twist",
rate_max: float = SPIN_RATE_MAX,
accel_end: float = SPIN_ACCEL_END,
hold_end: float = SPIN_HOLD_END,
brake_end: float = SPIN_BRAKE_END,
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
) -> torch.Tensor:
"""Bootstrap L1 : gradient constant vers la cible même quand la gaussienne
de `spin_rate_track` sature loin de la cible. À utiliser avec un poids
POSITIF (la valeur retournée est déjà négative)."""
asset: Entity = env.scene[asset_cfg.name]
omega_z = asset.data.root_link_ang_vel_b[:, 2]
target = _spin_target_rate(env, command_name, rate_max, accel_end, hold_end, brake_end)
return -torch.abs(omega_z - target)
def spin_stay_in_place(
env: ManagerBasedRlEnv, asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG
) -> torch.Tensor:
"""Coût ‖v_xy‖² du tronc : tourner SUR PLACE, et tuer l'élan d'entrée.
Pas d'état de référence (contrairement à une dérive mesurée depuis le reset),
donc reste valide sur les 5 cycles d'un épisode. À utiliser avec un poids
NÉGATIF."""
asset: Entity = env.scene[asset_cfg.name]
v_xy = asset.data.root_link_lin_vel_b[:, :2]
return torch.sum(torch.square(v_xy), dim=1)
```
- [ ] **Step 4 : Lancer les tests pour vérifier qu'ils passent**
Run: `uv run --with pytest pytest tests/test_spin.py -q`
Expected: PASS (13 tests)
- [ ] **Step 5 : Commit**
```bash
git add src/mjlab_microduck/tasks/mdp.py tests/test_spin.py
git commit -m "spin: rewards de suivi de lacet (gaussienne + L1) et sur-place"
```
---
### Task 3 : Amorces du roulement différentiel
Les deux rewards qui injectent la physique connue : les patins roulent en sens opposés (`spin_wheel_differential`) et les deux lames restent au sol (`spin_grounded`). Toutes deux portées par la porte `gate(φ)`.
**Rappel des signes** (dérivé dans le spec) : pour une rotation anti-horaire, le patin gauche va vers l'**arrière** et le droit vers l'**avant** ; les 4 roues tournent positif en marche avant, donc **`ω_D − ω_G > 0`**.
**Files:**
- Modify: `src/mjlab_microduck/tasks/mdp.py` (à la suite de la Task 2)
- Modify: `tests/test_spin.py`
**Interfaces:**
- Consumes: `_spin_gate`, `SPIN_*` (Tasks 1-2), `_FakeEntity` / `_FakeSensor` / `_FakeEnv` (Task 2)
- Produces:
- `spin_wheel_differential_from_values(diff: torch.Tensor, gate: torch.Tensor, omega_scale: float) -> torch.Tensor`
- `spin_wheel_differential(env, command_name="twist", omega_scale=20.0, rate_max=..., accel_end=..., hold_end=..., brake_end=...) -> torch.Tensor`
- `spin_grounded(env, sensor_name: str, command_name="twist", rate_max=..., accel_end=..., hold_end=..., brake_end=...) -> torch.Tensor`
- [ ] **Step 1 : Écrire les tests qui échouent**
Ajouter à `tests/test_spin.py` :
```python
# ── 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)
```
- [ ] **Step 2 : Lancer les tests pour vérifier qu'ils échouent**
Run: `uv run --with pytest pytest tests/test_spin.py -q`
Expected: FAIL — `AttributeError: ... has no attribute 'spin_wheel_differential'`
- [ ] **Step 3 : Implémenter les deux rewards**
Ajouter à la suite dans `src/mjlab_microduck/tasks/mdp.py` :
```python
SPIN_WHEEL_OMEGA_SCALE = 20.0 # rad/s ; voir le calibrage dans le spec
def spin_wheel_differential_from_values(
diff: torch.Tensor, gate: torch.Tensor, omega_scale: float
) -> torch.Tensor:
"""Fonction pure : tanh du différentiel de roues, portée par gate, clampée ≥ 0."""
return gate * torch.tanh(torch.clamp(diff, min=0.0) / omega_scale)
def spin_wheel_differential(
env: ManagerBasedRlEnv,
command_name: str = "twist",
omega_scale: float = SPIN_WHEEL_OMEGA_SCALE,
rate_max: float = SPIN_RATE_MAX,
accel_end: float = SPIN_ACCEL_END,
hold_end: float = SPIN_HOLD_END,
brake_end: float = SPIN_BRAKE_END,
) -> torch.Tensor:
"""Récompense la rotation EN ROULEMENT (et non en patinage).
Pour un spin anti-horaire, le patin gauche recule et le droit avance ; les 4
roues tournant positif en marche avant, cela donne ω_D − ω_G > 0. Le tanh
sature à `omega_scale` pour éviter la course à la vitesse de roue.
"""
asset: Entity = env.scene["robot"]
lf_ids, _ = asset.find_joints("passive_LF_?wheel")
lr_ids, _ = asset.find_joints("passive_LR_?wheel")
rf_ids, _ = asset.find_joints("passive_RF_?wheel")
rr_ids, _ = asset.find_joints("passive_RR_?wheel")
vel = asset.data.joint_vel
omega_left = (vel[:, lf_ids[0]] + vel[:, lr_ids[0]]) / 2.0
omega_right = (vel[:, rf_ids[0]] + vel[:, rr_ids[0]]) / 2.0
gate = _spin_gate(env, command_name, rate_max, accel_end, hold_end, brake_end)
return spin_wheel_differential_from_values(
omega_right - omega_left, gate, omega_scale
)
def spin_grounded(
env: ManagerBasedRlEnv,
sensor_name: str,
command_name: str = "twist",
rate_max: float = SPIN_RATE_MAX,
accel_end: float = SPIN_ACCEL_END,
hold_end: float = SPIN_HOLD_END,
brake_end: float = SPIN_BRAKE_END,
) -> torch.Tensor:
"""Les deux lames au sol pendant le spin — empêche « je saute et je vrille ».
Variante de `grounded_reward` du swizzle, qui n'est pas réutilisable ici :
elle se pondère par cmd_x, qui vaut cos(2πφ) sur la commande de phase.
"""
from mjlab.sensor import ContactSensor
sensor: ContactSensor = env.scene[sensor_name]
contact_time = sensor.data.current_contact_time # (num_envs, num_feet)
assert contact_time is not None
n_contact = torch.sum((contact_time > 0.0).float(), dim=1)
grounded = (n_contact >= 2).float()
gate = _spin_gate(env, command_name, rate_max, accel_end, hold_end, brake_end)
return grounded * gate
```
- [ ] **Step 4 : Lancer les tests pour vérifier qu'ils passent**
Run: `uv run --with pytest pytest tests/test_spin.py -q`
Expected: PASS (18 tests)
- [ ] **Step 5 : Mesurer la demi-voie réelle et ajuster `omega_scale`**
Le spec fixe `omega_scale = 20.0` sur une demi-voie **estimée** à 0.03 m. Mesurer la vraie valeur sur le modèle rollers, à la pose HOME :
```bash
uv run python -c "
import mujoco, numpy as np
m = mujoco.MjModel.from_xml_path('src/mjlab_microduck/robot/microduck/robot_allcollisions_rollers.xml')
d = mujoco.MjData(m)
mujoco.mj_forward(m, d)
ys = {}
for name in ('left_foot', 'right_foot'):
sid = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_SITE, name)
ys[name] = d.site_xpos[sid][1]
half = abs(ys['left_foot'] - ys['right_foot']) / 2.0
print('sites y =', ys)
print('demi-voie =', round(half, 4), 'm')
print('differentiel attendu a 6 rad/s =', round(2 * 6.0 * half / 0.0175, 1), 'rad/s')
"
```
Si le différentiel attendu diffère de plus de 30 % de 20.0, mettre `SPIN_WHEEL_OMEGA_SCALE` à la valeur mesurée (arrondie à l'entier) et noter la mesure en commentaire au-dessus de la constante. Sinon laisser 20.0 et noter la mesure en commentaire. Dans les deux cas, le commentaire doit contenir la demi-voie mesurée.
- [ ] **Step 6 : Relancer les tests**
Run: `uv run --with pytest pytest tests/test_spin.py -q`
Expected: PASS (18 tests — les tests passent `omega_scale=20.0` explicitement, donc ils sont indépendants de la constante)
- [ ] **Step 7 : Commit**
```bash
git add src/mjlab_microduck/tasks/mdp.py tests/test_spin.py
git commit -m "spin: amorces du roulement differentiel (roues opposees + lames au sol)"
```
---
### Task 4 : Ciseau des jambes et tête partiellement libre
L'amorce de pose (`leg_antisymmetry`) et l'ajustement de `neck_joint_pos_l2` pour laisser `head_yaw` libre de servir de volant d'inertie.
**Piège de convention** : le robot a des conventions de signe **miroir** gauche/droite. Une pose *symétrique* satisfait `q_G + q_D ≈ 0` (c'est ce que mesure `leg_symmetry_reward`). Donc le **ciseau** (une jambe vers l'avant, l'autre vers l'arrière) satisfait `q_G ≈ q_D`, et se mesure par `−|q_G − q_D|`.
**Files:**
- Modify: `src/mjlab_microduck/tasks/mdp.py` (`neck_joint_pos_l2` à la ligne ~1237, puis ajout à la suite de la Task 3)
- Modify: `tests/test_spin.py`
**Interfaces:**
- Consumes: `_spin_gate`, `SPIN_*`, `_FakeEntity` / `_FakeEnv`
- Produces:
- `leg_antisymmetry(env, command_name="twist", asset_cfg=_DEFAULT_ASSET_CFG, joint_bases=("hip_pitch", "knee"), rate_max=..., accel_end=..., hold_end=..., brake_end=...) -> torch.Tensor`
- `neck_joint_pos_l2(env, asset_cfg=_NECK_JOINT_CFG, pattern: str = r".*(neck|head).*") -> torch.Tensor` (signature **élargie**, comportement par défaut inchangé)
- [ ] **Step 1 : Écrire les tests qui échouent**
Ajouter à `tests/test_spin.py` :
```python
# ── 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,
)
```
- [ ] **Step 2 : Lancer les tests pour vérifier qu'ils échouent**
Run: `uv run --with pytest pytest tests/test_spin.py -q`
Expected: FAIL — `AttributeError: ... has no attribute 'leg_antisymmetry'`
- [ ] **Step 3 : Ajouter le paramètre `pattern` à `neck_joint_pos_l2`**
Dans `src/mjlab_microduck/tasks/mdp.py`, remplacer la fonction existante (vers la ligne 1237) par :
```python
def neck_joint_pos_l2(
env: ManagerBasedRlEnv,
asset_cfg: SceneEntityCfg = _NECK_JOINT_CFG,
pattern: str = r".*(neck|head).*",
) -> torch.Tensor:
"""Penalize neck/head joint position deviation from default (L2 squared).
Uses find_joints() every call to avoid stale cached indices when the same
SceneEntityCfg singleton is reused across robots with different joint layouts
(e.g. walk robot vs rollers robot where passive wheels shift neck indices).
``pattern`` sélectionne les joints comptés (défaut : toute la nuque + la tête).
La tâche spin passe un motif qui EXCLUT `head_yaw`, pour laisser la tête servir
de volant d'inertie au lancement de la rotation.
"""
asset: Entity = env.scene[asset_cfg.name]
joint_ids, _ = asset.find_joints(pattern)
error = asset.data.joint_pos[:, joint_ids] - asset.data.default_joint_pos[:, joint_ids]
return torch.sum(torch.square(error), dim=1)
```
- [ ] **Step 4 : Implémenter `leg_antisymmetry`**
Ajouter à la suite de la Task 3 dans `src/mjlab_microduck/tasks/mdp.py` :
```python
def leg_antisymmetry(
env: ManagerBasedRlEnv,
command_name: str = "twist",
asset_cfg: SceneEntityCfg = _DEFAULT_ASSET_CFG,
joint_bases: tuple = ("hip_pitch", "knee"),
rate_max: float = SPIN_RATE_MAX,
accel_end: float = SPIN_ACCEL_END,
hold_end: float = SPIN_HOLD_END,
brake_end: float = SPIN_BRAKE_END,
) -> torch.Tensor:
"""Amorce le CISEAU des jambes (une avant / une arrière) pendant le spin.
Le robot a des conventions de signe MIROIR gauche/droite : une pose
symétrique satisfait q_G + q_D ≈ 0 (cf. `leg_symmetry_reward`), donc le
ciseau satisfait q_G ≈ q_D. On retourne `gate(φ) · (−mean|q_G − q_D|)` — à
utiliser avec un poids POSITIF, décroissant par curriculum : l'amorce
s'efface pour laisser la policy affiner son propre geste.
"""
asset: Entity = env.scene[asset_cfg.name]
left, right = [], []
for base in joint_bases:
li, _ = asset.find_joints([f"left_{base}"])
ri, _ = asset.find_joints([f"right_{base}"])
left.append(li[0])
right.append(ri[0])
lids = torch.tensor(left, device=env.device)
rids = torch.tensor(right, device=env.device)
q = asset.data.joint_pos
scissor = -torch.abs(q[:, lids] - q[:, rids]).mean(dim=-1)
gate = _spin_gate(env, command_name, rate_max, accel_end, hold_end, brake_end)
return gate * scissor
```
- [ ] **Step 5 : Lancer les tests pour vérifier qu'ils passent**
Run: `uv run --with pytest pytest tests/test_spin.py -q`
Expected: PASS (21 tests)
- [ ] **Step 6 : Vérifier la non-régression des autres tests**
`neck_joint_pos_l2` est utilisée par les envs roller / slope / swizzle ; sa signature a changé (ajout d'un paramètre avec défaut, donc compatible).
Run: `uv run --with pytest pytest tests/ -q`
Expected: PASS — aucun test existant cassé
- [ ] **Step 7 : Commit**
```bash
git add src/mjlab_microduck/tasks/mdp.py tests/test_spin.py
git commit -m "spin: amorce ciseau des jambes + head_yaw libre (pattern sur neck_joint_pos_l2)"
```
---
### Task 5 : Env cfg et enregistrement de la tâche
Le fichier d'environnement complet et son enregistrement, avec les tests de configuration.
**Files:**
- Create: `src/mjlab_microduck/tasks/microduck_spin_env_cfg.py`
- Modify: `src/mjlab_microduck/tasks/__init__.py`
- Test: `tests/test_spin_cfg.py` (créer)
**Interfaces:**
- Consumes: `spin_rate_track`, `spin_rate_l1`, `spin_stay_in_place`, `spin_wheel_differential`, `spin_grounded`, `leg_antisymmetry`, `neck_joint_pos_l2(pattern=...)`, `SPIN_PERIOD`, `SPIN_RATE_MAX`, `SPIN_ACCEL_END`, `SPIN_HOLD_END`, `SPIN_BRAKE_END` (Tasks 1-4)
- Produces:
- `make_microduck_spin_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg`
- `MicroduckSpinRlCfg: RslRlOnPolicyRunnerCfg`
- task id `Mjlab-Spin-Flat-MicroDuck`
- [ ] **Step 1 : Écrire les tests qui échouent**
Créer `tests/test_spin_cfg.py` :
```python
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
# 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 en Task 6 Step 1 (il faut
# construire l'env pour compter les dims ; 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
```
- [ ] **Step 2 : Lancer les tests pour vérifier qu'ils échouent**
Run: `uv run --with pytest pytest tests/test_spin_cfg.py -q`
Expected: FAIL avec `ModuleNotFoundError: No module named 'mjlab_microduck.tasks.microduck_spin_env_cfg'`
- [ ] **Step 3 : Créer le fichier d'environnement**
Créer `src/mjlab_microduck/tasks/microduck_spin_env_cfg.py` :
```python
"""Microduck SPIN task — rotation rapide sur place, sur rollers.
Geste cyclique déclenché au bouton A via le slot --ground-pick du runtime :
~2 tours anti-horaire à ~6 rad/s puis arrêt propre debout.
Hybride :
- physique / robot roller ← microduck_velocity_rollers_env_cfg.py
- machinerie phase cyclique ← microduck_roller_crouch_env_cfg.py
(commande GroundPickPhaseCommand : [cos(2πφ), sin(2πφ), 0], période 4 s)
Différence de fond avec le crouch : la phase pilote une VITESSE DE LACET cible
(objectif de résultat) et non une pose articulaire. Deux amorces décroissantes
poussent vers le roulement différentiel — le seul mécanisme physique certain sur
4 roues passives : patin gauche vers l'arrière, patin droit vers l'avant.
Obs 61D unifié → interchangeable au runtime avec roller / ground_pick / crouch.
Voir docs/superpowers/specs/2026-08-04-spin-env-design.md.
"""
import math
from copy import deepcopy
# La symétrie G/D transformerait un spin à gauche en spin à droite : interdit ici.
ENABLE_SYMMETRY = False
# DR — repris du roller env
ENABLE_COM_RANDOMIZATION = True
ENABLE_HEAD_COM_RANDOMIZATION = True
ENABLE_MASS_INERTIA_RANDOMIZATION = True
ENABLE_JOINT_FRICTION_RANDOMIZATION = True
ENABLE_ARMATURE_RANDOMIZATION = True
ENABLE_WHEEL_FRICTION_RANDOMIZATION = True
ENABLE_VELOCITY_PUSHES = True
ENABLE_IMU_ORIENTATION_RANDOMIZATION = True
ENABLE_ENCODER_BIAS = True
COM_RANDOMIZATION_RANGE = 0.003
HEAD_COM_RANDOMIZATION_RANGE = 0.003
MASS_INERTIA_RANDOMIZATION_RANGE = (0.95, 1.05)
JOINT_FRICTION_RANDOMIZATION_RANGE = (0.9, 1.1)
ARMATURE_RANDOMIZATION_RANGE = (0.9, 1.1)
VELOCITY_PUSH_INTERVAL_S = (3.0, 6.0)
VELOCITY_PUSH_RANGE = (-0.2, 0.2)
IMU_ORIENTATION_RANDOMIZATION_ANGLE = 6.0
ENCODER_BIAS_RANGE = (-0.015, 0.015)
# Le bouton peut être pressé à l'arrêt OU en roulement lent : la policy apprend
# à tuer l'élan résiduel avant/pendant le lancement de la rotation.
ENTRY_VELOCITY_X = (0.0, 0.3)
from mjlab.envs import ManagerBasedRlEnvCfg
from mjlab.envs.mdp import dr
from mjlab.envs.mdp.actions import JointPositionActionCfg
from mjlab.managers import (
CurriculumTermCfg,
EventTermCfg,
ObservationTermCfg,
RewardTermCfg,
TerminationTermCfg,
)
from mjlab.managers.scene_entity_config import SceneEntityCfg
from mjlab.rl import RslRlOnPolicyRunnerCfg, RslRlModelCfg
from mjlab.sensor import ContactMatch, ContactSensorCfg
from mjlab.tasks.velocity import mdp
from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
from mjlab.tasks.velocity.velocity_env_cfg import make_velocity_env_cfg
from mjlab.utils.noise import UniformNoiseCfg as Unoise
from mjlab_microduck.robot.microduck_constants import MICRODUCK_WALK_ROLLERS_ROBOT_CFG
from mjlab_microduck.tasks import mdp as microduck_mdp
from mjlab_microduck.tasks.microduck_velocity_env_cfg import HEAD_BODY_NAMES
from mjlab_microduck.tasks.symmetry import PpoWithSymmetryCfg, SYMMETRY_CFG
# Enveloppe de phase : constantes canoniques définies dans mdp.py.
SPIN_PERIOD = microduck_mdp.SPIN_PERIOD
_ENVELOPE = {
"rate_max": microduck_mdp.SPIN_RATE_MAX,
"accel_end": microduck_mdp.SPIN_ACCEL_END,
"hold_end": microduck_mdp.SPIN_HOLD_END,
"brake_end": microduck_mdp.SPIN_BRAKE_END,
}
# Nuque/tête tenues près du neutre SAUF head_yaw, laissé libre : il peut servir
# de volant d'inertie pour lancer la rotation.
NECK_PATTERN_NO_YAW = r"^(neck_pitch|head_pitch|head_roll)$"
def make_microduck_spin_env_cfg(play: bool = False) -> ManagerBasedRlEnvCfg:
"""Env spin sur rollers, piloté par la phase du slot ground-pick."""
feet_ground_cfg = ContactSensorCfg(
name="feet_ground_contact",
primary=ContactMatch(
mode="subtree",
pattern=r"^(ankle_l_v1|ankle_r_v1)$",
entity="robot",
),
secondary=ContactMatch(mode="body", pattern="terrain"),
fields=("found", "force"),
reduce="netforce",
num_slots=1,
track_air_time=True,
)
self_collision_cfg = ContactSensorCfg(
name="self_collision",
primary=ContactMatch(mode="subtree", pattern="trunk_base", entity="robot"),
secondary=ContactMatch(mode="subtree", pattern="trunk_base", entity="robot"),
fields=("found",),
reduce="none",
num_slots=1,
)
cfg = make_velocity_env_cfg()
cfg.scene.entities = {"robot": MICRODUCK_WALK_ROLLERS_ROBOT_CFG}
cfg.scene.sensors = (feet_ground_cfg, self_collision_cfg)
cfg.viewer.body_name = "trunk_base"
joint_pos_action = cfg.actions["joint_pos"]
assert isinstance(joint_pos_action, JointPositionActionCfg)
joint_pos_action.scale = 1.0
# === REWARDS ===
# ⚠️ angular_momentum n'est PAS gardée : elle pénalise la norme 3D du moment
# angulaire, donc elle combattrait directement le spin. body_ang_vel, elle,
# ne pénalise que x/y (« Don't penalize z-angular velocity » dans mjlab) →
# gardée, elle mate le ballant roulis/tangage sans gêner la rotation.
keep = {"upright", "body_ang_vel", "action_rate_l2"}
for name in list(cfg.rewards.keys()):
if name not in keep:
del cfg.rewards[name]
cfg.rewards["upright"].params["asset_cfg"].body_names = ("trunk_base",)
cfg.rewards["upright"].weight = 2.0
cfg.rewards["body_ang_vel"].params["asset_cfg"].body_names = ("trunk_base",)
cfg.rewards["body_ang_vel"].weight = -0.05
cfg.rewards["action_rate_l2"].weight = -1.0
# Objectif principal : suivre la vitesse de lacet cible ω*(φ) (trapèze).
cfg.rewards["spin_rate_track"] = RewardTermCfg(
func=microduck_mdp.spin_rate_track,
weight=6.0,
params={"command_name": "twist", "std": 1.5, **_ENVELOPE},
)
# Bootstrap L1 : gradient constant quand la gaussienne sature loin de la cible.
cfg.rewards["spin_rate_l1"] = RewardTermCfg(
func=microduck_mdp.spin_rate_l1,
weight=0.5,
params={"command_name": "twist", **_ENVELOPE},
)
# Tourner SUR PLACE, et tuer l'élan d'entrée.
cfg.rewards["spin_stay_in_place"] = RewardTermCfg(
func=microduck_mdp.spin_stay_in_place,
weight=-1.0,
params={},
)
# Amorce 1 : tourner EN ROULEMENT (patins en sens opposés), pas en patinage.
cfg.rewards["spin_wheel_differential"] = RewardTermCfg(
func=microduck_mdp.spin_wheel_differential,
weight=1.0,
params={
"command_name": "twist",
"omega_scale": microduck_mdp.SPIN_WHEEL_OMEGA_SCALE,
**_ENVELOPE,
},
)
# Amorce 2 : ciseau des jambes (décroît par curriculum, voir plus bas).
cfg.rewards["leg_antisymmetry"] = RewardTermCfg(
func=microduck_mdp.leg_antisymmetry,
weight=1.0,
params={
"command_name": "twist",
"joint_bases": ("hip_pitch", "knee"),
**_ENVELOPE,
},
)
# Les deux lames au sol pendant le spin (pas de vrille en l'air).
cfg.rewards["spin_grounded"] = RewardTermCfg(
func=microduck_mdp.spin_grounded,
weight=0.5,
params={
"sensor_name": "feet_ground_contact",
"command_name": "twist",
**_ENVELOPE,
},
)
# Stabilité / sim2real
cfg.rewards["feet_flat"] = RewardTermCfg(
func=microduck_mdp.feet_flat_penalty,
weight=-2.0,
params={
"asset_cfg": SceneEntityCfg("robot", site_names=("left_foot", "right_foot")),
"sensor_name": "feet_ground_contact",
},
)
cfg.rewards["self_collisions"] = RewardTermCfg(
func=mdp.self_collision_cost,
weight=-1.0,
params={"sensor_name": "self_collision"},
)
cfg.rewards["neck_action_rate_l2"] = RewardTermCfg(
func=microduck_mdp.neck_action_rate_l2, weight=-0.5
)
cfg.rewards["neck_joint_pos_l2"] = RewardTermCfg(
func=microduck_mdp.neck_joint_pos_l2,
weight=-0.2,
params={"pattern": NECK_PATTERN_NO_YAW},
)
cfg.rewards["joint_torques_l2"] = RewardTermCfg(
func=microduck_mdp.joint_torques_l2, weight=-1e-3
)
# === TERMINATIONS ===
cfg.terminations["nan_state"] = TerminationTermCfg(
func=microduck_mdp.robot_state_is_nan, time_out=False,
)
# === EVENTS ===
cfg.events["reset_action_history"] = EventTermCfg(
func=microduck_mdp.reset_action_history, mode="reset",
)
del cfg.events["foot_friction"]
if ENABLE_VELOCITY_PUSHES:
cfg.events["push_robot"] = EventTermCfg(
func=mdp.push_by_setting_velocity,
mode="interval",
interval_range_s=VELOCITY_PUSH_INTERVAL_S,
params={
"velocity_range": {"x": VELOCITY_PUSH_RANGE, "y": VELOCITY_PUSH_RANGE},
"asset_cfg": SceneEntityCfg("robot"),
},
)
cfg.events["reset_base"].params["pose_range"]["z"] = (0.1335, 0.1435)
# Élan d'entrée : injecté via reset_root_state_uniform (état par défaut PROPRE
# + range), et NON via push_by_setting_velocity en mode reset, qui additionne à
# une vitesse racine potentiellement divergente et fait exploser le free-joint
# de la base -> NaN. Régression connue du roller_crouch.
cfg.events["reset_base"].params["velocity_range"] = {"x": ENTRY_VELOCITY_X}
if ENABLE_WHEEL_FRICTION_RANDOMIZATION:
cfg.events["randomize_wheel_friction"] = EventTermCfg(
func=dr.dof_frictionloss,
mode="reset",
params={
"asset_cfg": SceneEntityCfg("robot", joint_names=(r"^passive_.*",)),
"operation": "abs",
"ranges": (0.000, 0.000),
},
)
if ENABLE_COM_RANDOMIZATION:
cfg.events["randomize_com"] = EventTermCfg(
func=dr.body_ipos, mode="reset",
params={
"asset_cfg": SceneEntityCfg("robot", body_names=("trunk_base",)),
"operation": "add",
"ranges": (-COM_RANDOMIZATION_RANGE, COM_RANDOMIZATION_RANGE),
},
)
if ENABLE_HEAD_COM_RANDOMIZATION:
cfg.events["randomize_head_com"] = EventTermCfg(
func=dr.body_ipos, mode="reset",
params={
"asset_cfg": SceneEntityCfg("robot", body_names=HEAD_BODY_NAMES),
"operation": "add",
"ranges": (-HEAD_COM_RANDOMIZATION_RANGE, HEAD_COM_RANDOMIZATION_RANGE),
},
)
if ENABLE_MASS_INERTIA_RANDOMIZATION:
_mi_lo, _mi_hi = MASS_INERTIA_RANDOMIZATION_RANGE
cfg.events["randomize_mass_inertia"] = EventTermCfg(
func=dr.pseudo_inertia, mode="startup",
params={
"asset_cfg": SceneEntityCfg("robot", body_names=("trunk_base",)),
"alpha_range": (math.log(_mi_lo) / 2.0, math.log(_mi_hi) / 2.0),
},
)
if ENABLE_JOINT_FRICTION_RANDOMIZATION:
cfg.events["randomize_joint_friction"] = EventTermCfg(
func=microduck_mdp.randomize_bam_friction, mode="reset",
params={
"asset_cfg": SceneEntityCfg("robot"),
"scale_range": JOINT_FRICTION_RANDOMIZATION_RANGE,
},
)
if ENABLE_ARMATURE_RANDOMIZATION:
cfg.events["randomize_armature"] = EventTermCfg(
func=dr.joint_armature, mode="reset",
params={
"asset_cfg": SceneEntityCfg("robot", joint_names=(r"^(?!passive_).*",)),
"operation": "scale",
"ranges": ARMATURE_RANDOMIZATION_RANGE,
},
)
# === OBSERVATIONS (unified 61D layout) ===
del cfg.observations["actor"].terms["base_lin_vel"]
del cfg.observations["critic"].terms["foot_height"]
del cfg.observations["actor"].terms["height_scan"]
del cfg.observations["critic"].terms["height_scan"]
cfg.observations["critic"].terms["base_lin_vel"] = ObservationTermCfg(
func=mdp.base_lin_vel, scale=1.0,
)
gravity_term_name = "projected_gravity"
cfg.observations["actor"].terms[gravity_term_name] = deepcopy(
cfg.observations["actor"].terms[gravity_term_name]
)
cfg.observations["actor"].terms["base_ang_vel"] = deepcopy(
cfg.observations["actor"].terms["base_ang_vel"]
)
cfg.observations["actor"].terms["base_ang_vel"].delay_min_lag = 0
cfg.observations["actor"].terms["base_ang_vel"].delay_max_lag = 1
cfg.observations["actor"].terms["base_ang_vel"].delay_update_period = 64
cfg.observations["actor"].terms[gravity_term_name].delay_min_lag = 0
cfg.observations["actor"].terms[gravity_term_name].delay_max_lag = 1
cfg.observations["actor"].terms[gravity_term_name].delay_update_period = 64
cfg.observations["actor"].terms["base_ang_vel"].noise = Unoise(n_min=-0.03, n_max=0.03)
cfg.observations["actor"].terms[gravity_term_name].noise = Unoise(n_min=-0.01, n_max=0.01)
cfg.observations["actor"].terms["joint_pos"].noise = Unoise(n_min=-0.001, n_max=0.001)
cfg.observations["actor"].terms["joint_vel"].noise = Unoise(n_min=-0.25, n_max=0.25)
if ENABLE_IMU_ORIENTATION_RANDOMIZATION:
av = cfg.observations["actor"].terms["base_ang_vel"]
av.func = microduck_mdp.base_ang_vel_imu_misaligned
av.params = {"max_angle_deg": IMU_ORIENTATION_RANDOMIZATION_ANGLE}
g = cfg.observations["actor"].terms[gravity_term_name]
g.func = microduck_mdp.projected_gravity_imu_misaligned
g.params = {"max_angle_deg": IMU_ORIENTATION_RANDOMIZATION_ANGLE}
cfg.observations["actor"].terms["joint_vel"] = deepcopy(
cfg.observations["actor"].terms["joint_vel"]
)
cfg.observations["actor"].terms["joint_vel"].delay_min_lag = 1
cfg.observations["actor"].terms["joint_vel"].delay_max_lag = 1
cfg.observations["actor"].terms["joint_vel"].delay_update_period = 0
passive_excluded = SceneEntityCfg("robot", joint_names=(r"^(?!passive_).*",))
for grp in ("actor", "critic"):
for term in ("joint_pos", "joint_vel"):
cfg.observations[grp].terms[term] = deepcopy(cfg.observations[grp].terms[term])
cfg.observations[grp].terms[term].params["asset_cfg"] = deepcopy(passive_excluded)
if ENABLE_ENCODER_BIAS:
cfg.events["encoder_bias"].params["bias_range"] = ENCODER_BIAS_RANGE
cfg.observations["actor"].terms["joint_pos"].params["biased"] = True
cfg.observations["critic"].terms["joint_pos"].params["biased"] = False
else:
cfg.events.pop("encoder_bias", None)
wheel_cfg = SceneEntityCfg("robot", joint_names=(r"^passive_.*",))
cfg.observations["critic"].terms["wheel_vel"] = ObservationTermCfg(
func=mdp.joint_vel_rel, scale=1.0, params={"asset_cfg": wheel_cfg},
)
for group in ("actor", "critic"):
cfg.observations[group].terms["head_command"] = ObservationTermCfg(
func=microduck_mdp.zero_command_padding, params={"dim": 4},
)
cfg.observations[group].terms["body_command"] = ObservationTermCfg(
func=microduck_mdp.zero_command_padding, params={"dim": 6},
)
# === COMMAND: phase (comme ground_pick / roller_crouch) ===
command: UniformVelocityCommandCfg = cfg.commands["twist"]
command.rel_standing_envs = 0.0
command.rel_heading_envs = 0.0
# period=4.0 = défaut de --ground-pick-period (rien à passer au runtime) ;
# randomize_phase=False -> chaque épisode démarre debout à phase 0, comme le
# bouton au déploiement. Épisode 20 s = 5 cycles complets du geste.
cfg.commands["twist"] = microduck_mdp.GroundPickPhaseCommandCfg(
**{
**vars(command),
"class_type": microduck_mdp.GroundPickPhaseCommand,
"period": SPIN_PERIOD,
"randomize_phase": False,
}
)
cfg.scene.terrain.terrain_type = "plane"
cfg.scene.terrain.terrain_generator = None
# === CURRICULUM ===
del cfg.curriculum["terrain_levels"]
del cfg.curriculum["command_vel"]
cfg.curriculum["action_rate_weight"] = CurriculumTermCfg(
func=microduck_mdp.reward_weight,
params={
"reward_name": "action_rate_l2",
"weight_stages": [
{"step": 0, "weight": -0.5},
{"step": 250 * 24, "weight": -0.8},
{"step": 500 * 24, "weight": -1.0},
],
},
)
# L'amorce ciseau s'efface : elle lance le bon mécanisme puis laisse la policy
# affiner son propre geste (fréquence de pompage libre).
cfg.curriculum["leg_antisym_weight"] = CurriculumTermCfg(
func=microduck_mdp.reward_weight,
params={
"reward_name": "leg_antisymmetry",
"weight_stages": [
{"step": 0, "weight": 1.0},
{"step": 1500 * 24, "weight": 0.5},
{"step": 3000 * 24, "weight": 0.25},
],
},
)
if ENABLE_COM_RANDOMIZATION:
cfg.curriculum["com_range"] = CurriculumTermCfg(
func=microduck_mdp.com_range_curriculum,
params={
"event_name": "randomize_com",
"range_stages": [
{"step": 0, "range": 0.003},
{"step": 500 * 24, "range": 0.005},
{"step": 1000 * 24, "range": 0.01},
],
},
)
if ENABLE_HEAD_COM_RANDOMIZATION:
cfg.curriculum["head_com_range"] = CurriculumTermCfg(
func=microduck_mdp.com_range_curriculum,
params={
"event_name": "randomize_head_com",
"range_stages": [
{"step": 0, "range": 0.003},
{"step": 500 * 24, "range": 0.005},
{"step": 1000 * 24, "range": 0.01},
],
},
)
return cfg
MicroduckSpinRlCfg = RslRlOnPolicyRunnerCfg(
actor=RslRlModelCfg(
hidden_dims=(512, 256, 128),
activation="elu",
obs_normalization=True,
distribution_cfg={
"class_name": "GaussianDistribution",
"init_std": 1.0,
"std_type": "scalar",
},
),
critic=RslRlModelCfg(
hidden_dims=(512, 256, 128),
activation="elu",
obs_normalization=True,
),
algorithm=PpoWithSymmetryCfg(
value_loss_coef=1.0,
use_clipped_value_loss=True,
clip_param=0.2,
entropy_coef=0.01,
num_learning_epochs=5,
num_mini_batches=4,
learning_rate=1.0e-3,
schedule="adaptive",
gamma=0.99,
lam=0.95,
desired_kl=0.01,
max_grad_norm=1.0,
symmetry_cfg=SYMMETRY_CFG if ENABLE_SYMMETRY else None,
),
wandb_project="mjlab_microduck",
experiment_name="spin",
run_name="spin",
save_interval=250,
num_steps_per_env=24,
max_iterations=8_000,
)
```
- [ ] **Step 4 : Enregistrer la tâche**
Dans `src/mjlab_microduck/tasks/__init__.py`, ajouter l'import après le bloc `microduck_roller_slope_env_cfg` (vers la ligne 66) :
```python
from .microduck_spin_env_cfg import (
make_microduck_spin_env_cfg,
MicroduckSpinRlCfg,
)
```
Puis l'enregistrement à la fin du fichier, après le bloc `Mjlab-RollerSlope-Flat-MicroDuck` :
```python
register_mjlab_task(
task_id="Mjlab-Spin-Flat-MicroDuck",
env_cfg=make_microduck_spin_env_cfg(),
play_env_cfg=make_microduck_spin_env_cfg(play=True),
rl_cfg=MicroduckSpinRlCfg,
runner_cls=MicroduckOnPolicyRunner,
)
print("✓ Spin task registered: Mjlab-Spin-Flat-MicroDuck")
```
- [ ] **Step 5 : Lancer les tests de config**
Run: `uv run --with pytest pytest tests/test_spin_cfg.py -q`
Expected: PASS (8 tests)
- [ ] **Step 6 : Vérifier que la tâche est bien enregistrée**
Run: `uv run python -c "import mjlab_microduck.tasks"`
Expected: la sortie contient `✓ Spin task registered: Mjlab-Spin-Flat-MicroDuck`, sans exception
- [ ] **Step 7 : Vérifier la non-régression de toute la suite**
Run: `uv run --with pytest pytest tests/ -q`
Expected: PASS
- [ ] **Step 8 : Commit**
```bash
git add src/mjlab_microduck/tasks/microduck_spin_env_cfg.py src/mjlab_microduck/tasks/__init__.py tests/test_spin_cfg.py
git commit -m "spin: env cfg Mjlab-Spin-Flat-MicroDuck + enregistrement"
```
---
### Task 6 : Vérification en simulation (smoke run)
Les tests unitaires ne prouvent pas que l'env tourne réellement : les 61D de l'obs, les capteurs de contact, les résolutions de joints et l'absence de NaN ne se voient qu'en lançant le simulateur. Cette tâche est la porte de sortie du plan.
**Files:**
- Modify: `docs/superpowers/specs/2026-08-04-spin-env-design.md` (noter la mesure de demi-voie et le résultat du smoke run)
**Interfaces:**
- Consumes: la tâche enregistrée `Mjlab-Spin-Flat-MicroDuck` (Task 5)
- Produces: rien de code — une vérification et une note
- [ ] **Step 1 : Vérifier la dimension réelle de l'obs actor**
Run:
```bash
uv run python -c "
import mjlab_microduck.tasks # enregistre les tâches
from mjlab_microduck.tasks.microduck_spin_env_cfg import make_microduck_spin_env_cfg
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()
print('spin actor terms:', list(spin.observations['actor'].terms.keys()))
print('crouch actor terms:', list(crouch.observations['actor'].terms.keys()))
assert list(spin.observations['actor'].terms.keys()) == list(crouch.observations['actor'].terms.keys()), 'layout obs != crouch -> ONNX ne chargera pas dans le slot'
print('OK: layout obs actor identique au crouch')
"
```
Expected: `OK: layout obs actor identique au crouch`. Si les listes diffèrent, corriger le bloc OBSERVATIONS de l'env cfg avant de continuer — c'est la condition pour que l'ONNX charge dans le slot du runtime.
- [ ] **Step 2 : Lancer un entraînement très court avec garde NaN**
Run:
```bash
uv run train Mjlab-Spin-Flat-MicroDuck \
--env.scene.num-envs 64 \
--agent.max_iterations 5 \
--enable-nan-guard
```
Expected: 5 itérations sans exception, sans NaN, et les logs `Episode_Reward/` listent bien `spin_rate_track`, `spin_rate_l1`, `spin_stay_in_place`, `spin_wheel_differential`, `spin_grounded`, `leg_antisymmetry`.
Si des NaN apparaissent : les dumps sont dans `/tmp/mjlab/nan_dumps/`. Le suspect n°1 est la vitesse d'entrée — vérifier qu'elle passe bien par `reset_base.velocity_range` et pas par un push en `mode="reset"`.
- [ ] **Step 3 : Lancer un entraînement de calibrage (500 itérations)**
Run:
```bash
uv run train Mjlab-Spin-Flat-MicroDuck --env.scene.num-envs 4096 --agent.max_iterations 500
```
Expected: `Episode_Reward/spin_rate_track` **monte** sur les 500 itérations. C'est le signal que l'objectif est apprenable. Noter la valeur atteinte.
Si la courbe reste plate, appliquer le « Plan B » du spec dans l'ordre : (1) curriculum de vitesse 3 → 6 rad/s, (2) monter `spin_wheel_differential` et retarder la décroissance de `leg_antisymmetry`, (3) élargir `std` de 1.5 à 2.5.
- [ ] **Step 4 : Regarder le geste**
Run: `uv run scripts/play_latest.py`
Expected: le robot tente une rotation anti-horaire sur place. À 500 itérations le geste sera brut ; ce qu'on vérifie c'est qu'il tourne **dans le bon sens**, qu'il ne part pas en translation, et qu'il ne finit pas systématiquement par tomber.
- [ ] **Step 5 : Noter les mesures dans le spec**
Ajouter une section `## Résultats de la vérification initiale` à la fin de `docs/superpowers/specs/2026-08-04-spin-env-design.md`, contenant : la demi-voie mesurée (Task 3 Step 5), la valeur retenue pour `omega_scale`, le résultat du smoke run 5 itérations, et la valeur de `Episode_Reward/spin_rate_track` à 500 itérations.
- [ ] **Step 6 : Commit**
```bash
git add docs/superpowers/specs/2026-08-04-spin-env-design.md
git commit -m "spin: notes de verification initiale (demi-voie, smoke run, 500 it.)"
```
---
## Notes pour l'implémenteur
- **Ordre obligatoire** : Task 1 → 2 → 3 → 4 → 5 → 6. Les tâches 2 à 4 dépendent du faux env créé dans la Task 2 ; la Task 5 consomme toutes les fonctions des tâches 1 à 4.
- **Le piège n°1 de cette tâche** est la reward `angular_momentum` : si elle reste dans l'env, elle pénalise la norme 3D du moment angulaire et le spin ne décollera jamais. Le test `test_angular_momentum_reward_is_removed` la garde.
- **Le piège n°2** est la convention de signe miroir gauche/droite du robot : le ciseau se mesure par `|q_G − q_D|`, pas par `|q_G + q_D|`. Comparer avec `leg_symmetry_reward` (ligne ~3769 de `mdp.py`) en cas de doute.
- **Le piège n°3** est la parité d'obs 61D : toute divergence par rapport au layout du crouch rend l'ONNX inutilisable dans le slot du runtime. La Task 6 Step 1 la vérifie explicitement.
- Ne pas activer la symétrie PPO, quelle que soit la tentation d'accélérer l'apprentissage.
- **Hors périmètre de ce plan** : l'entraînement complet (8000 itérations), l'export ONNX et le déploiement dans le slot du runtime. Les commandes sont dans le spec ; le plan s'arrête quand l'env est vérifié apprenable (Task 6).
Reference in New Issue View Git Blame Copy Permalink