refactor(a2arm): complete manager-based migration cleanup

Author: lemon-star608

conf/ppo_cse/task/a2arm_pos_force/base.yaml

@@ -21,13 +21,19 @@ env:
         history:
           _target_: unilab.tasks.locomotion.a2arm.observations.A2ArmActorHistoryCfg
           func: unilab.tasks.locomotion.a2arm.observations.A2ArmActorHistory
-          params: {clip: 100.0, actor_noise: true}
+          params:
+            clip: 100.0
+            actor_noise: true
+            history_length: "${oc.select:algo.num_actor_history,32}"
     critic:
       terms:
         history:
           _target_: unilab.tasks.locomotion.a2arm.observations.A2ArmCriticHistoryCfg
           func: unilab.tasks.locomotion.a2arm.observations.A2ArmCriticHistory
-          params: {clip: 100.0, actor_noise: false}
+          params:
+            clip: 100.0
+            actor_noise: false
+            history_length: "${oc.select:algo.num_critic_history,3}"
   actions:
     joint_pd:
       _target_: unilab.tasks.locomotion.a2arm.actions.A2ArmPdActionCfg
@@ -85,11 +91,11 @@ env:
   critic_observation_group: critic
 
 reward:
-  tracking_lin_vel_force_world: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 2.0, params: {name: tracking_lin_vel_force_world}}
-  tracking_ee_force_world: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 2.0, params: {name: tracking_ee_force_world}}
-  tracking_ang_vel: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 1.0, params: {name: tracking_ang_vel}}
+  tracking_lin_vel_force_world: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 2.0, params: {name: tracking_lin_vel_force_world, sigma: 0.25}}
+  tracking_ee_force_world: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 2.0, params: {name: tracking_ee_force_world, sigma: 1.0}}
+  tracking_ang_vel: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 1.0, params: {name: tracking_ang_vel, sigma: 0.25}}
   orientation: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -1.5, params: {name: orientation}}
-  ref_dof_leg: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 3.0, params: {name: ref_dof_leg}}
+  ref_dof_leg: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 3.0, params: {name: ref_dof_leg, scale: 0.1}}
   alive: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 1.5, params: {name: alive}}
   lin_vel_z: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -1.5, params: {name: lin_vel_z}}
   ang_vel_xy: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -0.1, params: {name: ang_vel_xy}}
@@ -100,16 +106,16 @@ reward:
   dof_vel_arm: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -2.0e-4, params: {name: dof_vel_arm}}
   dof_acc: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -2.5e-7, params: {name: dof_acc}}
   dof_acc_arm: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -4.5e-7, params: {name: dof_acc_arm}}
-  base_height: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -2.0, params: {name: base_height}}
+  base_height: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -2.0, params: {name: base_height, target: 0.435}}
   hip_pos: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -0.5, params: {name: hip_pos}}
-  torque_limits: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -0.005, params: {name: torque_limits}}
+  torque_limits: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -0.005, params: {name: torque_limits, soft_limit: 0.9}}
   dof_pos_limits: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -10.0, params: {name: dof_pos_limits}}
-  stand_still: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 0.5, params: {name: stand_still}}
+  stand_still: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 0.5, params: {name: stand_still, scale: 0.05}}
   collision: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -5.0, params: {name: collision}}
   feet_contact_number: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 2.0, params: {name: feet_contact_number}}
-  feet_air_time: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 1.0, params: {name: feet_air_time}}
-  feet_height: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 1.0, params: {name: feet_height}}
-  feet_height_high: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -15.0, params: {name: feet_height_high}}
+  feet_air_time: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 1.0, params: {name: feet_air_time, threshold: 0.5}}
+  feet_height: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: 1.0, params: {name: feet_height, target: 0.12}}
+  feet_height_high: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -15.0, params: {name: feet_height_high, target: 0.24}}
   feet_pos_xy: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -3.0, params: {name: feet_pos_xy}}
   feet_drag: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -8.0e-4, params: {name: feet_drag}}
-  feet_contact_forces: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -1.0e-3, params: {name: feet_contact_forces}}
+  feet_contact_forces: {func: unilab.tasks.locomotion.a2arm.rewards.a2arm_reward, weight: -1.0e-3, params: {name: feet_contact_forces, threshold: 200.0}}

scripts/benchmark/rl/benchmark_offpolicy_collector_active.py

