fix: align timeout bootstrap with final observation contract

conf/offpolicy/algo/flashsac.yaml

@@ -0,0 +1,43 @@
+# @package algo
+algo: flashsac
+algo_log_name: flash_sac
+load_run: "-1"
+seed: 1
+num_envs: 2048
+batch_size: 2048
+replay_buffer_n: 512
+updates_per_step: 1
+warmup_steps: 10000
+policy_frequency: 2
+env_steps_per_sync: 1
+max_iterations: 5000
+save_interval: 1000
+gamma: 0.99
+tau: 0.01
+actor_lr: 3.0e-4
+critic_lr: 3.0e-4
+actor_hidden_dim: 128
+critic_hidden_dim: 256
+num_atoms: 101
+obs_normalization: false
+use_layer_norm: false
+algo_params:
+  normalize_reward: true
+  normalized_g_max: 5.0
+  actor_num_blocks: 2
+  critic_num_blocks: 2
+  actor_bc_alpha: 0.0
+  actor_noise_zeta_mu: 2.0
+  actor_noise_zeta_max: 16
+  critic_min_v: -5.0
+  critic_max_v: 5.0
+  temp_initial_value: 0.01
+  temp_target_sigma: 0.15
+  temp_target_entropy: null
+  learning_rate_init: 3.0e-4
+  learning_rate_peak: 3.0e-4
+  learning_rate_end: 1.5e-4
+  learning_rate_warmup_steps: 0
+  learning_rate_decay_steps: 500000
+  n_step: 1
+  use_compile: false

conf/offpolicy/task/flashsac/g1_sac/motrix.yaml

@@ -0,0 +1,28 @@
+# @package _global_
+training:
+  task_name: G1WalkTaskMjSAC
+  sim_backend: motrix
+algo:
+  num_envs: 2048
+  max_iterations: 5000
+  save_interval: 1000
+reward:
+  scales:
+    tracking_lin_vel: 2.2
+    tracking_ang_vel: 1.8
+    penalty_ang_vel_xy: -1.2
+    penalty_orientation: -12.0
+    penalty_action_rate: -2.5
+    pose: -0.6
+    penalty_feet_ori: -30.0
+    feet_phase: 6.0
+    alive: 12.0
+  tracking_sigma: 0.25
+  base_height_target: 0.754
+  min_base_height: 0.3
+  max_tilt_deg: 65.0
+  gait_frequency: 1.5
+  feet_phase_swing_height: 0.09
+  feet_phase_tracking_sigma: 0.008
+  close_feet_threshold: 0.15
+  pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0]

conf/offpolicy/task/flashsac/g1_sac/mujoco.yaml

@@ -0,0 +1,28 @@
+# @package _global_
+training:
+  task_name: G1WalkTaskMjSAC
+  sim_backend: mujoco
+algo:
+  num_envs: 2048
+  max_iterations: 5000
+  save_interval: 1000
+reward:
+  scales:
+    tracking_lin_vel: 2.0
+    tracking_ang_vel: 1.5
+    penalty_ang_vel_xy: -1.0
+    penalty_orientation: -10.0
+    penalty_action_rate: -2.0
+    pose: -0.5
+    penalty_feet_ori: -25.0
+    feet_phase: 5.0
+    alive: 10.0
+  tracking_sigma: 0.25
+  base_height_target: 0.754
+  min_base_height: 0.3
+  max_tilt_deg: 65.0
+  gait_frequency: 1.5
+  feet_phase_swing_height: 0.09
+  feet_phase_tracking_sigma: 0.008
+  close_feet_threshold: 0.15
+  pose_weights: [0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 0.01, 1.0, 5.0, 0.01, 5.0, 5.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0, 50.0]

conf/ppo/config.yaml

@@ -28,7 +28,7 @@ algo:
     activation: elu
     class_name: ActorCritic
   algorithm:
