Skip to content

Commit 5121bb5

Browse files
committed
chore: merge feat/enhance-tests into feat/reward-config-injection
- veryslow marker + comprehensive training integration tests (APPO, RSL-RL, MLX) - env reset+step tests with dict obs key/shape/reward/done verification - APPORunner early crash detection for dead collector subprocess - RslRlVecEnvWrapper: add "privileged" key to TensorDict for rsl_rl v5 compat Conflict resolution: - test_appo_runner: use APPOConfig() instead of deleted appo_config() factory - test_rsl_rl_runner: use PPOConfig() instead of deleted ppo_config() factory
2 parents 265e1ea + df55ca3 commit 5121bb5

10 files changed

Lines changed: 407 additions & 54 deletions

File tree

‎CLAUDE.md‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ make check # make format && make type
3535
make test # Run all non-slow tests (default)
3636
make test-cov # Run non-slow tests with coverage report
3737
make test-slow # Run slow integration tests (requires MuJoCo)
38+
make test-veryslow # Run full training iteration tests (minutes per test)
3839
make test-all # make check && make test-cov
3940
```
4041

@@ -56,7 +57,10 @@ uv run pytest -m "not slow"
5657
uv run pytest -m "not slow" --cov=unilab --cov-report=term-missing
5758

5859
# Slow integration tests (need MuJoCo installed)
59-
uv run pytest -m slow -v
60+
uv run pytest -m "slow and not veryslow" -v
61+
62+
# Very slow: full training iteration tests (minutes per test)
63+
uv run pytest -m veryslow -v
6064
```
6165

6266
## Test Structure
@@ -91,6 +95,9 @@ tests/
9195
Tests marked `@pytest.mark.slow` require a real MuJoCo environment and are excluded from CI
9296
by default. Run them locally when working on runner/learner code.
9397

98+
Tests marked `@pytest.mark.veryslow` run full training iterations (minutes per test). They are
99+
excluded by default even when running `make test-slow`. Run `make test-veryslow` explicitly.
100+
94101
## Testing
95102

96103
**New features must ship with tests.** When developing a new feature or refactoring, design and write comprehensive unit tests alongside the feature code — not as an afterthought. Tests should cover:

‎Makefile‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,11 @@ test-fast:
2929

3030
.PHONY: test-slow
3131
test-slow:
32-
uv run pytest -m "slow"
32+
uv run pytest -m "slow and not veryslow" -v
33+
34+
.PHONY: test-veryslow
35+
test-veryslow:
36+
uv run pytest -m "veryslow" -v
3337

3438
.PHONY: test-all
3539
test-all: check test-cov