@@ -33,12 +33,13 @@
 import sys
 import time
 from collections import defaultdict
+from collections.abc import Sequence
 from dataclasses import asdict, dataclass, field
 from datetime import datetime, timezone
 from importlib.util import find_spec
 from pathlib import Path
 from statistics import mean, median, pstdev
-from typing import Any, Sequence, cast
+from typing import Any, cast
 
 import numpy as np
 import torch
@@ -1329,89 +1330,53 @@ def _format_dr_reset_timing_table(results: list[CollectorResult]) -> str:
 )
 
 
-def _format_set_state_detail_table(results: list[CollectorResult]) -> str:
-    """Backend set_state sub-timing table (motrix keyset).
-
-    Renders the 14 motrix-oriented sub-keys next to the outer
-    ``dr_reset_set_state_ms``. Backends that don't populate a key emit 0.0 so
-    columns stay stable across backends. MuJoCo runs will show 0.0 for the
-    motrix-only sub-keys; use :func:`_format_set_state_mujoco_table` for the
-    MuJoCo-oriented view instead.
-    """
+def _format_set_state_backend_table(
+    results: list[CollectorResult],
+    keys: Sequence[tuple[str, str]],
+) -> str:
     headers = (
         "Algo",
         "Task",
         "Backend",
         "Set state ms (%env, %active)",
-        *(label for _, label in _SET_STATE_MOTRIX_KEYS),
+        *(label for _, label in keys),
     )
     rows = []
     for result in results:
-        env_step = result.phase_ms_per_vector_step.get("env_step_ms")
-        if env_step is None:
+        if result.phase_ms_per_vector_step.get("env_step_ms") is None:
             continue
         rows.append(
             (
                 result.case.algo,
                 result.case.task,
                 result.case.runtime_sim_backend,
                 _format_np_env_timing(result, "dr_reset_set_state_ms"),
-                *(_format_set_state_sub_ms(result, key) for key, _ in _SET_STATE_MOTRIX_KEYS),
+                *(_format_set_state_sub_ms(result, key) for key, _ in keys),
             )
         )
     return _format_table(headers, rows)
 
 
+def _format_set_state_detail_table(results: list[CollectorResult]) -> str:
+    """Backend set_state sub-timing table (motrix keyset).
+
+    Renders the 14 motrix-oriented sub-keys next to the outer
+    ``dr_reset_set_state_ms``. Backends that don't populate a key emit 0.0 so
+    columns stay stable across backends. MuJoCo runs will show 0.0 for the
+    motrix-only sub-keys; use :func:`_format_set_state_mujoco_table` for the
+    MuJoCo-oriented view instead.
+    """
+    return _format_set_state_backend_table(results, _SET_STATE_MOTRIX_KEYS)
+
+
 def _format_set_state_mujoco_table(results: list[CollectorResult]) -> str:
     """Backend set_state sub-timing table (mujoco keyset)."""
-    headers = (
-        "Algo",
-        "Task",
-        "Backend",
-        "Set state ms (%env, %active)",
-        *(label for _, label in _SET_STATE_MUJOCO_KEYS),
-    )
-    rows = []
-    for result in results:
-        env_step = result.phase_ms_per_vector_step.get("env_step_ms")
-        if env_step is None:
-            continue
-        rows.append(
-            (
-                result.case.algo,
-                result.case.task,
-                result.case.runtime_sim_backend,
-                _format_np_env_timing(result, "dr_reset_set_state_ms"),
-                *(_format_set_state_sub_ms(result, key) for key, _ in _SET_STATE_MUJOCO_KEYS),
-            )
-        )
-    return _format_table(headers, rows)
+    return _format_set_state_backend_table(results, _SET_STATE_MUJOCO_KEYS)
 
 
 def _format_set_state_mjwarp_table(results: list[CollectorResult]) -> str:
     """Backend set_state sub-timing table (mjwarp keyset)."""