-    class_name: PPO
+    class_name: unilab.algos.torch.rsl_rl_ppo:FinalObservationAwarePPO
     value_loss_coef: 1.0
     use_clipped_value_loss: true
     clip_param: 0.2

scripts/train_mlx_ppo.py

@@ -128,6 +128,28 @@ def build_model(cfg, obs_dim: int, action_dim: int, dtype=mx.float32) -> MLPActo
     )
 
 
+def get_time_limit_bootstrap_values(
+    state: Any, model: MLPActorCritic, model_dtype=mx.float32
+) -> mx.array | None:
+    """Return V(final_observation) for current timeout envs when available."""
+    if not hasattr(state, "truncated"):
+        return None
+    timeout_mask = np.asarray(state.truncated, dtype=bool)
+    if not np.any(timeout_mask):
+        return None
+    final_observation = getattr(state, "final_observation", None)
+    if final_observation is None:
+        info = getattr(state, "info", None)
+        if isinstance(info, dict):
+            final_observation = info.get("final_observation")
+    if not isinstance(final_observation, dict):
+        return None
+    final_obs = mx.array(flatten_obs_dict(final_observation))
+    if getattr(final_obs, "dtype", None) != model_dtype:
+        final_obs = final_obs.astype(model_dtype)
+    return model.value(final_obs)
+
+
 def _get_log_root(cfg: DictConfig) -> Path:
     return cast(Path, get_log_root(ROOT_DIR, cfg))
 
@@ -464,7 +486,15 @@ def log(msg: str) -> None:
             next_obs = mx.nan_to_num(raw_obs, nan=0.0, posinf=0.0, neginf=0.0)
             if hasattr(state, "truncated"):
                 timeouts = mx.array(state.truncated, dtype=dtype)
-                rewards = rewards + ppo_cfg.gamma * values_mx.astype(rewards.dtype) * timeouts
+                timeout_bootstrap_values = get_time_limit_bootstrap_values(
+                    state, model, model_dtype
+                )
+                if timeout_bootstrap_values is None:
+                    timeout_bootstrap_values = values_mx
+                rewards = (
+                    rewards
+                    + ppo_cfg.gamma * timeout_bootstrap_values.astype(rewards.dtype) * timeouts
+                )
             if rewards.dtype != dtype:
                 rewards = rewards.astype(dtype)
 

scripts/train_offpolicy.py

@@ -86,6 +86,9 @@ def build_runner(algo_name: str, cfg: DictConfig):
     """Build algorithm runner from unified Hydra config."""
     env_cfg_override = build_offpolicy_env_cfg_override(algo_name, cfg)
 
+    if algo_name == "flashsac" and cfg.training.num_gpus > 1:
+        raise ValueError("FlashSAC does not support training.num_gpus > 1")
+
     if algo_name == "sac":
         from unilab.algos.torch.fast_sac.learner import FastSACLearner
         from unilab.algos.torch.fast_sac.runner import FastSACRunner
@@ -218,6 +221,52 @@ def build_runner(algo_name: str, cfg: DictConfig):
             sim_backend=cfg.training.sim_backend,
         )
 