‎pyproject.toml‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -77,8 +77,11 @@ exclude = [
7777

7878
[tool.pytest.ini_options]
7979
testpaths = ["tests"]
80-
markers = ["slow: requires real env or long-running processes"]
81-
addopts = "--tb=short"
80+
markers = [
81+
"slow: requires real env or long-running processes",
82+
"veryslow: full training iteration tests (minutes per test)",
83+
]
84+
addopts = "--tb=short -m 'not veryslow'"
8285

8386
[tool.coverage.run]
8487
source = ["unilab"]

‎scripts/train_rsl_rl.py‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -71,15 +71,13 @@ def __init__(self, env, device="cuda"):
7171

7272
def _obs_to_tensordict(self, obs: dict[str, np.ndarray]) -> TensorDict:
7373
actor = to_torch(obs["actor"], self.device)
74+
td = {"actor": actor}
7475
if "privileged" in obs:
75-
policy = to_torch(flatten_obs_dict(obs), self.device)
76+
td["privileged"] = to_torch(obs["privileged"], self.device)
77+
td["policy"] = to_torch(flatten_obs_dict(obs), self.device)
7678
else:
77-
policy = actor
78-
return TensorDict(
79-
{"policy": policy, "actor": actor},
80-
batch_size=self.num_envs,
81-
device=self.device,
82-
)
79+
td["policy"] = actor
80+
return TensorDict(td, batch_size=self.num_envs, device=self.device)
8381

8482
def step(self, actions):
8583
if isinstance(actions, torch.Tensor):

‎src/unilab/algos/torch/appo/runner.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,16 @@ def learn(
235235

236236
data_ready = shared_storage.wait_for_data(timeout=60.0)
237237
if not data_ready:
238+
# Check if the collector subprocess died — fail fast instead of
239+
# burning through remaining iterations with 60s timeouts each.
240+
if not self._check_collector_alive():
241+
self._drain_metrics(
242+
metrics_queue, reward_history, latest_reward_components, logger
243+
)
244+
raise RuntimeError(
245+
"APPO collector process died before producing data. "
246+
"Check stderr for [APPO WORKER CRASH] messages."
247+
)
238248
logger.log_status(
239249
f"[yellow]Warning: Timeout waiting for data at iteration {iteration}[/]"
240250
)

‎tests/algos/test_appo_runner.py‎

Lines changed: 26 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
"""Slow integration tests for APPORunner.
22
33
Requires MuJoCo to be installed. Run with:
4-
uv run pytest -m slow -v
4+
uv run pytest -m slow -v # init + close tests
5+
uv run pytest -m veryslow -v # full training iteration
56
"""
67

78
from __future__ import annotations
89

9-
import os
1010
import tempfile
1111

1212
import pytest
@@ -33,24 +33,33 @@ def test_appo_runner_init_no_crash(mock_env_name):
3333
runner.close()
3434

3535

36-
# @pytest.mark.slow
37-
# def test_appo_runner_learn_two_iterations(mock_env_name):
38-
# cfg = APPOConfig().to_dict()
39-
# cfg["num_envs"] = 4
40-
# cfg["steps_per_env"] = 4
36+
@pytest.mark.slow
37+
@pytest.mark.veryslow
38+
@pytest.mark.parametrize("env_name", ["Go2JoystickFlatTerrain"])
39+
def test_appo_runner_learn_two_iterations(env_name):
40+
"""APPO learn test must use a real env — DummyFlatTest is not registered in
41+
the collector subprocess (mp.spawn) so registry.make() would fail there."""
42+
cfg = APPOConfig().to_dict()
43+
cfg["num_envs"] = 128
44+
cfg["steps_per_env"] = 8
45+
# Small network for smoke test speed
46+
cfg["actor"]["hidden_dims"] = [64, 64]
47+
cfg["critic"]["hidden_dims"] = [64, 64]
48+
cfg["algorithm"]["num_learning_epochs"] = 1
49+
cfg["algorithm"]["num_mini_batches"] = 2
4150

42-
# runner = APPORunner(
43-
# env_name=mock_env_name,
44-
# env_cfg_overrides={},
45-
# rl_cfg=cfg,
46-
# num_envs=4,
47-
# steps_per_env=4,
48-
# )
51+
runner = APPORunner(
52+
env_name=env_name,
53+
env_cfg_overrides={},
54+
rl_cfg=cfg,
55+
num_envs=128,
56+
steps_per_env=8,
57+
)
4958

50-
# with tempfile.TemporaryDirectory() as tmpdir:
51-
# runner.learn(max_iterations=2, save_interval=0, log_dir=tmpdir)
59+
with tempfile.TemporaryDirectory() as tmpdir:
60+
runner.learn(max_iterations=2, save_interval=0, log_dir=tmpdir)
5261

53-
# runner.close()
62+
runner.close()
5463

5564

5665
@pytest.mark.slow

‎tests/algos/test_mlx_ppo.py‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -300,3 +300,100 @@ def test_ppo_trainer_update_decreases_loss():
300300
metrics = trainer.update(buf, iteration=0)
301301
assert np.isfinite(metrics["surrogate"])
302302
assert np.isfinite(metrics["value"])
303+
304+
305+
# ---------------------------------------------------------------------------
306+
# Full training iteration on real env (veryslow)
307+
# ---------------------------------------------------------------------------
308+
309+
310+
@pytest.mark.slow
311+
@pytest.mark.veryslow
312+
def test_mlx_ppo_one_iteration_real_env():
313+
"""Run 1 full MLX PPO iteration (collect rollout + update) on a real env."""
314+
_mujoco = pytest.importorskip("mujoco")
315+
316+
from unilab.base import registry
317+
from unilab.config import locomotion_params
318+
from unilab.utils.algo_utils import ensure_registries
319+
from unilab.utils.obs_utils import flatten_obs_dict
320+
321+
ensure_registries()
322+
323+
env_name = "Go2JoystickFlatTerrain"
324+
num_envs = 4
325+
num_steps = 8
326+
327+
cfg = locomotion_params.ppo_config(env_name)
328+
algo_cfg = cfg.algorithm
329+
330+
env = registry.make(env_name, num_envs=num_envs, sim_backend="mujoco")
331+
obs_dim = sum(env.obs_groups_spec.values())
332+
action_dim = env.action_space.shape[0]
333+
334+
model = MLPActorCritic(
335+
obs_dim=obs_dim,
336+
action_dim=action_dim,
337+
actor_hidden_dims=[64, 64],
338+
critic_hidden_dims=[64, 64],
339+
)
340+
ppo_cfg = PPOConfig(
341+
num_learning_epochs=1,
342+
num_mini_batches=2,
343+
clip_param=float(algo_cfg.clip_param),
344+
gamma=float(algo_cfg.gamma),
345+
lam=float(algo_cfg.lam),
346+
)
347+
trainer = PPOTrainer(model, ppo_cfg)
348+
349+
# Init env and get first obs
350+
if env.state is None:
351+
env.init_state()
352+
reset_indices = np.arange(num_envs, dtype=np.int32)
353+
_, obs_dict, _ = env.reset(reset_indices)
354+
obs = mx.array(flatten_obs_dict(obs_dict))
355+
356+
# Collect rollout
357+
buffer = RolloutBuffer(
358+
num_steps=num_steps,
359+
num_envs=num_envs,
360+
obs_dim=obs_dim,
361+
action_dim=action_dim,
362+
gamma=ppo_cfg.gamma,
363+
lam=ppo_cfg.lam,
364+
)
365+
366+
for _ in range(num_steps):
367+
actions, log_probs, values, mean, std = model.act(obs)
368+
mx.eval(actions, log_probs, values, mean, std)
369+
370+
env_actions = np.asarray(actions)
371+
state = env.step(env_actions)
372+
raw_obs = flatten_obs_dict(state.obs)
373+
rewards = mx.array(state.reward)
374+
dones = mx.array(state.done.astype(np.float32))
375+
376+
buffer.add(
377+
obs=obs,
378+
actions=actions,
379+
log_probs=log_probs,
380+
action_mean=mean,
381+
action_std=std,
382+
rewards=rewards,
383+
dones=dones,
384+
values=values,
385+
)
386+
obs = mx.array(raw_obs)
387+
388+
last_values = model.value(obs)
389+
buffer.compute_returns_and_advantages(last_values)
390+
391+
# PPO update
392+
metrics = trainer.update(buffer, iteration=0)
393+
assert isinstance(metrics, dict)
394+
assert np.isfinite(metrics["surrogate"]), (
395+
f"surrogate loss is not finite: {metrics['surrogate']}"
396+
)
397+
assert np.isfinite(metrics["value"]), f"value loss is not finite: {metrics['value']}"
398+
399+
env.close()

‎tests/algos/test_offpolicy_runner.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
"""Slow integration tests for OffPolicyRunner (FastSAC/FastTD3).
22
33
Requires MuJoCo to be installed. Run with:
4-
uv run pytest -m slow -v
4+
uv run pytest -m slow -v # init + close tests
5+
uv run pytest -m veryslow -v # full training iteration
56
"""
67

78
from __future__ import annotations
@@ -53,6 +54,7 @@ def test_offpolicy_runner_sac_init_no_crash(mock_env_name):
5354

5455

5556
@pytest.mark.slow
57+
@pytest.mark.veryslow
5658
def test_offpolicy_runner_sac_learn_two_iterations(mock_env_name):
5759
runner = _make_sac_runner(mock_env_name)
5860
with tempfile.TemporaryDirectory() as tmpdir:
@@ -109,6 +111,7 @@ def test_offpolicy_runner_td3_init_no_crash(mock_env_name):
109111

110112

111113
@pytest.mark.slow
114+
@pytest.mark.veryslow
112115
def test_offpolicy_runner_td3_learn_two_iterations(mock_env_name):
113116
runner = _make_td3_runner(mock_env_name)
114117
with tempfile.TemporaryDirectory() as tmpdir:

0 commit comments

Comments
 (0)