feat: 补全 Manager-Based 与 FlashSAC 通用训练和评估能力

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,

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)

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:

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()
 

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

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":