+    if algo_name == "flashsac":
+        from unilab.algos.torch.flash_sac.runner import FlashSACRunner
+
+        return FlashSACRunner(
+            env_name=cfg.training.task_name,
+            env_cfg_override=env_cfg_override,
+            device=cfg.training.device,
+            num_envs=cfg.algo.num_envs,
+            replay_buffer_n=cfg.algo.replay_buffer_n,
+            batch_size=cfg.algo.batch_size,
+            warmup_steps=cfg.algo.warmup_steps,
+            updates_per_step=cfg.algo.updates_per_step,
+            policy_frequency=cfg.algo.policy_frequency,
+            gamma=cfg.algo.gamma,
+            tau=cfg.algo.tau,
+            actor_lr=cfg.algo.actor_lr,
+            critic_lr=cfg.algo.critic_lr,
+            obs_normalization=cfg.algo.obs_normalization,
+            actor_hidden_dim=cfg.algo.actor_hidden_dim,
+            critic_hidden_dim=cfg.algo.critic_hidden_dim,
+            num_atoms=cfg.algo.num_atoms,
+            use_amp=cfg.training.use_amp,
+            sync_collection=not cfg.training.no_sync_collection,
+            env_steps_per_sync=cfg.training.env_steps_per_sync,
+            sim_backend=cfg.training.sim_backend,
+            actor_num_blocks=cfg.algo.algo_params.actor_num_blocks,
+            critic_num_blocks=cfg.algo.algo_params.critic_num_blocks,
+            actor_bc_alpha=cfg.algo.algo_params.actor_bc_alpha,
+            actor_noise_zeta_mu=cfg.algo.algo_params.actor_noise_zeta_mu,
+            actor_noise_zeta_max=cfg.algo.algo_params.actor_noise_zeta_max,
+            critic_min_v=cfg.algo.algo_params.critic_min_v,
+            critic_max_v=cfg.algo.algo_params.critic_max_v,
+            target_sigma=cfg.algo.algo_params.temp_target_sigma,
+            target_entropy=cfg.algo.algo_params.temp_target_entropy,
+            temp_initial_value=cfg.algo.algo_params.temp_initial_value,
+            learning_rate_init=cfg.algo.algo_params.learning_rate_init,
+            learning_rate_peak=cfg.algo.algo_params.learning_rate_peak,
+            learning_rate_end=cfg.algo.algo_params.learning_rate_end,
+            learning_rate_warmup_steps=cfg.algo.algo_params.learning_rate_warmup_steps,
+            learning_rate_decay_steps=cfg.algo.algo_params.learning_rate_decay_steps,
+            normalize_reward=cfg.algo.algo_params.normalize_reward,
+            normalized_g_max=cfg.algo.algo_params.normalized_g_max,
+            n_step=cfg.algo.algo_params.n_step,
+            use_compile=cfg.algo.algo_params.use_compile,
+        )
+
     raise ValueError(f"Unsupported algo: {algo_name}")
 
 
@@ -273,6 +322,22 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None:
             torch.device(device),
         )
         if cfg.algo.obs_normalization:
+            normalizer = EmpiricalNormalization(shape=obs_dim, device=device)
+    elif algo_name == "flashsac":
+        actor = build_actor(
+            "flashsac",
+            obs_dim,
+            action_dim,
+            cfg.algo.actor_hidden_dim,
+            cfg.algo.use_layer_norm,
+            device,
+            actor_num_blocks=cfg.algo.algo_params.actor_num_blocks,
+            actor_noise_zeta_mu=cfg.algo.algo_params.actor_noise_zeta_mu,
+            actor_noise_zeta_max=cfg.algo.algo_params.actor_noise_zeta_max,
+        )
+        if cfg.algo.obs_normalization:
+            from unilab.algos.torch.common.normalization import EmpiricalNormalization
+
             normalizer = EmpiricalNormalization(shape=obs_dim, device=device)
     else:
         raise ValueError(f"Unsupported algo: {algo_name}")
@@ -291,8 +356,11 @@ def play_offpolicy(algo_name: str, cfg: DictConfig) -> str | None:
 
     print(f"Loading model: {load_path}")
     checkpoint = torch.load(load_path, map_location=device, weights_only=True)
-    if algo_name == "sac":
+    if algo_name in ("sac", "flashsac"):
         actor.load_state_dict(checkpoint["actor"])
+        if normalizer and checkpoint.get("obs_normalizer"):
+            normalizer.load_state_dict(checkpoint["obs_normalizer"])
+            normalizer.eval()
     else:
         actor_state = {k: v for k, v in checkpoint["actor"].items() if k not in ("noise_scales",)}
         actor.load_state_dict(actor_state, strict=False)
