From 0a43bc4be22bb4fc0bcb3d887c469c287f9e1a9a Mon Sep 17 00:00:00 2001 From: lemon-star608 Date: Mon, 26 Oct 2026 16:41:47 +0800 Subject: [PATCH 1/6] fix(base): bind mapped bodies to backend state rows --- src/unilab/base/entity.py | 49 ++++++++++--- tests/base/test_entity_facade.py | 121 +++++++++++++++++++++++++++++++ 2 files changed, 161 insertions(+), 9 deletions(-) diff --git a/src/unilab/base/entity.py b/src/unilab/base/entity.py index 407cce505..ee430863d 100644 --- a/src/unilab/base/entity.py +++ b/src/unilab/base/entity.py @@ -2414,7 +2414,6 @@ def _validate_mapped_state_addresses( immutable scene layout rather than copying live CUDA state to host. """ assert physical_entity is not None - owner = layout.get_entity(physical_entity) if pos_ids is not None and pos_ids.size and int(np.max(pos_ids)) >= layout.nq: raise ValueError( f"Entity '{self.name}' joint position index {int(np.max(pos_ids))} " @@ -2425,20 +2424,52 @@ def _validate_mapped_state_addresses( f"Entity '{self.name}' joint velocity index {int(np.max(vel_ids))} " f"exceeds mapped qvel width {layout.nv}" ) + + # SceneLayout body IDs are model addresses and may include unowned world + # rows. Backend body IDs address public body-state storage and need not + # use the same origin. Bind the two cold-path layouts by public name, + # then keep using backend rows for state reads, reset and randomization. + layout_body_names = tuple( + f"{entity.name}/{body_name}" + for entity in layout.entities + for body_name in entity.body_names + ) + layout_body_ids = layout.get_body_ids(layout_body_names) + body_state_rows = self._resolve_ids( + "mapped body state", layout_body_names, self._backend.get_body_ids + ) + body_state_width = int(np.max(body_state_rows)) + 1 if body_state_rows.size else 0 + row_by_layout_id = dict(zip(layout_body_ids, body_state_rows, strict=True)) for label, ids in (("root body", root_body_ids), ("body", body_ids)): - if ids is not None and ids.size and int(np.max(ids)) >= layout.nbody: + if ids is not None and ids.size and int(np.max(ids)) >= body_state_width: raise ValueError( - f"Entity '{self.name}' {label} id {int(np.max(ids))} exceeds " - f"mapped body count {layout.nbody}" + f"Entity '{self.name}' {label} row {int(np.max(ids))} exceeds " + f"mapped body-state count {body_state_width}" ) - expected_root = owner.body_ids[owner.body_names.index(owner.root_body)] if root_body_ids is not None and len(root_body_ids) != 1: raise ValueError(f"Entity '{self.name}' root body selection must contain one row") - if root_body_ids is not None and int(root_body_ids[0]) != expected_root: - raise ValueError( - f"Entity '{self.name}' selected root body {int(root_body_ids[0])}; " - f"mapped layout declares {expected_root}" + + selections = ( + ( + "root body", + (self.root_body_name,) if self.root_body_name is not None else None, + root_body_ids, + ), + ("body", self._body_names, body_ids), + ) + for label, names, ids in selections: + if ids is None: + continue + assert names is not None + expected_rows = np.asarray( + [row_by_layout_id[body_id] for body_id in layout.get_body_ids(names)], + dtype=np.intp, ) + if not np.array_equal(ids, expected_rows): + raise ValueError( + f"Entity '{self.name}' selected {label} rows {ids.tolist()}; " + f"mapped body-state layout declares {expected_rows.tolist()}" + ) def _validate_body_state( self, diff --git a/tests/base/test_entity_facade.py b/tests/base/test_entity_facade.py index 0b42a71f7..ba03359d6 100644 --- a/tests/base/test_entity_facade.py +++ b/tests/base/test_entity_facade.py @@ -13,9 +13,11 @@ MocapPoseBinding, RootStateLayout, SimBackend, + TensorExecution, TrackedBodyStateViews, ) from unisim.dr.types import DomainRandomizationCapabilities +from unisim.scene_layout import CompiledSceneLayout, EntityLayout import unilab.base.entity as entity_module from unilab.assets import ASSETS_ROOT_PATH @@ -894,3 +896,122 @@ def test_scene_exposes_read_only_community_entities_and_zero_origins() -> None: assert origins.shape == (3, 3) torch.testing.assert_close(origins, torch.zeros_like(origins)) assert not origins.requires_grad + + +@pytest.mark.parametrize("excluded_world_rows", [0, 1], ids=["whole-mjcf", "packed-state"]) +def test_mapped_multi_entity_body_addresses_use_backend_state_rows( + excluded_world_rows: int, +) -> None: + def physical_entity(name: str, first_body_id: int) -> EntityLayout: + return EntityLayout( + name=name, + kind="rigid", + root_mode="fixed", + root_body="wrist", + body_names=("wrist", "finger"), + body_ids=(first_body_id, first_body_id + 1), + body_parent_names=(None, "wrist"), + joints=(), + actuator_names=(), + actuator_joint_names=(), + actuator_indices=(), + ) + + layout = CompiledSceneLayout( + entities=(physical_entity("left_hand", 1), physical_entity("right_hand", 3)), + nq=0, + nv=0, + nu=0, + nbody=5, + ) + + class Backend: + backend_type = "fake" + num_envs = 2 + num_actuators = 0 + + def __init__(self) -> None: + self.model_ids = { + f"{entity.name}/{body_name}": body_id + for entity in layout.entities + for body_name, body_id in zip(entity.body_names, entity.body_ids, strict=True) + } + self.state_rows = { + name: body_id - excluded_world_rows for name, body_id in self.model_ids.items() + } + self.body_pos = np.zeros( + (self.num_envs, layout.nbody - excluded_world_rows, 3), dtype=np.float32 + ) + for name, row in self.state_rows.items(): + self.body_pos[:, row, 0] = self.model_ids[name] + + def get_tensor_capabilities(self): + from unisim.backend.base import TensorDataPlane, TensorLifecycleCapabilities + + return TensorLifecycleCapabilities( + execution=TensorExecution.DEVICE_RESIDENT, + data_plane=TensorDataPlane.DIRECT, + stream_event_ownership="test", + torch_devices=("cpu",), + ) + + def get_entity_default_state(self, name: str): + return { + "root_pose": np.tile( + np.asarray([0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0], dtype=np.float32), + (self.num_envs, 1), + ), + "root_velocity": np.zeros((self.num_envs, 6), dtype=np.float32), + "joint_positions": np.zeros((self.num_envs, 0), dtype=np.float32), + "joint_velocities": np.zeros((self.num_envs, 0), dtype=np.float32), + } + + def get_body_ids(self, names) -> np.ndarray: + return np.asarray([self.state_rows[name] for name in names], dtype=np.int32) + + def get_tracked_body_views(self, body_names, device=None) -> TrackedBodyStateViews: + del device + names = tuple(body_names) + rows = np.asarray([self.state_rows[name] for name in names], dtype=np.intp) + quat = np.zeros((self.num_envs, len(names), 4), dtype=np.float32) + quat[..., 0] = 1.0 + return TrackedBodyStateViews( + body_names=names, + pos_w=torch.from_numpy(self.body_pos[:, rows].copy()), + quat_w=torch.from_numpy(quat), + lin_vel_w=torch.zeros((self.num_envs, len(names), 3)), + ang_vel_w=torch.zeros((self.num_envs, len(names), 3)), + ) + + backend = Backend() + transaction = ResetStateTransaction(cast(SimBackend, backend), scene_layout=layout) + entities = {} + for name, first_body_id in (("left_hand", 1), ("right_hand", 3)): + entities[name] = Entity( + name, + EntityCfg( + root_body_name=f"{name}/wrist", + body_names=(f"{name}/wrist", f"{name}/finger"), + physical_entity=name, + ), + cast(SimBackend, backend), + reset_state=transaction, + ) + expected_rows = np.asarray( + [first_body_id - excluded_world_rows, first_body_id + 1 - excluded_world_rows], + dtype=np.int32, + ) + np.testing.assert_array_equal(entities[name]._root_body_ids, expected_rows[:1]) + np.testing.assert_array_equal(entities[name]._body_ids, expected_rows) + _, _, reset_rows = entities[name]._bind_body_randomization( + None, capability="test body reset rows" + ) + np.testing.assert_array_equal(reset_rows, expected_rows) + + copy_body_state = entities[name].bind_body_state_copy() + out_pos = np.zeros((2, 2, 3), dtype=np.float32) + out_quat = np.zeros((2, 2, 4), dtype=np.float32) + out_lin_vel = np.zeros((2, 2, 3), dtype=np.float32) + out_ang_vel = np.zeros((2, 2, 3), dtype=np.float32) + copy_body_state(out_pos, out_quat, out_lin_vel, out_ang_vel) + np.testing.assert_array_equal(out_pos[0, :, 0], [first_body_id, first_body_id + 1]) From e7b7e65ebbf49e9edb0ab6de71e60781528141a6 Mon Sep 17 00:00:00 2001 From: lemon-star608 Date: Mon, 26 Oct 2026 16:41:59 +0800 Subject: [PATCH 2/6] feat(manager): support state-feedback physics substeps --- src/unilab/base/torch_env.py | 40 ++++++++- src/unilab/envs/manager_based_rl_env.py | 25 ++++-- tests/base/test_torch_env.py | 55 ++++++++++++ tests/envs/test_manager_based_rl_env.py | 106 ++++++++++++++++++++---- 4 files changed, 201 insertions(+), 25 deletions(-) diff --git a/src/unilab/base/torch_env.py b/src/unilab/base/torch_env.py index 1a977bb4b..7052134f9 100644 --- a/src/unilab/base/torch_env.py +++ b/src/unilab/base/torch_env.py @@ -269,7 +269,7 @@ def step(self, actions: torch.Tensor) -> TorchEnvState: # previous one before the backend call so unsupported or failed steps # cannot expose stale completion metadata. self._state.info.pop("completion", None) - backend_result = self._backend.step_tensor(ctrl, self._cfg.sim_substeps) + backend_result = self._step_physics(ctrl) step_core_ms = (time.perf_counter() - phase) * 1000.0 step_core_cpu_ms = (_cpu_time() - phase_cpu) * 1000.0 if isinstance(backend_result, dict): @@ -369,6 +369,44 @@ def step(self, actions: torch.Tensor) -> TorchEnvState: self._state.info["completion"] = completion return self._state + @property + def requires_substep_state_feedback(self) -> bool: + """Whether controls must be recomputed from state between physics substeps.""" + return False + + def apply_substep_action(self) -> torch.Tensor: + """Refresh state-dependent control before the next physics substep.""" + raise NotImplementedError( + f"{type(self).__name__} must implement apply_substep_action when " + "requires_substep_state_feedback is True" + ) + + def _step_physics(self, ctrl: torch.Tensor) -> dict[str, Any] | None: + if not self.requires_substep_state_feedback: + return cast( + dict[str, Any] | None, + self._backend.step_tensor(ctrl, self._cfg.sim_substeps), + ) + + backend_result: dict[str, Any] | None = None + backend_timing: dict[str, Any] = {} + for substep in range(self._cfg.sim_substeps): + if substep > 0: + ctrl = self.apply_substep_action() + self._validate_control(ctrl) + result = self._backend.step_tensor(ctrl, 1) + if not isinstance(result, dict): + continue + backend_result = result + for key, value in result.get("timing", {}).items(): + backend_timing[key] = backend_timing.get(key, 0.0) + value + + if backend_result is None: + return None + # The final receipt orders the whole substep sequence for transition + # consumers; intermediate receipts are superseded by the next step. + return {"timing": backend_timing, "completion": backend_result.get("completion")} + def _validate_action(self, actions: torch.Tensor) -> None: self._validate_action_contract(actions) self._record_action_finite(actions) diff --git a/src/unilab/envs/manager_based_rl_env.py b/src/unilab/envs/manager_based_rl_env.py index 5253fe4d0..dfd45cf0d 100644 --- a/src/unilab/envs/manager_based_rl_env.py +++ b/src/unilab/envs/manager_based_rl_env.py @@ -1151,15 +1151,7 @@ def _validate_substep_capabilities(self) -> None: "declare a post-substep hook." ) - if self._cfg.sim_substeps > 1 and self.action_manager.requires_substep_state_feedback: - raise NotImplementedError( - "ManagerBasedRlEnv Torch lifecycle P1 does not support state-feedback " - "actions on every physics substep" - ) - def _configure_action_control(self) -> None: - # Tensor Manager stepping consumes one control tensor per control step in - # P1; multi-substep state feedback fails closed during validation. return def _control_to_backend_boundary(self) -> torch.Tensor: @@ -1222,6 +1214,19 @@ def step(self, actions: torch.Tensor) -> TorchEnvState: self.recorder_manager.record_post_step() return state + @property + def requires_substep_state_feedback(self) -> bool: + return self.action_manager.requires_substep_state_feedback + + def apply_substep_action(self) -> torch.Tensor: + self.scene._invalidate_state_reads() + read_plan = self.scene.read_plan + if read_plan is not None and self.action_manager.uses_tensor_state_packet: + read_plan.refresh() + self._sim_step_counter += 1 + self.action_manager.apply_action() + return self._control_to_backend_boundary() + def apply_action(self, actions: torch.Tensor, state: TorchEnvState) -> torch.Tensor: del state read_plan = self.scene.read_plan @@ -1232,7 +1237,9 @@ def apply_action(self, actions: torch.Tensor, state: TorchEnvState) -> torch.Ten # state once after physics. read_plan.refresh() self.action_manager.process_public_action(actions) - self._sim_step_counter += self._cfg.sim_substeps + self._sim_step_counter += ( + 1 if self.requires_substep_state_feedback else self._cfg.sim_substeps + ) self.action_manager.apply_action() return self._control_to_backend_boundary() diff --git a/tests/base/test_torch_env.py b/tests/base/test_torch_env.py index 14dae3169..6927e7d91 100644 --- a/tests/base/test_torch_env.py +++ b/tests/base/test_torch_env.py @@ -14,6 +14,8 @@ import pytest import torch from unisim.backend.base import ( + COMPLETION_CONTRACT_VERSION, + CompletionToken, SimBackend, TensorDataPlane, TensorExecution, @@ -204,6 +206,59 @@ def test_step_uses_backend_tensor_contract_and_autoreset() -> None: assert "_final_observation" not in state.info +def test_substep_feedback_preserves_control_timing_and_final_completion() -> None: + class _FeedbackEnv(_StubTorchEnv): + @property + def requires_substep_state_feedback(self) -> bool: + return True + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.feedback_control = torch.zeros((self.num_envs, 4), dtype=torch.float32) + self.substep_applies = 0 + + def apply_action(self, actions: torch.Tensor, state: TorchEnvState) -> torch.Tensor: + self.feedback_control = super().apply_action(actions, state) + return self.feedback_control + + def apply_substep_action(self) -> torch.Tensor: + self.substep_applies += 1 + self.feedback_control.add_(1.0) + return self.feedback_control + + backend = _backend() + completions = [ + CompletionToken( + contract_version=COMPLETION_CONTRACT_VERSION, + operation="step", + sequence=index, + producer_id="test", + ) + for index in range(5) + ] + controls: list[torch.Tensor] = [] + + def step_tensor(ctrl: torch.Tensor, nsteps: int) -> dict[str, Any]: + assert nsteps == 1 + controls.append(ctrl.clone()) + index = len(controls) - 1 + return {"timing": {"physics_ms": float(index + 1)}, "completion": completions[index]} + + backend.step_tensor.side_effect = step_tensor + cfg = _StubCfg(ctrl_dt=0.5) + env = _FeedbackEnv(cfg=cfg, backend=backend, terminate=False) + env.init_state() + state = env.step(_actions(env)) + + assert env.substep_applies == 4 + assert backend.step_tensor.call_count == 5 + expected = torch.stack([torch.full((3, 4), float(index)) for index in range(5)]) + torch.testing.assert_close(torch.stack(controls), expected) + assert state.info["timing"]["backend_physics_ms"] == 15.0 + assert state.info["completion"] is completions[-1] + assert env._pending_backend_step_completion is None + + def test_unsupported_backend_tensor_runtime_fails_closed() -> None: cfg = _StubCfg() env = _StubTorchEnv(cfg=cfg, backend=_unsupported_backend()) diff --git a/tests/envs/test_manager_based_rl_env.py b/tests/envs/test_manager_based_rl_env.py index 6ba330f36..b6748d53e 100644 --- a/tests/envs/test_manager_based_rl_env.py +++ b/tests/envs/test_manager_based_rl_env.py @@ -795,6 +795,21 @@ def apply_actions(self) -> None: self._entity.data.write_ctrl(self._processed) +class _FeedbackTensorBodyAction(_TensorBodyAction): + requires_substep_state_feedback = True + + def apply_actions(self) -> None: + self._env.trace.append("action_apply") + self._env.action_sim_steps.append(self._env._sim_step_counter) + super().apply_actions() + + +@dataclass(kw_only=True) +class _FeedbackTensorBodyActionCfg(_TensorBodyActionCfg): + def build(self, env) -> ActionTerm: + return _FeedbackTensorBodyAction(self, env) + + @dataclass(kw_only=True) class _CommandCfg(CommandTermCfg): def build(self, env) -> CommandTerm: @@ -3541,35 +3556,96 @@ def test_per_substep_metrics_fail_without_post_substep_backend_hook() -> None: _make_env(cfg) -def test_multisubstep_state_independent_action_skips_callback() -> None: +def test_multisubstep_state_independent_action_keeps_batched_step() -> None: env, backend = _make_env(reject_pre_step=True) try: + env.init_state() + env.step(torch.zeros((2, 1), dtype=torch.float32)) + assert backend.pre_step_control is None assert not env.action_manager.requires_substep_state_feedback + assert backend.step_nsteps == [2] + assert len(backend.tensor_controls) == 1 + assert env.action_sim_steps == [2] finally: env.close() -def test_multisubstep_state_feedback_action_uses_callback() -> None: +def test_multisubstep_state_feedback_action_steps_without_host_callback() -> None: cfg = _make_cfg( + sim_substeps=5, actions={"drive": _FeedbackDriveCfg(entity_name="robot")}, ) - with pytest.raises( - NotImplementedError, - match="Torch lifecycle P1 does not support state-feedback actions", - ): - _make_env(cfg) + env, backend = _make_env(cfg, reject_pre_step=True) + try: + env.init_state() + state = env.step(torch.zeros((2, 1), dtype=torch.float32)) + assert backend.pre_step_control is None + assert backend.step_nsteps == [1] * 5 + assert len(backend.tensor_controls) == 5 + assert env.action_sim_steps == [1, 2, 3, 4, 5] + assert state.info["completion"] is backend.completions[-1] + finally: + env.close() -def test_multisubstep_state_feedback_action_fails_when_backend_rejects_callback() -> None: - cfg = _make_cfg( - actions={"drive": _FeedbackDriveCfg(entity_name="robot")}, + +def test_multisubstep_state_feedback_refreshes_public_state_and_control( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cfg = _make_cfg(sim_substeps=5, include_optional_managers=False) + cfg.actions = { + "tilt": _FeedbackTensorBodyActionCfg( + entity_name="robot", + top_body_name="platform", + ball_body_name="ball", + ) + } + cfg.scene.entities["robot"] = EntityCfg( + root_body_name="base", + joint_names=("joint",), + body_names=("platform", "ball"), + actuator_names=("motor",), ) - with pytest.raises( - NotImplementedError, - match="Torch lifecycle P1 does not support state-feedback actions", - ): - _make_env(cfg, reject_pre_step=True) + + class _AdvancingBackend(_ScenePlanBackend): + def step(self, ctrl: np.ndarray, nsteps: int = 1) -> None: + super().step(ctrl, nsteps) + self.sensors["track_pos_w_ball"][:, 2] += 1.0 + self.trace.append("physics") + + backend = _AdvancingBackend(2) + env = _TestEnv(cfg, cast(SimBackend, backend), 2) + backend.trace = env.trace + try: + env.init_state() + env.trace.clear() + plan = env.scene.read_plan + assert plan is not None + original_refresh = plan.refresh + + def refresh() -> None: + env.trace.append("read") + original_refresh() + + monkeypatch.setattr(plan, "refresh", refresh) + state = env.step(torch.zeros((2, 1), dtype=torch.float32)) + + assert env.trace[:15] == ["read", "action_apply", "physics"] * 5 + assert backend.step_nsteps == [1] * 5 + controls = torch.stack([control[:, 0] for control in backend.tensor_controls]) + expected = torch.tensor([1.2, 2.2, 3.2, 4.2, 5.2])[:, None].expand(5, 2) + torch.testing.assert_close(controls, expected) + assert env.action_sim_steps == [1, 2, 3, 4, 5] + assert state.info["completion"] is backend.completions[-1] + + state = env.step(torch.tensor([[1.0], [0.0]], dtype=torch.float32)) + assert state.final_observation is not None + torch.testing.assert_close(state.obs["obs"][0], torch.zeros(2)) + torch.testing.assert_close(state.info["steps"], torch.tensor([0, 1])) + assert state.info["completion"] is backend.completions[-1] + finally: + env.close() def test_single_substep_feedback_action_does_not_require_callback() -> None: From cbfe0878b59a95b9868cc7e139ccb58912add8c1 Mon Sep 17 00:00:00 2001 From: lemon-star608 Date: Mon, 26 Oct 2026 16:42:28 +0800 Subject: [PATCH 3/6] feat(flashsac): add custom actor optimizer groups --- src/unilab/rl/algos/flash_sac/actor_spec.py | 109 +++++++++++++++++- src/unilab/rl/algos/flash_sac/learner.py | 35 ++++-- tests/rl/algos/test_flash_sac_custom_actor.py | 64 ++++++++++ 3 files changed, 197 insertions(+), 11 deletions(-) diff --git a/src/unilab/rl/algos/flash_sac/actor_spec.py b/src/unilab/rl/algos/flash_sac/actor_spec.py index 7768749ca..d08532bc8 100644 --- a/src/unilab/rl/algos/flash_sac/actor_spec.py +++ b/src/unilab/rl/algos/flash_sac/actor_spec.py @@ -11,6 +11,15 @@ actor_config: builder: "some_package.actors:build_actor" kwargs: {hidden_dim: 256} + optimizer_param_groups: + residual_mean: + parameters: [residual_mean.gain, residual_mean.bias] + lr_multiplier: 0.1 + +``optimizer_param_groups`` is optional. Its exact ``named_parameters()`` names +are disjoint, unknown names fail closed, and all unassigned actor parameters +form the default group. Multipliers scale each Adam group's learning rate (not +its gradients) and are preserved by scheduling and checkpoint state. The builder is a dotted import path (``pkg.module:callable`` or ``pkg.module.callable``) resolved with :func:`importlib.import_module`, so the @@ -52,6 +61,7 @@ from __future__ import annotations import json +import math import re from collections.abc import Mapping from dataclasses import dataclass, field @@ -79,12 +89,22 @@ ) +@dataclass(frozen=True) +class OptimizerParamGroupSpec: + """Exact actor parameter names sharing one Adam LR multiplier.""" + + name: str + parameters: tuple[str, ...] + lr_multiplier: float + + @dataclass(frozen=True) class FlashSACActorSpec: """Spawn-safe declarative description of a custom FlashSAC actor.""" builder: str kwargs: dict[str, Any] = field(default_factory=dict) + optimizer_param_groups: tuple[OptimizerParamGroupSpec, ...] = () @classmethod def from_config(cls, actor_config: Any) -> "FlashSACActorSpec": @@ -118,7 +138,43 @@ def from_config(cls, actor_config: Any) -> "FlashSACActorSpec": raise ValueError( f"FlashSAC actor_config.kwargs must be JSON-serializable: {exc}" ) from exc - return cls(builder=builder.strip(), kwargs=kwargs) + raw_groups = actor_config.get("optimizer_param_groups") or {} + if not isinstance(raw_groups, Mapping): + raise TypeError("actor_config.optimizer_param_groups must be a mapping") + groups: list[OptimizerParamGroupSpec] = [] + assigned: set[str] = set() + for name, group in raw_groups.items(): + if not isinstance(name, str) or not _AUX_LOSS_NAME_PATTERN.fullmatch(name): + raise ValueError(f"Invalid optimizer parameter-group name {name!r}") + if not isinstance(group, Mapping): + raise TypeError(f"Optimizer parameter group {name!r} must be a mapping") + parameters = group.get("parameters") + if ( + not isinstance(parameters, (list, tuple)) + or not parameters + or not all(isinstance(parameter, str) and parameter for parameter in parameters) + ): + raise ValueError(f"Optimizer parameter group {name!r} needs exact parameter names") + duplicates = assigned.intersection(parameters) + if duplicates or len(set(parameters)) != len(parameters): + raise ValueError( + f"Optimizer parameter names must be unique across groups: {duplicates}" + ) + assigned.update(parameters) + multiplier = group.get("lr_multiplier") + if ( + isinstance(multiplier, bool) + or not isinstance(multiplier, (int, float)) + or not math.isfinite(multiplier) + or multiplier <= 0.0 + ): + raise ValueError(f"Optimizer group {name!r} lr_multiplier must be positive") + groups.append(OptimizerParamGroupSpec(name, tuple(parameters), float(multiplier))) + return cls( + builder=builder.strip(), + kwargs=kwargs, + optimizer_param_groups=tuple(groups), + ) def _resolve_actor_builder(path: str) -> Any: @@ -281,6 +337,45 @@ def build_custom_actor( return actor +def build_actor_optimizer_param_groups( + actor: nn.Module, + spec: FlashSACActorSpec | None, +) -> list[dict[str, Any]]: + """Resolve an exhaustive, disjoint actor optimizer layout.""" + named_parameters = dict(actor.named_parameters()) + if spec is None or not spec.optimizer_param_groups: + return [ + {"params": list(named_parameters.values()), "name": "default", "lr_multiplier": 1.0} + ] + groups: list[dict[str, Any]] = [] + assigned: set[str] = set() + for group in spec.optimizer_param_groups: + missing = set(group.parameters).difference(named_parameters) + if missing: + raise ValueError( + f"Optimizer group {group.name!r} names unknown actor parameters: {sorted(missing)}" + ) + assigned.update(group.parameters) + groups.append( + { + "params": [named_parameters[name] for name in group.parameters], + "name": group.name, + "lr_multiplier": group.lr_multiplier, + } + ) + remaining = [name for name in named_parameters if name not in assigned] + if remaining: + groups.insert( + 0, + { + "params": [named_parameters[name] for name in remaining], + "name": "default", + "lr_multiplier": 1.0, + }, + ) + return groups + + def build_builtin_actor_semantics( *, obs_dim: int, @@ -320,6 +415,13 @@ def build_custom_actor_semantics( "kind": "custom", "builder": spec.builder, "kwargs": json.loads(json.dumps(spec.kwargs, sort_keys=True)), + "optimizer_param_groups": { + group.name: { + "parameters": list(group.parameters), + "lr_multiplier": group.lr_multiplier, + } + for group in spec.optimizer_param_groups + }, "obs_dim": int(obs_dim), "action_dim": int(action_dim), "action_transform": str(action_transform_id), @@ -360,7 +462,10 @@ def validate_actor_semantics(saved: Any, expected: Mapping[str, Any]) -> None: f"{saved.get('schema_version')!r} (expected {ACTOR_SEMANTICS_SCHEMA_VERSION})" ) mismatches = [] - for key in ("kind", "builder", "kwargs", "obs_dim", "action_dim", "action_transform"): + keys = ["kind", "builder", "kwargs", "obs_dim", "action_dim", "action_transform"] + if expected.get("kind") == "custom": + keys.append("optimizer_param_groups") + for key in keys: saved_value = saved.get(key) expected_value = expected.get(key) if key == "kwargs": diff --git a/src/unilab/rl/algos/flash_sac/learner.py b/src/unilab/rl/algos/flash_sac/learner.py index 8f80da23d..9e4a09447 100644 --- a/src/unilab/rl/algos/flash_sac/learner.py +++ b/src/unilab/rl/algos/flash_sac/learner.py @@ -24,6 +24,7 @@ from unilab.rl.algos.common.performance import resolve_learner_performance_tier from unilab.rl.algos.flash_sac.actor_spec import ( FlashSACActorSpec, + build_actor_optimizer_param_groups, build_builtin_actor_semantics, build_custom_actor, build_custom_actor_semantics, @@ -402,11 +403,19 @@ def __init__( optimizer_kwargs: dict[str, Any] = {"fused": _fused} if _fused and self._performance_tier.supports_capturable_optimizer: optimizer_kwargs["capturable"] = True - self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=lr_peak, **optimizer_kwargs) + actor_param_groups = build_actor_optimizer_param_groups(self.actor, self._actor_spec) + self._optimizer_lr_multipliers: dict[int, tuple[float, ...]] = {} + for group in actor_param_groups: + group["lr"] = lr_peak * group["lr_multiplier"] + self.actor_optimizer = optim.Adam(actor_param_groups, lr=lr_peak, **optimizer_kwargs) self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=lr_peak, **optimizer_kwargs) self.temperature_optimizer = optim.Adam( self.temperature.parameters(), lr=lr_peak, **optimizer_kwargs ) + for optimizer in (self.actor_optimizer, self.critic_optimizer, self.temperature_optimizer): + self._optimizer_lr_multipliers[id(optimizer)] = tuple( + float(group.get("lr_multiplier", 1.0)) for group in optimizer.param_groups + ) self._pending_actor_metric_values: torch.Tensor | None = None scheduler_fn = build_lr_lambda( @@ -433,7 +442,7 @@ def __init__( self._update_cycle_static_batch: dict[str, torch.Tensor] | None = None self._update_cycle_graph_metric_values: torch.Tensor | None = None self._update_cycle_lr_specs: list[ - tuple[optim.Optimizer, optim.lr_scheduler.LambdaLR, torch.Tensor, int] + tuple[optim.Optimizer, optim.lr_scheduler.LambdaLR, torch.Tensor, int, int] ] = [] self._update_cycle_lr_cursors: dict[int, int] = {} self._capture_update_cycle = False @@ -603,9 +612,13 @@ def _stage_update_cycle_lr(self, optimizer: optim.Optimizer) -> None: return cursor = self._update_cycle_lr_cursors.get(id(optimizer), 0) specs = [spec for spec in self._update_cycle_lr_specs if spec[0] is optimizer] - if cursor >= len(specs): + group_count = len(optimizer.param_groups) + start = cursor * group_count + staged = specs[start : start + group_count] + if len(staged) != group_count: raise RuntimeError("FlashSAC update-cycle LR staging exhausted") - optimizer.param_groups[0]["lr"] = specs[cursor][2] + for spec in staged: + optimizer.param_groups[spec[4]]["lr"] = spec[2] self._update_cycle_lr_cursors[id(optimizer)] = cursor + 1 def _maybe_normalize_obs(self, obs: torch.Tensor, *, update: bool) -> torch.Tensor: @@ -979,20 +992,24 @@ def _prepare_update_cycle_lr_specs( (self.temperature_scheduler, self.temperature_optimizer, actor_updates), ): for offset in range(count): - tensor = torch.empty((), device=self.device, dtype=torch.float32) - self._update_cycle_lr_specs.append((optimizer, scheduler, tensor, offset)) + for group_index in range(len(optimizer.param_groups)): + tensor = torch.empty((), device=self.device, dtype=torch.float32) + self._update_cycle_lr_specs.append( + (optimizer, scheduler, tensor, offset, group_index) + ) def _fill_update_cycle_lr_specs(self, *, zero: bool) -> None: - for _optimizer, scheduler, tensor, offset in self._update_cycle_lr_specs: + for optimizer, scheduler, tensor, offset, group_index in self._update_cycle_lr_specs: if zero: tensor.zero_() continue - lr = self._lr_peak * self._lr_schedule_fn(scheduler.last_epoch + offset) + multiplier = self._optimizer_lr_multipliers[id(optimizer)][group_index] + lr = multiplier * self._lr_peak * self._lr_schedule_fn(scheduler.last_epoch + offset) tensor.fill_(float(lr)) def _advance_update_cycle_schedulers(self) -> None: counts: dict[int, tuple[optim.lr_scheduler.LambdaLR, int]] = {} - for _optimizer, scheduler, _tensor, offset in self._update_cycle_lr_specs: + for _optimizer, scheduler, _tensor, offset, _group_index in self._update_cycle_lr_specs: key = id(scheduler) _scheduler, current = counts.get(key, (scheduler, 0)) counts[key] = (scheduler, max(current, offset + 1)) diff --git a/tests/rl/algos/test_flash_sac_custom_actor.py b/tests/rl/algos/test_flash_sac_custom_actor.py index fed5e6568..9ebfb2b4d 100644 --- a/tests/rl/algos/test_flash_sac_custom_actor.py +++ b/tests/rl/algos/test_flash_sac_custom_actor.py @@ -82,6 +82,7 @@ def test_builds_custom_actor_and_records_semantics(self) -> None: "kind": "custom", "builder": _BUILDER, "kwargs": {"action_scale": 2.0, "aux_weight": 0.5, "hidden_dim": 8}, + "optimizer_param_groups": {}, "obs_dim": _OBS_DIM, "action_dim": _ACTION_DIM, "action_transform": "scaled_tanh_test_v1", @@ -157,6 +158,69 @@ def test_rejects_actor_normalize_parameters_opt_out(self) -> None: ) +class TestOptimizerParamGroups: + @staticmethod + def _config() -> dict[str, Any]: + return { + "builder": _BUILDER, + "kwargs": {"hidden_dim": 8}, + "optimizer_param_groups": { + "residual_mean": { + "parameters": ["mean_head.weight", "mean_head.bias"], + "lr_multiplier": 0.1, + } + }, + } + + def test_multiplier_scales_adam_update_not_gradient(self) -> None: + learner = _make_learner(actor_config=self._config()) + optimizer = learner.actor_optimizer + default = learner.actor.trunk[0].weight + scaled = learner.actor.mean_head.weight + for parameter in learner.actor.parameters(): + parameter.grad = torch.ones_like(parameter) + default_before = default.detach().clone() + scaled_before = scaled.detach().clone() + + optimizer.step() + + default_delta = (default - default_before).abs().mean() + scaled_delta = (scaled - scaled_before).abs().mean() + torch.testing.assert_close(scaled_delta, default_delta * 0.1, rtol=1e-5, atol=1e-8) + torch.testing.assert_close(scaled.grad, torch.ones_like(scaled)) + group = next(group for group in optimizer.param_groups if group["name"] == "residual_mean") + assert group["lr"] == pytest.approx(optimizer.param_groups[0]["lr"] * 0.1) + + learner.actor_scheduler.step() + assert optimizer.param_groups[0]["lr"] > 0 + assert group["lr"] == pytest.approx(optimizer.param_groups[0]["lr"] * 0.1) + + def test_optimizer_group_checkpoint_round_trip(self) -> None: + learner = _make_learner(actor_config=self._config()) + learner.update_actor(_make_batch()) + state = learner.get_state_dict() + + restored = _make_learner(actor_config=self._config()) + restored.load_state_dict(state) + + assert [group["name"] for group in restored.actor_optimizer.param_groups] == [ + "default", + "residual_mean", + ] + assert ( + restored.get_state_dict()["actor_semantics"]["optimizer_param_groups"]["residual_mean"][ + "lr_multiplier" + ] + == 0.1 + ) + + def test_unknown_parameter_fails_closed(self) -> None: + config = self._config() + config["optimizer_param_groups"]["residual_mean"]["parameters"] = ["missing.weight"] + with pytest.raises(ValueError, match="unknown actor parameters"): + _make_learner(actor_config=config) + + class TestAuxiliaryLosses: def test_aux_loss_gradients_reach_actor_parameters(self) -> None: torch.manual_seed(7) From 28450d9e712b566c4def690de03afc4cd2f1d22a Mon Sep 17 00:00:00 2001 From: lemon-star608 Date: Mon, 26 Oct 2026 16:43:03 +0800 Subject: [PATCH 4/6] fix(flashsac): retain gradients for compiled AMP updates --- src/unilab/rl/algos/common/learner_base.py | 11 +- tests/rl/algos/test_flash_sac_amp_compile.py | 100 +++++++++++++++++++ 2 files changed, 108 insertions(+), 3 deletions(-) create mode 100644 tests/rl/algos/test_flash_sac_amp_compile.py diff --git a/src/unilab/rl/algos/common/learner_base.py b/src/unilab/rl/algos/common/learner_base.py index f466ea0f1..08bbb2122 100644 --- a/src/unilab/rl/algos/common/learner_base.py +++ b/src/unilab/rl/algos/common/learner_base.py @@ -565,10 +565,15 @@ def _materialize_capturable_optimizer_state(self) -> None: value.zero_() def _zero_optimizer_gradients(self, optimizer: optim.Optimizer) -> None: - """Zero gradients without the graph-skipped optimizer Python entrypoint.""" + """Zero gradients, retaining stable buffers required by compiled backward.""" if not self._compile_full_update_cycle: - optimizer.zero_grad(set_to_none=True) - return + if not getattr(self, "use_compile", False): + optimizer.zero_grad(set_to_none=True) + return + for group in optimizer.param_groups: + for parameter in group["params"]: + if parameter.grad is None: + parameter.grad = torch.zeros_like(parameter) gradients = [ parameter.grad for group in optimizer.param_groups diff --git a/tests/rl/algos/test_flash_sac_amp_compile.py b/tests/rl/algos/test_flash_sac_amp_compile.py new file mode 100644 index 000000000..2671a12b2 --- /dev/null +++ b/tests/rl/algos/test_flash_sac_amp_compile.py @@ -0,0 +1,100 @@ +"""CUDA compile and FP16 gradient-lifecycle regressions for FlashSAC.""" + +from __future__ import annotations + +import math +from dataclasses import replace + +import pytest +import torch + +from unilab.rl.algos.flash_sac.learner import FlashSACLearner +from unilab.rl.runtime.offpolicy.warmup import OffPolicyWarmupContext, build_offpolicy_warmup_batch + + +def _batch(batch_size: int = 4) -> dict[str, torch.Tensor]: + return { + "obs": torch.randn(batch_size, 4, device="cuda:0"), + "critic": torch.randn(batch_size, 6, device="cuda:0"), + "actions": torch.tanh(torch.randn(batch_size, 2, device="cuda:0")), + "rewards": torch.randn(batch_size, device="cuda:0"), + "next_obs": torch.randn(batch_size, 4, device="cuda:0"), + "next_critic": torch.randn(batch_size, 6, device="cuda:0"), + "dones": torch.zeros(batch_size, device="cuda:0"), + "truncated": torch.zeros(batch_size, device="cuda:0"), + } + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA compile and FP16 required") +def test_compile_fp16_prepare_and_repeated_updates() -> None: + learner = FlashSACLearner( + obs_dim=4, + action_dim=2, + critic_obs_dim=6, + device="cuda:0", + actor_hidden_dim=16, + critic_hidden_dim=16, + actor_num_blocks=1, + critic_num_blocks=1, + num_atoms=5, + use_compile=True, + compile_full_objectives=True, + use_whole_cycle_cuda_graph=False, + use_amp=True, + amp_dtype="fp16", + normalize_reward=False, + actor_config={ + "builder": "tests.rl.algos.flash_sac_actor_fixtures:build_scaled_tanh_actor", + "kwargs": {"hidden_dim": 16}, + "optimizer_param_groups": { + "mean": { + "parameters": ["mean_head.weight", "mean_head.bias"], + "lr_multiplier": 0.1, + } + }, + }, + ) + context = OffPolicyWarmupContext( + inference_observations=torch.zeros(4, 4, device="cuda:0"), + inference_dones=torch.zeros(4, device="cuda:0"), + batch_size=4, + updates_per_step=2, + policy_frequency=1, + target_frequency=1, + policy_before_critic=False, + ) + context = replace( + context, + replay_batch=build_offpolicy_warmup_batch( + context, + obs_dim=4, + critic_obs_dim=6, + action_dim=2, + device="cuda:0", + ), + ) + + learner.prepare_for_collection(context) + optimizers = (learner.actor_optimizer, learner.critic_optimizer, learner.temperature_optimizer) + gradient_buffers = { + parameter: parameter.grad + for optimizer in optimizers + for group in optimizer.param_groups + for parameter in group["params"] + } + assert all(buffer is not None for buffer in gradient_buffers.values()) + for _ in range(3): + batch = _batch() + critic_metrics = learner.update_critic(batch) + actor_metrics = learner.update_actor(batch) + assert all(math.isfinite(value) for value in critic_metrics.values()) + assert all(math.isfinite(value) for value in actor_metrics.values()) + + assert learner.actor_optimizer.state + assert learner.critic_optimizer.state + assert learner.temperature_optimizer.state + assert all(parameter.grad is buffer for parameter, buffer in gradient_buffers.items()) + groups = learner.actor_optimizer.param_groups + scaled = next(group for group in groups if group["name"] == "mean") + default = next(group for group in groups if group["name"] == "default") + assert scaled["lr"] == pytest.approx(default["lr"] * 0.1) From 7c2036a0fd0be7489e74dcca4115c451a50ab58d Mon Sep 17 00:00:00 2001 From: lemon-star608 Date: Mon, 26 Oct 2026 19:29:30 +0800 Subject: [PATCH 5/6] feat(visualization): hold completed Viser playback --- src/unilab/conf/flashsac/config.yaml | 1 + src/unilab/visualization/viser_playback.py | 31 ++++++++++++++++- tests/visualization/test_viser_playback.py | 40 ++++++++++++++++++++++ 3 files changed, 71 insertions(+), 1 deletion(-) create mode 100644 tests/visualization/test_viser_playback.py diff --git a/src/unilab/conf/flashsac/config.yaml b/src/unilab/conf/flashsac/config.yaml index 8b2bf90f0..7058eb7b4 100644 --- a/src/unilab/conf/flashsac/config.yaml +++ b/src/unilab/conf/flashsac/config.yaml @@ -176,6 +176,7 @@ viser: env_idx: 0 display_mode: all max_envs: 16 + pause_on_done: false hydra: run: diff --git a/src/unilab/visualization/viser_playback.py b/src/unilab/visualization/viser_playback.py index 5ce90eab1..b1b093240 100644 --- a/src/unilab/visualization/viser_playback.py +++ b/src/unilab/visualization/viser_playback.py @@ -154,6 +154,24 @@ def _close_scene_entries(entries: list[dict[str, Any]]) -> None: entry["scene"].close() +def _advance_session( + session: PlaybackSession, + controls: PlaybackControls, + *, + pause_on_done: bool, + completed: bool, +) -> bool: + if completed: + return True + advanced = session.advance(controls) + if not pause_on_done: + return False + if not advanced: + return False + state = session.env.state + return bool((state.terminated | state.truncated).any().item()) + + def run_viser_playback( session: PlaybackSession, *, @@ -164,6 +182,7 @@ def run_viser_playback( render_spacing: float = 1.0, start_paused: bool = False, speed: float = 1.0, + pause_on_done: bool = False, entrypoint: str = "viser playback", log: Callable[[str], None] = print, ) -> None: @@ -294,11 +313,20 @@ def _on_env_switch(event: Any) -> None: # --- Main loop ----------------------------------------------------------- try: with torch.inference_mode(): + completed = False while True: t0 = time.perf_counter() controls.set_speed(float(speed_slider.value)) - session.advance(controls) + was_completed = completed + completed = _advance_session( + session, + controls, + pause_on_done=pause_on_done, + completed=completed, + ) + if completed and not was_completed: + log("Playback completed; holding the final physics state.") physics_batch = session.physics_state() for entry in scene_entries["value"]: @@ -356,6 +384,7 @@ def run_viser_playback_from_cfg( render_spacing=render_spacing, start_paused=start_paused, speed=speed, + pause_on_done=bool(OmegaConf.select(cfg, "viser.pause_on_done", default=False)), entrypoint=entrypoint, log=log, ) diff --git a/tests/visualization/test_viser_playback.py b/tests/visualization/test_viser_playback.py new file mode 100644 index 000000000..9d5fc05a5 --- /dev/null +++ b/tests/visualization/test_viser_playback.py @@ -0,0 +1,40 @@ +from types import SimpleNamespace + +import torch + +from unilab.visualization.interactive_playback import PlaybackControls +from unilab.visualization.viser_playback import _advance_session + + +class _Session: + def __init__(self, terminated: bool, truncated: bool) -> None: + self.env = SimpleNamespace( + state=SimpleNamespace( + terminated=torch.tensor([False, terminated]), + truncated=torch.tensor([False, truncated]), + ) + ) + self.advances = 0 + + def advance(self, controls: PlaybackControls) -> bool: + del controls + self.advances += 1 + return True + + +def test_pause_on_done_latches_any_completed_environment() -> None: + session = _Session(terminated=False, truncated=True) + controls = PlaybackControls() + completed = _advance_session(session, controls, pause_on_done=True, completed=False) + assert completed + assert session.advances == 1 + assert _advance_session(session, controls, pause_on_done=True, completed=completed) + assert session.advances == 1 + + +def test_pause_on_done_disabled_preserves_playback() -> None: + session = _Session(terminated=True, truncated=False) + controls = PlaybackControls() + assert not _advance_session(session, controls, pause_on_done=False, completed=False) + assert not _advance_session(session, controls, pause_on_done=False, completed=False) + assert session.advances == 2 From b5438e074f937cd593788c539f1131d9bb2dce16 Mon Sep 17 00:00:00 2001 From: lemon-star608 Date: Mon, 26 Oct 2026 19:29:30 +0800 Subject: [PATCH 6/6] feat(logging): publish curriculum metrics --- src/unilab/rl/logging/metric_schema.py | 8 ++++++++ src/unilab/rl/logging/metrics_drain.py | 2 ++ src/unilab/rl/logging/offpolicy.py | 11 +++++++++++ src/unilab/rl/runtime/offpolicy/worker.py | 9 +++++++-- tests/rl/logging/test_metric_schema.py | 10 ++++++++-- tests/rl/logging/test_offpolicy_logger.py | 5 +++++ tests/rl/runtime/offpolicy/test_worker.py | 8 ++++++++ 7 files changed, 49 insertions(+), 4 deletions(-) diff --git a/src/unilab/rl/logging/metric_schema.py b/src/unilab/rl/logging/metric_schema.py index 7ad870b07..5e258f96e 100644 --- a/src/unilab/rl/logging/metric_schema.py +++ b/src/unilab/rl/logging/metric_schema.py @@ -543,6 +543,14 @@ def _collector_ms(tag: str, description: str) -> MetricSpec: ), distributed_aggregation="cross-rank mean of per-rank report means", ), + MetricFamilySpec( + prefix="Curriculum/", + owner="environment/curriculum", + unit="curriculum state", + aggregation="latest reset-published state", + step_axis=_STEP_AXIS, + description="Latest curriculum values published through env info['log'].", + ), MetricFamilySpec( prefix="Loss/aux_", owner="learner", diff --git a/src/unilab/rl/logging/metrics_drain.py b/src/unilab/rl/logging/metrics_drain.py index 2d5315978..129677786 100644 --- a/src/unilab/rl/logging/metrics_drain.py +++ b/src/unilab/rl/logging/metrics_drain.py @@ -81,6 +81,8 @@ def drain_collector_metrics( reward_history.append(metrics["return_mean_ep100"]) if "reward_components" in metrics: reward_components.update(metrics["reward_components"]) + if "curriculum_metrics" in metrics: + logger.update_curriculum_metrics(metrics["curriculum_metrics"]) if "mean_episode_length" in metrics: logger.update_mean_episode_length(metrics["mean_episode_length"]) if "collector_timing_ms" in metrics: diff --git a/src/unilab/rl/logging/offpolicy.py b/src/unilab/rl/logging/offpolicy.py index 037364dce..f714baa79 100644 --- a/src/unilab/rl/logging/offpolicy.py +++ b/src/unilab/rl/logging/offpolicy.py @@ -297,6 +297,7 @@ def __init__( self._has_iteration_extra_info: bool = False self._collector_timing: dict[str, float] = {} self._collector_inference: dict[str, float] = {} + self._curriculum_metrics: dict[str, float] = {} if timing_profile not in _LEARNER_TIMING_PROFILES: raise ValueError("timing_profile must be 'sac_family' or 'appo'") self._timing_profile = timing_profile @@ -533,6 +534,14 @@ def _record_terminal_sample( batch_size_per_rank=self._batch_size_per_rank, ) + def update_curriculum_metrics(self, metrics: dict[str, float]) -> None: + for key, value in metrics.items(): + if not key.startswith("Curriculum/"): + raise ValueError(f"invalid curriculum metric key: {key}") + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"curriculum metric {key} must be numeric") + self._curriculum_metrics.update(metrics) + def update_collector_timing(self, timing_ms: dict[str, float]): normalized = dict(timing_ms) unknown = sorted(set(normalized) - set(_COLLECTOR_TIMING_SPECS)) @@ -719,6 +728,8 @@ def _build_backend_scalars( _set_backend_scalar(scalars, "Train/mean_episode_length", self._mean_ep_length) if self._timeout_rate is not None: _set_backend_scalar(scalars, "Episode/timeout_rate", self._timeout_rate) + for key, value in self._curriculum_metrics.items(): + _set_backend_scalar(scalars, key, value) for key, value in self._collector_timing.items(): if key == "rollout_ms": _set_backend_scalar(scalars, "Perf/collection_time", value / 1000.0) diff --git a/src/unilab/rl/runtime/offpolicy/worker.py b/src/unilab/rl/runtime/offpolicy/worker.py index addd80f97..fe94ca937 100644 --- a/src/unilab/rl/runtime/offpolicy/worker.py +++ b/src/unilab/rl/runtime/offpolicy/worker.py @@ -475,6 +475,7 @@ def _run_collector( max_publication_lag_since_metric = 0 final_metrics_flushed = False pending_metric_flush: MetricFlush | None = None + latest_curriculum_metrics: dict[str, float] = {} def enqueue_collector_metrics( *, @@ -495,6 +496,8 @@ def enqueue_collector_metrics( "total_steps": total_steps, "buffer_size": int(replay_buffer.size[0]), } + if latest_curriculum_metrics: + msg["curriculum_metrics"] = dict(latest_curriculum_metrics) if final: msg["metric_flush"] = "final" rewards_for_report = reward_history if reward_history is not None else ep_rewards @@ -870,12 +873,14 @@ def flush_final_collector_metrics() -> None: critic_t = next_critic_t total_steps += num_envs - # Extract reward components from env info - log_info = state.info.get("log", {}) if reward_component_names else {} + # Extract reward components and latest curriculum state from env info. + log_info = state.info.get("log", {}) if log_info: for k, v in log_info.items(): if k.startswith("reward/"): ep_reward_components[k.removeprefix("reward/")].append(v) + elif k.startswith("Curriculum/") and isinstance(v, (int, float)): + latest_curriculum_metrics[k] = float(v) phase_start_ns = _record_phase_ms(cycle_timing_ms, "metrics_publish_ms", phase_start_ns) diff --git a/tests/rl/logging/test_metric_schema.py b/tests/rl/logging/test_metric_schema.py index afa1e4939..248b5bba7 100644 --- a/tests/rl/logging/test_metric_schema.py +++ b/tests/rl/logging/test_metric_schema.py @@ -87,10 +87,16 @@ def test_replay_ingress_diagnostics_are_registered_as_cumulative_pipeline_metric def test_schema_is_closed_except_reward_terms() -> None: validate_metric_tags( - ["Loss/surrogate", "reward/term", "Train/rollouts_read", "Perf/learner_inference_ms"] + [ + "Loss/surrogate", + "reward/term", + "Curriculum/voc/stage", + "Train/rollouts_read", + "Perf/learner_inference_ms", + ] ) - assert set(METRIC_FAMILIES) == {"reward/", "Loss/aux_"} + assert set(METRIC_FAMILIES) == {"reward/", "Curriculum/", "Loss/aux_"} assert "Perf/learner_other_ms" not in METRIC_SPECS assert "Train/staging_pool_capacity" not in METRIC_SPECS assert "Train/updates" not in METRIC_SPECS diff --git a/tests/rl/logging/test_offpolicy_logger.py b/tests/rl/logging/test_offpolicy_logger.py index be3ab07a8..ec1e5c531 100644 --- a/tests/rl/logging/test_offpolicy_logger.py +++ b/tests/rl/logging/test_offpolicy_logger.py @@ -725,6 +725,9 @@ def test_offpolicy_backend_log_step_writes_one_batched_event() -> None: writer = _BatchedWriter() logger._tb_writer = writer + logger.update_curriculum_metrics( + {"Curriculum/voc/stage": 2.0, "Curriculum/voc/target_voc_scale": 0.5} + ) logger.log_step( iteration=1, metrics={"Loss/critic": 2.0}, @@ -743,6 +746,8 @@ def test_offpolicy_backend_log_step_writes_one_batched_event() -> None: assert scalars["Loss/critic"] == pytest.approx(2.0) assert scalars["Train/mean_reward"] == pytest.approx(1.25) assert scalars["reward/alive"] == pytest.approx(1.0) + assert scalars["Curriculum/voc/stage"] == pytest.approx(2.0) + assert scalars["Curriculum/voc/target_voc_scale"] == pytest.approx(0.5) assert scalars["Perf/total_fps"] == pytest.approx(400.0) assert "Perf/iteration_time" in scalars assert "Perf/learning_time" in scalars diff --git a/tests/rl/runtime/offpolicy/test_worker.py b/tests/rl/runtime/offpolicy/test_worker.py index 6820bb0fe..6ef6026d3 100644 --- a/tests/rl/runtime/offpolicy/test_worker.py +++ b/tests/rl/runtime/offpolicy/test_worker.py @@ -1176,6 +1176,10 @@ def step(self, actions): "reward_components": torch.tensor( [[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32 ), + "log": { + "Curriculum/voc/stage": 2.0, + "Curriculum/voc/target_voc_scale": 0.5, + }, }, ) return self.state @@ -1246,6 +1250,10 @@ def serve_learner_action(coordination_queue, tick_id, stop, **kwargs): report = metrics_queue.get_nowait() assert report["reward_components"] == {"left": 2.0, "right": 3.0} assert report["reward_component_means"] is True + assert report["curriculum_metrics"] == { + "Curriculum/voc/stage": 2.0, + "Curriculum/voc/target_voc_scale": 0.5, + } def test_worker_shutdown_flushes_partial_metrics_before_abnormal_cleanup(