-    headers = (
-        "Algo",
-        "Task",
-        "Backend",
-        "Set state ms (%env, %active)",
-        *(label for _, label in _SET_STATE_MJWARP_KEYS),
-    )
-    rows = []
-    for result in results:
-        env_step = result.phase_ms_per_vector_step.get("env_step_ms")
-        if env_step is None:
-            continue
-        rows.append(
-            (
-                result.case.algo,
-                result.case.task,
-                result.case.runtime_sim_backend,
-                _format_np_env_timing(result, "dr_reset_set_state_ms"),
-                *(_format_set_state_sub_ms(result, key) for key, _ in _SET_STATE_MJWARP_KEYS),
-            )
-        )
-    return _format_table(headers, rows)
+    return _format_set_state_backend_table(results, _SET_STATE_MJWARP_KEYS)
 
 
 def _format_np_env_step_timing_table(results: list[CollectorResult]) -> str:

scripts/play_a2arm_pos_force_interactive.py

@@ -25,10 +25,8 @@
 from unilab.base.backend.mujoco.playback import resolve_render_play_model_files
 from unilab.base.config_adapter import BackendAdapter, create_env
 from unilab.tasks.locomotion.a2arm.state import A2ArmPosForceState
-from unilab.training import ensure_registries, parse_checkpoint_path
+from unilab.training import algo_config_dict, ensure_registries, parse_checkpoint_path
 from unilab.visualization.a2arm_pos_force import (
-    KEY_BACKSPACE,
-    KEY_SPACE,
     clear_teleop_override,
     draw_markers,
     install_teleop_override,
@@ -49,13 +47,6 @@ def _select_device(cfg: DictConfig) -> str:
     return "cpu"
 
 
-def _algo_config_dict(cfg: DictConfig) -> dict[str, Any]:
-    resolved = OmegaConf.to_container(cfg.algo, resolve=True)
-    if not isinstance(resolved, dict):
-        raise TypeError("cfg.algo must resolve to a mapping")
-    return cast(dict[str, Any], resolved)
-
-
 def _backend_adapter(cfg: DictConfig) -> BackendAdapter:
     return BackendAdapter(
         cfg,
@@ -131,10 +122,9 @@ def play_interactive(cfg: DictConfig, device: str) -> None:
         sim_backend="mujoco",
     )
     wrapped = RslRlVecEnvWrapper(env, device=device)
-    runner = CSEOnPolicyRunner(wrapped, _algo_config_dict(cfg), log_dir=None, device=device)
+    runner = CSEOnPolicyRunner(wrapped, algo_config_dict(cfg), log_dir=None, device=device)
     runner.load(str(checkpoint))
     policy = runner.get_inference_policy(device=device)
-    env.set_autoreset(False)
 
     state: A2ArmPosForceState = env.command_manager.get_term("task_state")
     teleop = make_teleop_from_state(state)

scripts/train_cse_ppo.py

@@ -46,11 +46,6 @@
 EXPORT_POLICY = False
 
 
-def _algo_config_dict(cfg: DictConfig) -> dict[str, Any]:
-    """Compatibility helper matching the historical CSE entrypoint surface."""
-    return algo_config_dict(cfg)
-
-
 def _backend_adapter(cfg: DictConfig) -> BackendAdapter:
     return BackendAdapter(
         cfg,
@@ -64,30 +59,14 @@ def _get_log_root(cfg: DictConfig) -> str:
     return str(get_log_root(ROOT_DIR, cfg))
 
 
-def _format_play_checkpoint_error(
-    cfg: DictConfig,
-    *,
-    task_log_root: Path,
-    load_path: Path | None,
-    load_path_dir: Path | None,
-) -> str:
-    """Historical CSE helper retained while delegating to shared diagnostics."""
-    return format_play_checkpoint_error(
-        cfg,
-        task_log_root=task_log_root,
-        load_path=load_path,
-        load_path_dir=load_path_dir,
-    )
-
-
 def play_cse_ppo(cfg: DictConfig, device: str) -> str | None:
     """Resolve a checkpoint and render a CSE policy using the shared playback session."""
-    rl_cfg = _algo_config_dict(cfg)
+    rl_cfg = algo_config_dict(cfg)
     task_log_root = get_log_root(ROOT_DIR, cfg) / str(cfg.training.task_name)
     load_path, load_path_dir = parse_checkpoint_path(cfg, root_dir=ROOT_DIR)
     if load_path is None or load_path_dir is None or not load_path.exists():
         print(
-            _format_play_checkpoint_error(
+            format_play_checkpoint_error(
                 cfg, task_log_root=task_log_root, load_path=load_path, load_path_dir=load_path_dir
             )
         )
@@ -161,9 +140,6 @@ def env_factory(num_envs: int):
                 "cam_tracking_env_idx": getattr(cfg.training, "cam_tracking_env_idx", 0),
                 "cam_tracking_extra_envs": getattr(cfg.training, "cam_tracking_extra_envs", 2),
             },
-            extra_data_getter=(lambda: getattr(env, "curr_ee_goal_world", None))
-            if hasattr(env, "curr_ee_goal_world")
-            else None,
         )
     return str(output_video)
 
@@ -209,7 +185,7 @@ def main(cfg: DictConfig) -> None:
             apply_env_nan_guard(env, cfg.training)
             runner = CSEOnPolicyRunner(
                 RslRlVecEnvWrapper(env, device=device),
-                _algo_config_dict(cfg),
+                algo_config_dict(cfg),
                 log_dir=log_dir,
                 device=device,
             )

src/unilab/algos/cse_ppo/__init__.py

@@ -6,4 +6,4 @@
 from .runner import CSEOnPolicyRunner
 from .storage import CSERolloutStorage
 
-__all__ = ["CSEActorCritic", "CSEPPO", "CSEEstimator", "CSEOnPolicyRunner", "CSERolloutStorage"]
+__all__ = ["CSEPPO", "CSEActorCritic", "CSEEstimator", "CSEOnPolicyRunner", "CSERolloutStorage"]

src/unilab/algos/cse_ppo/actor_critic.py

@@ -3,13 +3,13 @@
 
 from __future__ import annotations
 
-from typing import Sequence
+from collections.abc import Sequence
 
 import torch
 from torch import nn
 from torch.distributions import Normal
 
-from .estimator import CSEEstimator, get_activation
+from .estimator import CSEEstimator, _mlp
 
 
 class CSEActorCritic(nn.Module):
@@ -48,13 +48,13 @@ def __init__(
             activation=activation,
             **dict(estimator or {}),
         )
-        self.actor = _build_mlp(
+        self.actor = _mlp(
             self.num_one_step_obs + self.estimator.num_latent,
             self.num_actions,
             actor_hidden_dims,
             activation,
         )
-        self.critic = _build_mlp(self.num_critic_obs, 1, critic_hidden_dims, activation)
+        self.critic = _mlp(self.num_critic_obs, 1, critic_hidden_dims, activation)
         self.std = nn.Parameter(float(init_noise_std) * torch.ones(self.num_actions))
         self.distribution: Normal | None = None
         Normal.set_default_validate_args(False)
@@ -105,19 +105,3 @@ def act_inference(self, obs_history: torch.Tensor, observations=None) -> torch.T
     def evaluate(self, critic_observations: torch.Tensor, **kwargs) -> torch.Tensor:
         del kwargs
         return self.critic(critic_observations)
-
-
-def _build_mlp(
-    input_dim: int,
-    output_dim: int,
-    hidden_dims: Sequence[int],
-    activation: str,
-) -> nn.Sequential:
-    """Build a plain feed-forward MLP (kept as a module-level compatibility helper)."""
-    layers: list[nn.Module] = []
-    last = int(input_dim)
-    for dim in hidden_dims:
-        layers.extend((nn.Linear(last, int(dim)), get_activation(activation)))
-        last = int(dim)
-    layers.append(nn.Linear(last, int(output_dim)))
-    return nn.Sequential(*layers)

src/unilab/algos/cse_ppo/algorithm.py

@@ -226,7 +226,7 @@ def update(self) -> tuple[float, float, float]:
 def _critic_obs(obs: TensorDict | torch.Tensor) -> torch.Tensor:
     if isinstance(obs, TensorDict):
         for key in ("critic", "policy", "actor"):
-            if key in obs.keys():
+            if key in obs:
                 return obs[key]
         raise KeyError("CSE-PPO TensorDict obs must contain critic, policy, or actor")
     return obs

src/unilab/algos/cse_ppo/estimator.py

@@ -121,7 +121,7 @@ def _regression_loss(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Te
             return (weights * F.mse_loss(pred, target, reduction="none")).mean()
         loss = pred.new_zeros(())
         offset = 0
-        for size, weight in zip(self.target_group_sizes, weights):
+        for size, weight in zip(self.target_group_sizes, weights, strict=True):
             part = slice(offset, offset + size)
             loss = loss + F.mse_loss(pred[:, part] * weight, target[:, part] * weight)
             offset += size