@@ -309,7 +377,7 @@ def _policy_step(obs_np: np.ndarray) -> np.ndarray:
             obs_torch = normalizer(obs_torch, update=False)
         actions_np = (
             actor.explore(obs_torch, deterministic=True).cpu().numpy()
-            if algo_name == "sac"
+            if algo_name in ("sac", "flashsac")
             else actor(obs_torch).cpu().numpy()
         )
         state = env.step(actions_np)

src/unilab/algos/torch/appo/learner.py

@@ -200,7 +200,6 @@ def process_batch(self, batch_dict):
         privileged = batch_dict.get("privileged", None)  # [T, N, P] or None
         rewards = batch_dict["rewards"]  # [T, N]
         dones = batch_dict["dones"].float()  # [T, N]
-        truncated = batch_dict["truncated"].float()  # [T, N]
         last_obs = batch_dict["last_obs"]  # [N, D]
         last_privileged = batch_dict.get("last_privileged", None)  # [N, P] or None
         behavior_log_probs = batch_dict["actions_log_prob"]  # [T, N]
@@ -255,10 +254,6 @@ def process_batch(self, batch_dict):
             batch_dict["_old_sigma"] = self.target_actor.output_std.clone()
         target_log_probs = target_log_probs_flat.view(T, N)
 
-        # Time-out bootstrap correction (matching rsl_rl):
-        #   rewards += gamma * V(s) * time_outs
-        rewards += self.gamma * values.detach() * truncated
-
         # V-trace targets and advantages
         vs, advantages = vtrace_advantages(
             behavior_log_probs=behavior_log_probs,

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

@@ -188,11 +188,21 @@ def learn(
         )
         self._shared_resources.append(shared_storage)
 
-        # Create weight sync
-        weight_sync = SharedWeightSync.from_state_dict(learner.actor.state_dict(), create=True)
-        self._shared_resources.append(weight_sync)
+        # Create weight sync for collector-side actor and critic bootstrap values.
+        actor_weight_sync = SharedWeightSync.from_state_dict(
+            learner.actor.state_dict(), create=True
+        )
+        critic_weight_sync = SharedWeightSync.from_state_dict(
+            learner.critic.state_dict(), create=True
+        )
+        self._shared_resources.extend([actor_weight_sync, critic_weight_sync])
 
-        weight_param_shapes = {name: p.shape for name, p in learner.actor.state_dict().items()}
+        actor_weight_param_shapes = {
+            name: p.shape for name, p in learner.actor.state_dict().items()
+        }
+        critic_weight_param_shapes = {
+            name: p.shape for name, p in learner.critic.state_dict().items()
+        }
 
         metrics_queue: mp.Queue = mp.get_context("spawn").Queue(maxsize=100)
 
@@ -210,8 +220,10 @@ def learn(
             "obs_dim": self.obs_dim,
             "action_dim": self.action_dim,
             "privileged_dim": self.privileged_dim,
-            "weight_sync_name": weight_sync.name,
-            "weight_param_shapes": weight_param_shapes,
+            "actor_weight_sync_name": actor_weight_sync.name,
+            "actor_weight_param_shapes": actor_weight_param_shapes,
+            "critic_weight_sync_name": critic_weight_sync.name,
+            "critic_weight_param_shapes": critic_weight_param_shapes,
             "metrics_queue": metrics_queue,
             "collector_device": self.collector_device,
             "sim_backend": self.extra_kwargs.get("sim_backend", "mujoco"),
@@ -317,7 +329,8 @@ def learn(
             train_start = time.time()
             learner.process_batch(combined)
             metrics = learner.update(combined)
-            weight_sync.write_weights(learner.actor.state_dict())
+            actor_weight_sync.write_weights(learner.actor.state_dict())
+            critic_weight_sync.write_weights(learner.critic.state_dict())
             train_time = time.time() - train_start
 
             metrics["replay_queue_len"] = float(len(replay_queue))