From be8edbf1c13a9c0074d5ed98f6f623a3192a23ba Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Sun, 12 Apr 2026 21:07:00 +0800 Subject: [PATCH 1/5] fix: align timeout bootstrap with final observation contract --- conf/ppo/config.yaml | 2 +- scripts/train_mlx_ppo.py | 28 ++++- src/unilab/algos/torch/appo/learner.py | 5 - src/unilab/algos/torch/appo/runner.py | 27 ++-- src/unilab/algos/torch/appo/worker.py | 136 ++++++++++++++++----- src/unilab/algos/torch/offpolicy/worker.py | 26 ++-- src/unilab/algos/torch/rsl_rl_ppo.py | 50 ++++++++ src/unilab/config/structured_configs.py | 2 +- src/unilab/utils/final_observation.py | 73 +++++++++++ src/unilab/utils/rsl_rl_vec_env_wrapper.py | 9 ++ tests/algos/test_appo_worker.py | 25 ++++ tests/algos/test_rsl_rl_ppo.py | 77 ++++++++++++ tests/config/test_locomotion_params.py | 1 + tests/scripts/test_train_scripts.py | 82 +++++++++++++ tests/utils/test_final_observation.py | 49 ++++++++ 15 files changed, 530 insertions(+), 62 deletions(-) create mode 100644 src/unilab/algos/torch/rsl_rl_ppo.py create mode 100644 src/unilab/utils/final_observation.py create mode 100644 tests/algos/test_appo_worker.py create mode 100644 tests/algos/test_rsl_rl_ppo.py create mode 100644 tests/utils/test_final_observation.py diff --git a/conf/ppo/config.yaml b/conf/ppo/config.yaml index d26869411..543456b18 100644 --- a/conf/ppo/config.yaml +++ b/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 diff --git a/scripts/train_mlx_ppo.py b/scripts/train_mlx_ppo.py index bed76004e..c9452e76b 100644 --- a/scripts/train_mlx_ppo.py +++ b/scripts/train_mlx_ppo.py @@ -128,6 +128,24 @@ 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 + info = getattr(state, "info", None) + if not isinstance(info, dict) or "final_observation" not in info: + return None + final_obs = mx.array(flatten_obs_dict(info["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 +482,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) diff --git a/src/unilab/algos/torch/appo/learner.py b/src/unilab/algos/torch/appo/learner.py index 6144afc84..5381df458 100644 --- a/src/unilab/algos/torch/appo/learner.py +++ b/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, diff --git a/src/unilab/algos/torch/appo/runner.py b/src/unilab/algos/torch/appo/runner.py index 7c5e1e16e..d67071d59 100644 --- a/src/unilab/algos/torch/appo/runner.py +++ b/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)) diff --git a/src/unilab/algos/torch/appo/worker.py b/src/unilab/algos/torch/appo/worker.py index c071cf4b3..fc04d59de 100644 --- a/src/unilab/algos/torch/appo/worker.py +++ b/src/unilab/algos/torch/appo/worker.py @@ -16,7 +16,38 @@ from rsl_rl.utils import resolve_callable from unilab.utils.algo_utils import ensure_registries -from unilab.utils.obs_utils import flatten_obs_dict, split_obs_dict +from unilab.utils.final_observation import resolve_transition_bootstrap_contract +from unilab.utils.obs_utils import split_obs_dict + + +def compute_timeout_bootstrap_correction( + critic: Any, + collector_device: str, + gamma: float, + timeout_mask: np.ndarray, + final_obs: np.ndarray, + final_privileged: np.ndarray | None, +) -> np.ndarray: + """Compute gamma * V(final_observation) for current timeout envs.""" + corrections = np.zeros(timeout_mask.shape, dtype=np.float32) + if not np.any(timeout_mask): + return corrections + + from tensordict import TensorDict + + critic_input_np = final_obs + if final_privileged is not None: + critic_input_np = np.concatenate([final_obs, final_privileged], axis=-1) + critic_input = torch.from_numpy(critic_input_np[timeout_mask]).to(collector_device) + critic_td = TensorDict( + {"policy": critic_input}, + batch_size=critic_input.shape[0], + device=collector_device, + ) + with torch.no_grad(): + bootstrap = critic(critic_td).squeeze(-1).cpu().numpy().astype(np.float32, copy=False) + corrections[timeout_mask] = float(gamma) * bootstrap + return corrections def appo_collector_fn( @@ -30,8 +61,10 @@ def appo_collector_fn( obs_dim: int, action_dim: int, privileged_dim: int, - weight_sync_name: str, - weight_param_shapes: dict, + actor_weight_sync_name: str, + actor_weight_param_shapes: dict, + critic_weight_sync_name: str, + critic_weight_param_shapes: dict, metrics_queue: Any, collector_device: str = "cpu", sim_backend: str = "mujoco", @@ -62,7 +95,12 @@ def appo_collector_fn( shm_name_prefix=shm_storage_name, ) storage.attach_sync_primitives(*sync_primitives) # (write_ptr, read_ptr) - weight_sync = SharedWeightSync(weight_param_shapes, create=False, shm_name=weight_sync_name) + actor_weight_sync = SharedWeightSync( + actor_weight_param_shapes, create=False, shm_name=actor_weight_sync_name + ) + critic_weight_sync = SharedWeightSync( + critic_weight_param_shapes, create=False, shm_name=critic_weight_sync_name + ) # Create environment env: Any = registry.make( @@ -94,11 +132,33 @@ def appo_collector_fn( actor = actor.to(collector_device) actor.eval() + critic_obs_dim = obs_dim + privileged_dim + critic_obs_example = torch.zeros((num_envs, critic_obs_dim), device=collector_device) + critic_td_example = TensorDict({"policy": critic_obs_example}, batch_size=num_envs) + critic_cfg = deepcopy(cfg.get("critic") or cfg.get("actor") or {}) + critic_cls = resolve_callable(critic_cfg.pop("class_name", "rsl_rl.models.MLPModel")) + critic_cfg.pop("num_actions", None) + critic_cfg.pop("distribution_cfg", None) + critic = critic_cls( + critic_td_example, + cfg.get("obs_groups", {"actor": {"policy": critic_obs_dim}}), + "actor", + 1, + **critic_cfg, + ) + critic = critic.to(collector_device) + critic.eval() + # Load initial weights - sd = dict(actor.state_dict()) - weight_sync.read_weights_into(sd) - actor.load_state_dict(sd) - local_weight_version = weight_sync.version + actor_sd = dict(actor.state_dict()) + actor_weight_sync.read_weights_into(actor_sd) + actor.load_state_dict(actor_sd) + local_actor_weight_version = actor_weight_sync.version + + critic_sd = dict(critic.state_dict()) + critic_weight_sync.read_weights_into(critic_sd) + critic.load_state_dict(critic_sd) + local_critic_weight_version = critic_weight_sync.version # Reset environment env_indices = np.arange(num_envs, dtype=np.int32) @@ -141,10 +201,14 @@ def to_float32_np(x): try: while not stop_event.is_set(): # Pull latest weights from learner - if weight_sync.version > local_weight_version: - sd = dict(actor.state_dict()) - local_weight_version = weight_sync.read_weights_into(sd) - actor.load_state_dict(sd) + if actor_weight_sync.version > local_actor_weight_version: + actor_sd = dict(actor.state_dict()) + local_actor_weight_version = actor_weight_sync.read_weights_into(actor_sd) + actor.load_state_dict(actor_sd) + if critic_weight_sync.version > local_critic_weight_version: + critic_sd = dict(critic.state_dict()) + local_critic_weight_version = critic_weight_sync.read_weights_into(critic_sd) + critic.load_state_dict(critic_sd) # Collect one rollout of length steps_per_env write_buf = storage.write_buffer @@ -175,33 +239,39 @@ def to_float32_np(x): next_obs_raw = state.obs reward_raw = np.asarray(state.reward, dtype=np.float32).ravel() - done_raw = np.asarray(state.terminated, dtype=np.float32).ravel() + terminated_raw = np.asarray(state.terminated, dtype=np.float32).ravel() truncated_raw = np.asarray(state.truncated, dtype=np.float32).ravel() + combined_done_raw = np.clip(terminated_raw + truncated_raw, 0, 1) + + next_actor_obs_np, next_actor_priv_np = split_obs_dict(next_obs_raw) + next_actor_obs_np = to_float32_np(next_actor_obs_np) + if next_actor_priv_np is not None: + next_actor_priv_np = to_float32_np(next_actor_priv_np) + transition_contract = resolve_transition_bootstrap_contract( + next_actor_obs_np, + next_actor_priv_np, + state.info, + truncated=truncated_raw, + ) + + reward_raw += compute_timeout_bootstrap_correction( + critic=critic, + collector_device=collector_device, + gamma=float(cfg["algorithm"].get("gamma", 0.99)), + timeout_mask=transition_contract.timeout_final_mask, + final_obs=transition_contract.storage_next_obs, + final_privileged=transition_contract.storage_next_privileged, + ) write_buf["rewards"][:, step] = reward_raw - write_buf["dones"][:, step] = done_raw + write_buf["dones"][:, step] = combined_done_raw write_buf["truncated"][:, step] = truncated_raw - obs_np, priv_np = split_obs_dict(next_obs_raw) - obs_np = to_float32_np(obs_np) - if priv_np is not None: - priv_np = to_float32_np(priv_np) - - # Bootstrap from true terminal obs so next rollout starts cleanly - if "_final_observation" in state.info: - has_final = np.asarray(state.info["_final_observation"], dtype=bool) - if np.any(has_final): - final_obs, final_priv = split_obs_dict(state.info["final_observation"]) - obs_np[has_final] = to_float32_np(final_obs)[has_final] - if priv_np is not None and final_priv is not None: - priv_np[has_final] = to_float32_np(final_priv)[has_final] - # Episode tracking (vectorized) total_steps += num_envs current_ep_rewards += reward_raw current_ep_lengths += 1 - combined_dones = np.clip(done_raw + truncated_raw, 0, 1) - reset_indices = np.where(combined_dones > 0.5)[0] + reset_indices = np.where(combined_done_raw > 0.5)[0] if len(reset_indices) > 0: ep_rewards.extend(current_ep_rewards[reset_indices].tolist()) ep_lengths.extend(current_ep_lengths[reset_indices].tolist()) @@ -246,6 +316,9 @@ def to_float32_np(x): except Exception as e: print(f"[APPOWorker] metrics enqueue error: {e}", file=sys.stderr) + obs_np = transition_contract.actor_next_obs + priv_np = transition_contract.actor_next_privileged + write_buf["last_obs"][:] = obs_np if priv_np is not None: write_buf["last_privileged"][:] = priv_np @@ -265,5 +338,6 @@ def to_float32_np(x): raise storage.close() - weight_sync.close() + actor_weight_sync.close() + critic_weight_sync.close() env.close() diff --git a/src/unilab/algos/torch/offpolicy/worker.py b/src/unilab/algos/torch/offpolicy/worker.py index baa8f50cb..f542c4fc3 100644 --- a/src/unilab/algos/torch/offpolicy/worker.py +++ b/src/unilab/algos/torch/offpolicy/worker.py @@ -11,6 +11,7 @@ import torch from unilab.utils.algo_utils import build_actor, ensure_registries +from unilab.utils.final_observation import resolve_transition_bootstrap_contract from unilab.utils.obs_utils import split_obs_dict @@ -279,29 +280,22 @@ def _run_collector( timeout_count_window += int(np.count_nonzero(timeout_mask_np)) terminated_count_window += int(np.count_nonzero(terminated_mask_np)) - # Handle true terminal observations - if "_final_observation" in state.info: - has_final = state.info["_final_observation"] - has_final_np = np.asarray(has_final, dtype=bool) - if np.any(has_final_np): - final_obs_np, final_priv_np = split_obs_dict(state.info["final_observation"]) - final_obs_np = np.asarray(final_obs_np, dtype=np.float32) - if final_priv_np is not None: - final_priv_np = np.asarray(final_priv_np, dtype=np.float32) - next_obs_np[has_final_np] = final_obs_np[has_final_np] - if final_priv_np is not None and next_priv_np is not None: - next_priv_np[has_final_np] = final_priv_np[has_final_np] + transition_contract = resolve_transition_bootstrap_contract( + next_obs_np, next_priv_np, state.info, truncated=truncated_np + ) # Write to replay buffer replay_buffer.add( torch.from_numpy(obs_np), torch.from_numpy(actions_np), torch.from_numpy(rewards_np), - torch.from_numpy(next_obs_np), + torch.from_numpy(transition_contract.storage_next_obs), torch.from_numpy(terminated_np), torch.from_numpy(truncated_np), torch.from_numpy(priv_np) if priv_np is not None else None, - torch.from_numpy(next_priv_np) if next_priv_np is not None else None, + torch.from_numpy(transition_contract.storage_next_privileged) + if transition_contract.storage_next_privileged is not None + else None, ) # Track episode rewards - vectorized @@ -315,8 +309,8 @@ def _run_collector( current_ep_rewards[reset_indices] = 0.0 current_ep_lengths[reset_indices] = 0 - obs_np = next_obs_np - priv_np = next_priv_np + obs_np = transition_contract.actor_next_obs + priv_np = transition_contract.actor_next_privileged total_steps += num_envs env_steps_since_sync += 1 diff --git a/src/unilab/algos/torch/rsl_rl_ppo.py b/src/unilab/algos/torch/rsl_rl_ppo.py new file mode 100644 index 000000000..b40e3aad6 --- /dev/null +++ b/src/unilab/algos/torch/rsl_rl_ppo.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import torch +from rsl_rl.algorithms import PPO +from tensordict import TensorDict + + +class FinalObservationAwarePPO(PPO): + """PPO variant that bootstraps time limits from env final_observation.""" + + def process_env_step( + self, + obs: TensorDict, + rewards: torch.Tensor, + dones: torch.Tensor, + extras: dict[str, torch.Tensor | TensorDict], + ) -> None: + self.actor.update_normalization(obs) + self.critic.update_normalization(obs) + if self.rnd: + self.rnd.update_normalization(obs) + + self.transition.rewards = rewards.clone() + self.transition.dones = dones + + if self.rnd: + self.intrinsic_rewards = self.rnd.get_intrinsic_reward(obs) + self.transition.rewards += self.intrinsic_rewards + + timeouts = extras.get("time_outs") + timeout_bootstrap_obs = extras.get("time_out_bootstrap_obs") + if isinstance(timeouts, torch.Tensor): + timeout_mask = timeouts.to(self.device).float() + if timeout_bootstrap_obs is not None and torch.count_nonzero(timeout_mask) > 0: + bootstrap_obs = timeout_bootstrap_obs.to(self.device) + bootstrap_values = self.critic(bootstrap_obs).detach() + self.transition.rewards += self.gamma * torch.squeeze( + bootstrap_values * timeout_mask.unsqueeze(1), 1 + ) + else: + transition_values = self.transition.values + assert transition_values is not None + self.transition.rewards += self.gamma * torch.squeeze( + transition_values * timeout_mask.unsqueeze(1), 1 + ) + + self.storage.add_transition(self.transition) + self.transition.clear() + self.actor.reset(dones) + self.critic.reset(dones) diff --git a/src/unilab/config/structured_configs.py b/src/unilab/config/structured_configs.py index 7f020728b..d8511640b 100644 --- a/src/unilab/config/structured_configs.py +++ b/src/unilab/config/structured_configs.py @@ -171,7 +171,7 @@ class PPOPolicyConfig: @dataclass class PPOAlgorithmConfig: - class_name: str = "PPO" + class_name: str = "unilab.algos.torch.rsl_rl_ppo:FinalObservationAwarePPO" value_loss_coef: float = 1.0 use_clipped_value_loss: bool = True clip_param: float = 0.2 diff --git a/src/unilab/utils/final_observation.py b/src/unilab/utils/final_observation.py new file mode 100644 index 000000000..602b61430 --- /dev/null +++ b/src/unilab/utils/final_observation.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from unilab.utils.obs_utils import split_obs_dict + + +@dataclass(frozen=True) +class TransitionBootstrapContract: + actor_next_obs: np.ndarray + actor_next_privileged: np.ndarray | None + storage_next_obs: np.ndarray + storage_next_privileged: np.ndarray | None + final_mask: np.ndarray + timeout_final_mask: np.ndarray + + +def patch_transition_next_obs( + next_obs: np.ndarray, + next_privileged: np.ndarray | None, + info: dict[str, Any] | None, +) -> tuple[np.ndarray, np.ndarray | None, np.ndarray]: + """Patch transition next obs with final_observation without mutating actor inputs.""" + if not isinstance(info, dict): + return next_obs, next_privileged, np.zeros((next_obs.shape[0],), dtype=bool) + + final_mask = np.asarray(info.get("_final_observation"), dtype=bool) + final_obs_dict = info.get("final_observation") + if ( + final_mask.shape != (next_obs.shape[0],) + or not np.any(final_mask) + or not isinstance(final_obs_dict, dict) + ): + return next_obs, next_privileged, np.zeros((next_obs.shape[0],), dtype=bool) + + final_obs, final_privileged = split_obs_dict(final_obs_dict) + storage_next_obs = next_obs.copy() + storage_next_obs[final_mask] = np.asarray(final_obs, dtype=next_obs.dtype)[final_mask] + + storage_next_privileged = next_privileged + if next_privileged is not None and final_privileged is not None: + storage_next_privileged = next_privileged.copy() + storage_next_privileged[final_mask] = np.asarray( + final_privileged, dtype=next_privileged.dtype + )[final_mask] + + return storage_next_obs, storage_next_privileged, final_mask + + +def resolve_transition_bootstrap_contract( + next_obs: np.ndarray, + next_privileged: np.ndarray | None, + info: dict[str, Any] | None, + truncated: np.ndarray | None = None, +) -> TransitionBootstrapContract: + """Resolve actor/storage observations and timeout bootstrap masks for a step.""" + storage_next_obs, storage_next_privileged, final_mask = patch_transition_next_obs( + next_obs, next_privileged, info + ) + timeout_final_mask = final_mask + if truncated is not None: + timeout_final_mask = np.logical_and(final_mask, np.asarray(truncated, dtype=bool).ravel()) + return TransitionBootstrapContract( + actor_next_obs=next_obs, + actor_next_privileged=next_privileged, + storage_next_obs=storage_next_obs, + storage_next_privileged=storage_next_privileged, + final_mask=final_mask, + timeout_final_mask=timeout_final_mask, + ) diff --git a/src/unilab/utils/rsl_rl_vec_env_wrapper.py b/src/unilab/utils/rsl_rl_vec_env_wrapper.py index e9ddd926f..4e720f5f7 100644 --- a/src/unilab/utils/rsl_rl_vec_env_wrapper.py +++ b/src/unilab/utils/rsl_rl_vec_env_wrapper.py @@ -107,6 +107,15 @@ def step(self, actions): if len(done_indices) > 0: if hasattr(state, "truncated"): infos["time_outs"] = to_torch(state.truncated, self.device).bool() + if ( + torch.any(infos["time_outs"]) + and hasattr(state, "info") + and isinstance(state.info, dict) + and "final_observation" in state.info + ): + infos["time_out_bootstrap_obs"] = self._obs_to_tensordict( + state.info["final_observation"] + ) self.episode_returns[done_indices] = 0 self.episode_lengths[done_indices] = 0 diff --git a/tests/algos/test_appo_worker.py b/tests/algos/test_appo_worker.py new file mode 100644 index 000000000..a97fb6ba2 --- /dev/null +++ b/tests/algos/test_appo_worker.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +import numpy as np +import torch + +from unilab.algos.torch.appo.worker import compute_timeout_bootstrap_correction + + +class _FakeCritic: + def __call__(self, obs): + policy = obs["policy"] + return policy.sum(dim=1, keepdim=True) + + +def test_compute_timeout_bootstrap_correction_uses_final_observation_value(): + correction = compute_timeout_bootstrap_correction( + critic=_FakeCritic(), + collector_device="cpu", + gamma=0.5, + timeout_mask=np.array([True, False]), + final_obs=np.array([[2.0, 3.0], [9.0, 9.0]], dtype=np.float32), + final_privileged=np.array([[5.0], [1.0]], dtype=np.float32), + ) + + np.testing.assert_allclose(correction, np.array([5.0, 0.0], dtype=np.float32)) diff --git a/tests/algos/test_rsl_rl_ppo.py b/tests/algos/test_rsl_rl_ppo.py new file mode 100644 index 000000000..4bb7ae9eb --- /dev/null +++ b/tests/algos/test_rsl_rl_ppo.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import torch +from tensordict import TensorDict + +from unilab.algos.torch.rsl_rl_ppo import FinalObservationAwarePPO + + +class _FakeActor: + def update_normalization(self, obs): + return None + + def reset(self, dones): + return None + + +class _FakeCritic: + def __init__(self, values: torch.Tensor): + self.values = values + self.last_obs = None + + def update_normalization(self, obs): + return None + + def reset(self, dones): + return None + + def __call__(self, obs): + self.last_obs = obs + return self.values + + +class _FakeTransition: + def __init__(self): + self.values = torch.tensor([[10.0], [20.0]]) + self.rewards = None + self.dones = None + + def clear(self): + return None + + +class _FakeStorage: + def __init__(self): + self.saved_rewards = None + + def add_transition(self, transition): + self.saved_rewards = transition.rewards.clone() + + +def test_final_observation_aware_ppo_bootstraps_from_final_observation(): + algo = object.__new__(FinalObservationAwarePPO) + algo.actor = _FakeActor() + algo.critic = _FakeCritic(torch.tensor([[3.0], [4.0]])) + algo.rnd = None + algo.gamma = 0.99 + algo.transition = _FakeTransition() + algo.storage = _FakeStorage() + algo.device = "cpu" + + obs = TensorDict({"policy": torch.zeros((2, 1))}, batch_size=[2]) + rewards = torch.tensor([1.0, 2.0]) + dones = torch.tensor([True, True]) + final_obs = TensorDict({"policy": torch.tensor([[30.0], [40.0]])}, batch_size=[2]) + + algo.process_env_step( + obs, + rewards, + dones, + { + "time_outs": torch.tensor([True, False]), + "time_out_bootstrap_obs": final_obs, + }, + ) + + assert torch.allclose(algo.storage.saved_rewards, torch.tensor([1.0 + 0.99 * 3.0, 2.0])) + assert torch.equal(algo.critic.last_obs["policy"], final_obs["policy"]) diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index fd39a2c22..ae2e2cfa6 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -43,6 +43,7 @@ def test_ppo_config_defaults(): assert cfg.algo == "ppo" assert cfg.max_iterations == 101 assert cfg.algorithm.clip_param == 0.2 + assert cfg.algorithm.class_name == "unilab.algos.torch.rsl_rl_ppo:FinalObservationAwarePPO" assert cfg.policy.class_name == "ActorCritic" diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 00264a358..8df129897 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -816,6 +816,42 @@ def test_mlx_get_latest_checkpoint_ignores_non_safetensors(tmp_path): assert mod.get_latest_checkpoint(tmp_path) is None +@pytest.mark.skipif(not _HAS_MLX, reason="mlx not installed") +def test_mlx_time_limit_bootstrap_values_use_final_observation(): + mod = _load_script("train_mlx_ppo") + + class FakeModel: + def __init__(self): + self.last_obs = None + + def value(self, obs): + self.last_obs = obs + return mod.mx.sum(obs, axis=1) + + state = type( + "State", + (), + { + "truncated": np.array([True, False]), + "info": { + "final_observation": { + "obs": np.array([[3.0, 4.0], [9.0, 9.0]], dtype=np.float32), + } + }, + }, + )() + model = FakeModel() + + values = mod.get_time_limit_bootstrap_values(state, model, mod.mx.float32) + + assert values is not None + np.testing.assert_allclose(np.array(values.tolist()), np.array([7.0, 18.0], dtype=np.float32)) + np.testing.assert_allclose( + np.array(model.last_obs.tolist()), + np.array([[3.0, 4.0], [9.0, 9.0]], dtype=np.float32), + ) + + # --------------------------------------------------------------------------- # play_interactive.py — resolve_checkpoint() # --------------------------------------------------------------------------- @@ -983,6 +1019,52 @@ def reset(self, env_indices): assert obs_td["actor"].shape == (1, 3) +def test_play_wrapper_step_exports_timeout_bootstrap_obs(): + import torch + + from unilab.utils.rsl_rl_vec_env_wrapper import RslRlVecEnvWrapper + + class FakeEnv: + def __init__(self): + self.num_envs = 1 + self.cfg = type("Cfg", (), {"max_episode_seconds": 10.0, "ctrl_dt": 0.02})() + self.observation_space = type("Space", (), {"shape": (3,)})() + self.action_space = type("Space", (), {"shape": (2,)})() + self.obs_groups_spec = {"obs": 3} + self.state = type("State", (), {"obs": {"obs": np.zeros((1, 3), dtype=np.float32)}})() + + def init_state(self): + pass + + def reset(self, env_indices): + return {"obs": np.zeros((1, 3), dtype=np.float32)}, {} + + def step(self, actions): + return type( + "StepState", + (), + { + "obs": {"obs": np.array([[1.0, 2.0, 3.0]], dtype=np.float32)}, + "reward": np.array([1.0], dtype=np.float32), + "done": np.array([True]), + "truncated": np.array([True]), + "info": { + "final_observation": {"obs": np.array([[7.0, 8.0, 9.0]], dtype=np.float32)} + }, + }, + )() + + wrapper = RslRlVecEnvWrapper(FakeEnv(), device="cpu", policy_obs_mode="flat") + + _, _, _, infos = wrapper.step(torch.zeros((1, 2))) + + assert torch.equal(infos["time_outs"], torch.tensor([True])) + np.testing.assert_allclose( + infos["time_out_bootstrap_obs"]["policy"].cpu().numpy(), + np.array([[7.0, 8.0, 9.0]], dtype=np.float32), + ) + + # --------------------------------------------------------------------------- # Issue #168: Unified log directory and load_run resolution # --------------------------------------------------------------------------- diff --git a/tests/utils/test_final_observation.py b/tests/utils/test_final_observation.py new file mode 100644 index 000000000..03a216084 --- /dev/null +++ b/tests/utils/test_final_observation.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +import numpy as np + +from unilab.utils.final_observation import ( + patch_transition_next_obs, + resolve_transition_bootstrap_contract, +) + + +def test_patch_transition_next_obs_uses_final_observation_without_mutating_actor_obs(): + next_obs = np.array([[1.0, 1.0], [2.0, 2.0]], dtype=np.float32) + next_privileged = np.array([[10.0], [20.0]], dtype=np.float32) + info = { + "_final_observation": np.array([True, False]), + "final_observation": { + "obs": np.array([[7.0, 8.0], [9.0, 9.0]], dtype=np.float32), + "privileged": np.array([[70.0], [90.0]], dtype=np.float32), + }, + } + + storage_next_obs, storage_next_privileged, final_mask = patch_transition_next_obs( + next_obs, next_privileged, info + ) + + np.testing.assert_array_equal(final_mask, np.array([True, False])) + np.testing.assert_array_equal(storage_next_obs, np.array([[7.0, 8.0], [2.0, 2.0]])) + np.testing.assert_array_equal(storage_next_privileged, np.array([[70.0], [20.0]])) + np.testing.assert_array_equal(next_obs, np.array([[1.0, 1.0], [2.0, 2.0]])) + np.testing.assert_array_equal(next_privileged, np.array([[10.0], [20.0]])) + + +def test_resolve_transition_bootstrap_contract_separates_actor_and_storage_paths(): + contract = resolve_transition_bootstrap_contract( + next_obs=np.array([[1.0, 1.0], [2.0, 2.0]], dtype=np.float32), + next_privileged=None, + info={ + "_final_observation": np.array([False, True]), + "final_observation": {"obs": np.array([[5.0, 5.0], [8.0, 9.0]], dtype=np.float32)}, + }, + truncated=np.array([False, True]), + ) + + np.testing.assert_array_equal(contract.actor_next_obs, np.array([[1.0, 1.0], [2.0, 2.0]])) + np.testing.assert_array_equal( + contract.storage_next_obs, np.array([[1.0, 1.0], [8.0, 9.0]], dtype=np.float32) + ) + np.testing.assert_array_equal(contract.final_mask, np.array([False, True])) + np.testing.assert_array_equal(contract.timeout_final_mask, np.array([False, True])) From 62c80da87b47c7266c70f1c281d603674d22769b Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Sun, 12 Apr 2026 23:04:14 +0800 Subject: [PATCH 2/5] refactor: unify final observation contract --- scripts/train_mlx_ppo.py | 10 +- src/unilab/algos/torch/appo/worker.py | 29 +++-- src/unilab/algos/torch/offpolicy/worker.py | 31 +++-- src/unilab/base/np_env.py | 85 +++++++++--- src/unilab/ipc/replay_buffer.py | 86 ++++++++++++ src/unilab/utils/final_observation.py | 144 ++++++++++++++++----- src/unilab/utils/rsl_rl_vec_env_wrapper.py | 10 +- tests/base/test_np_env.py | 9 +- tests/ipc/test_replay_buffer.py | 28 ++++ tests/scripts/test_train_scripts.py | 6 + tests/utils/test_final_observation.py | 58 ++++++--- 11 files changed, 401 insertions(+), 95 deletions(-) diff --git a/scripts/train_mlx_ppo.py b/scripts/train_mlx_ppo.py index c9452e76b..45c3fc30f 100644 --- a/scripts/train_mlx_ppo.py +++ b/scripts/train_mlx_ppo.py @@ -137,10 +137,14 @@ def get_time_limit_bootstrap_values( timeout_mask = np.asarray(state.truncated, dtype=bool) if not np.any(timeout_mask): return None - info = getattr(state, "info", None) - if not isinstance(info, dict) or "final_observation" not in info: + 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(info["final_observation"])) + 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) diff --git a/src/unilab/algos/torch/appo/worker.py b/src/unilab/algos/torch/appo/worker.py index fc04d59de..4db083152 100644 --- a/src/unilab/algos/torch/appo/worker.py +++ b/src/unilab/algos/torch/appo/worker.py @@ -16,7 +16,7 @@ from rsl_rl.utils import resolve_callable from unilab.utils.algo_utils import ensure_registries -from unilab.utils.final_observation import resolve_transition_bootstrap_contract +from unilab.utils.final_observation import resolve_terminal_observation_contract from unilab.utils.obs_utils import split_obs_dict @@ -247,10 +247,11 @@ def to_float32_np(x): next_actor_obs_np = to_float32_np(next_actor_obs_np) if next_actor_priv_np is not None: next_actor_priv_np = to_float32_np(next_actor_priv_np) - transition_contract = resolve_transition_bootstrap_contract( - next_actor_obs_np, - next_actor_priv_np, - state.info, + terminal_contract = resolve_terminal_observation_contract( + next_obs_batch_size=next_actor_obs_np.shape[0], + final_observation=getattr(state, "final_observation", None), + done=combined_done_raw > 0.5, + info=state.info, truncated=truncated_raw, ) @@ -258,9 +259,17 @@ def to_float32_np(x): critic=critic, collector_device=collector_device, gamma=float(cfg["algorithm"].get("gamma", 0.99)), - timeout_mask=transition_contract.timeout_final_mask, - final_obs=transition_contract.storage_next_obs, - final_privileged=transition_contract.storage_next_privileged, + timeout_mask=terminal_contract.timeout_terminal_mask, + final_obs=( + terminal_contract.terminal_obs + if terminal_contract.terminal_obs is not None + else next_actor_obs_np + ), + final_privileged=( + terminal_contract.terminal_privileged + if terminal_contract.terminal_privileged is not None + else next_actor_priv_np + ), ) write_buf["rewards"][:, step] = reward_raw @@ -316,8 +325,8 @@ def to_float32_np(x): except Exception as e: print(f"[APPOWorker] metrics enqueue error: {e}", file=sys.stderr) - obs_np = transition_contract.actor_next_obs - priv_np = transition_contract.actor_next_privileged + obs_np = next_actor_obs_np + priv_np = next_actor_priv_np write_buf["last_obs"][:] = obs_np if priv_np is not None: diff --git a/src/unilab/algos/torch/offpolicy/worker.py b/src/unilab/algos/torch/offpolicy/worker.py index f542c4fc3..3b1437798 100644 --- a/src/unilab/algos/torch/offpolicy/worker.py +++ b/src/unilab/algos/torch/offpolicy/worker.py @@ -11,7 +11,7 @@ import torch from unilab.utils.algo_utils import build_actor, ensure_registries -from unilab.utils.final_observation import resolve_transition_bootstrap_contract +from unilab.utils.final_observation import resolve_terminal_observation_contract from unilab.utils.obs_utils import split_obs_dict @@ -280,8 +280,12 @@ def _run_collector( timeout_count_window += int(np.count_nonzero(timeout_mask_np)) terminated_count_window += int(np.count_nonzero(terminated_mask_np)) - transition_contract = resolve_transition_bootstrap_contract( - next_obs_np, next_priv_np, state.info, truncated=truncated_np + terminal_contract = resolve_terminal_observation_contract( + next_obs_batch_size=next_obs_np.shape[0], + final_observation=getattr(state, "final_observation", None), + done=done_mask_np, + info=state.info, + truncated=truncated_np, ) # Write to replay buffer @@ -289,13 +293,22 @@ def _run_collector( torch.from_numpy(obs_np), torch.from_numpy(actions_np), torch.from_numpy(rewards_np), - torch.from_numpy(transition_contract.storage_next_obs), + torch.from_numpy(next_obs_np), torch.from_numpy(terminated_np), torch.from_numpy(truncated_np), torch.from_numpy(priv_np) if priv_np is not None else None, - torch.from_numpy(transition_contract.storage_next_privileged) - if transition_contract.storage_next_privileged is not None - else None, + torch.from_numpy(next_priv_np) if next_priv_np is not None else None, + terminal_mask=torch.from_numpy(terminal_contract.terminal_mask), + terminal_next_obs=( + torch.from_numpy(terminal_contract.terminal_obs) + if terminal_contract.terminal_obs is not None + else None + ), + terminal_next_privileged=( + torch.from_numpy(terminal_contract.terminal_privileged) + if terminal_contract.terminal_privileged is not None + else None + ), ) # Track episode rewards - vectorized @@ -309,8 +322,8 @@ def _run_collector( current_ep_rewards[reset_indices] = 0.0 current_ep_lengths[reset_indices] = 0 - obs_np = transition_contract.actor_next_obs - priv_np = transition_contract.actor_next_privileged + obs_np = next_obs_np + priv_np = next_priv_np total_steps += num_envs env_steps_since_sync += 1 diff --git a/src/unilab/base/np_env.py b/src/unilab/base/np_env.py index 4503ea8ee..88b8f34d5 100644 --- a/src/unilab/base/np_env.py +++ b/src/unilab/base/np_env.py @@ -20,7 +20,8 @@ class NpEnvState: reward: np.ndarray terminated: np.ndarray truncated: np.ndarray - info: dict + info: dict[str, Any] + final_observation: dict[str, np.ndarray] | None = None @property def done(self) -> np.ndarray: @@ -39,6 +40,7 @@ def __init__(self, cfg: EnvCfg, backend: SimBackend, num_envs: int): self._num_envs = num_envs self._state: Optional[NpEnvState] = None self._truncated_scratch: np.ndarray = np.zeros((self._num_envs,), dtype=bool) + self._final_observation_scratch: dict[str, np.ndarray] | None = None self.step_counter = 0 self._dr_manager: DomainRandomizationManager | None = None @@ -79,6 +81,7 @@ def init_state(self) -> NpEnvState: self._state = NpEnvState(obs, reward, terminated, truncated, info) self._reset_done_envs() + self._clear_step_final_observation() return self._state def step(self, actions: np.ndarray) -> NpEnvState: @@ -98,9 +101,7 @@ def step(self, actions: np.ndarray) -> NpEnvState: if self._dr_manager is not None: self._dr_manager.apply_interval_randomization_if_due(self.step_counter) self._state.truncated.fill(False) - final_obs_mask = self._state.info.get("_final_observation") - if isinstance(final_obs_mask, np.ndarray): - final_obs_mask.fill(False) + self._clear_step_final_observation() t0 = time.perf_counter() backend_result = self._backend.step(ctrl, self._cfg.sim_substeps) @@ -144,18 +145,17 @@ def _reset_done_envs(self) -> None: env_indices = np.flatnonzero(done).astype(np.int32) self._state.info["steps"][env_indices] = 0 - if "final_observation" not in self._state.info: - self._state.info["final_observation"] = { - k: np.zeros_like(v) for k, v in self._state.obs.items() - } - self._state.info["_final_observation"] = np.zeros((self._num_envs,), dtype=bool) - - self._state.info["_final_observation"][:] = False - self._state.info["_final_observation"][env_indices] = True + final_observation = self._ensure_final_observation_scratch() + compat_final_observation, compat_terminal_mask = ( + self._ensure_final_observation_compat_buffers() + ) + compat_terminal_mask[:] = False + compat_terminal_mask[env_indices] = True for key in self._state.obs: - self._state.info["final_observation"][key][env_indices] = self._state.obs[key][ - env_indices - ] + final_observation[key][env_indices] = self._state.obs[key][env_indices] + compat_final_observation[key][env_indices] = final_observation[key][env_indices] + + self._state.final_observation = final_observation new_obs, info1 = self.reset(env_indices) for key in self._state.obs: @@ -173,6 +173,61 @@ def _reset_done_envs(self) -> None: elif isinstance(value, np.ndarray): self._state.info[key][env_indices] = value + def _ensure_final_observation_scratch(self) -> dict[str, np.ndarray]: + assert self._state is not None + obs = self._state.obs + scratch = self._final_observation_scratch + if scratch is None or set(scratch) != set(obs): + scratch = {key: np.zeros_like(value) for key, value in obs.items()} + self._final_observation_scratch = scratch + else: + for key, value in obs.items(): + if scratch[key].shape != value.shape or scratch[key].dtype != value.dtype: + scratch = { + obs_key: np.zeros_like(obs_value) for obs_key, obs_value in obs.items() + } + self._final_observation_scratch = scratch + break + assert scratch is not None + return scratch + + def _ensure_final_observation_compat_buffers( + self, + ) -> tuple[dict[str, np.ndarray], np.ndarray]: + assert self._state is not None + obs = self._state.obs + compat_final_observation = self._state.info.get("final_observation") + if not isinstance(compat_final_observation, dict) or set(compat_final_observation) != set( + obs + ): + compat_final_observation = {key: np.zeros_like(value) for key, value in obs.items()} + self._state.info["final_observation"] = compat_final_observation + else: + for key, value in obs.items(): + if ( + compat_final_observation[key].shape != value.shape + or compat_final_observation[key].dtype != value.dtype + ): + compat_final_observation = { + obs_key: np.zeros_like(obs_value) for obs_key, obs_value in obs.items() + } + self._state.info["final_observation"] = compat_final_observation + break + compat_terminal_mask = self._state.info.get("_final_observation") + if not isinstance(compat_terminal_mask, np.ndarray) or compat_terminal_mask.shape != ( + self._num_envs, + ): + compat_terminal_mask = np.zeros((self._num_envs,), dtype=bool) + self._state.info["_final_observation"] = compat_terminal_mask + return compat_final_observation, compat_terminal_mask + + def _clear_step_final_observation(self) -> None: + assert self._state is not None + self._state.final_observation = None + compat_terminal_mask = self._state.info.get("_final_observation") + if isinstance(compat_terminal_mask, np.ndarray): + compat_terminal_mask.fill(False) + def _init_domain_randomization(self, provider: "DomainRandomizationProvider") -> None: from unilab.dr import DomainRandomizationManager diff --git a/src/unilab/ipc/replay_buffer.py b/src/unilab/ipc/replay_buffer.py index 39aee9338..44d029f54 100644 --- a/src/unilab/ipc/replay_buffer.py +++ b/src/unilab/ipc/replay_buffer.py @@ -143,6 +143,9 @@ def add( truncated, privileged=None, next_privileged=None, + terminal_mask=None, + terminal_next_obs=None, + terminal_next_privileged=None, ): """Add batch (called by collector).""" n = obs.shape[0] @@ -160,6 +163,15 @@ def add( assert next_privileged is not None self.privileged_obs[idx : idx + n] = privileged self.next_privileged_obs[idx : idx + n] = next_privileged + self._patch_terminal_next_observations( + self.next_obs[idx : idx + n], + terminal_mask, + terminal_next_obs, + self.next_privileged_obs[idx : idx + n] + if self._privileged_dim > 0 and privileged is not None + else None, + terminal_next_privileged, + ) else: split = self.capacity - idx self.obs[idx:] = obs[:split] @@ -180,6 +192,28 @@ def add( self.privileged_obs[: n - split] = privileged[split:] self.next_privileged_obs[idx:] = next_privileged[:split] self.next_privileged_obs[: n - split] = next_privileged[split:] + self._patch_terminal_next_observations( + self.next_obs[idx:], + terminal_mask[:split] if terminal_mask is not None else None, + terminal_next_obs[:split] if terminal_next_obs is not None else None, + self.next_privileged_obs[idx:] + if self._privileged_dim > 0 and privileged is not None + else None, + terminal_next_privileged[:split] + if terminal_next_privileged is not None + else None, + ) + self._patch_terminal_next_observations( + self.next_obs[: n - split], + terminal_mask[split:] if terminal_mask is not None else None, + terminal_next_obs[split:] if terminal_next_obs is not None else None, + self.next_privileged_obs[: n - split] + if self._privileged_dim > 0 and privileged is not None + else None, + terminal_next_privileged[split:] + if terminal_next_privileged is not None + else None, + ) else: parts = [ obs, @@ -196,14 +230,66 @@ def add( if idx + n <= self.capacity: self._storage[idx : idx + n] = row + self._patch_terminal_next_observations( + self._storage[idx : idx + n, self._nobs_sl], + terminal_mask, + terminal_next_obs, + self._storage[idx : idx + n, self._npriv_sl] + if self._privileged_dim > 0 and privileged is not None + else None, + terminal_next_privileged, + ) else: split = self.capacity - idx self._storage[idx:] = row[:split] self._storage[: n - split] = row[split:] + self._patch_terminal_next_observations( + self._storage[idx:, self._nobs_sl], + terminal_mask[:split] if terminal_mask is not None else None, + terminal_next_obs[:split] if terminal_next_obs is not None else None, + self._storage[idx:, self._npriv_sl] + if self._privileged_dim > 0 and privileged is not None + else None, + terminal_next_privileged[:split] + if terminal_next_privileged is not None + else None, + ) + self._patch_terminal_next_observations( + self._storage[: n - split, self._nobs_sl], + terminal_mask[split:] if terminal_mask is not None else None, + terminal_next_obs[split:] if terminal_next_obs is not None else None, + self._storage[: n - split, self._npriv_sl] + if self._privileged_dim > 0 and privileged is not None + else None, + terminal_next_privileged[split:] + if terminal_next_privileged is not None + else None, + ) self.ptr[0] += n self.size[0] = min(int(self.size[0]) + n, self.capacity) + @staticmethod + def _patch_terminal_next_observations( + target_next_obs, + terminal_mask, + terminal_next_obs, + target_next_privileged=None, + terminal_next_privileged=None, + ) -> None: + if terminal_mask is None or terminal_next_obs is None: + return + if terminal_mask.ndim != 1 or terminal_mask.shape[0] != target_next_obs.shape[0]: + return + if not torch.any(terminal_mask): + return + + target_next_obs[terminal_mask] = terminal_next_obs[terminal_mask] + + if target_next_privileged is None or terminal_next_privileged is None: + return + target_next_privileged[terminal_mask] = terminal_next_privileged[terminal_mask] + def sample(self, batch_size: int) -> Dict[str, torch.Tensor]: """Sample batch (called by learner).""" size = int(self.size[0]) diff --git a/src/unilab/utils/final_observation.py b/src/unilab/utils/final_observation.py index 602b61430..a96493765 100644 --- a/src/unilab/utils/final_observation.py +++ b/src/unilab/utils/final_observation.py @@ -12,62 +12,138 @@ class TransitionBootstrapContract: actor_next_obs: np.ndarray actor_next_privileged: np.ndarray | None - storage_next_obs: np.ndarray - storage_next_privileged: np.ndarray | None - final_mask: np.ndarray - timeout_final_mask: np.ndarray + transition_next_obs: np.ndarray + transition_next_privileged: np.ndarray | None + terminal_mask: np.ndarray + timeout_terminal_mask: np.ndarray + + +@dataclass(frozen=True) +class TerminalObservationContract: + terminal_obs: np.ndarray | None + terminal_privileged: np.ndarray | None + terminal_mask: np.ndarray + timeout_terminal_mask: np.ndarray def patch_transition_next_obs( next_obs: np.ndarray, next_privileged: np.ndarray | None, - info: dict[str, Any] | None, + final_observation: dict[str, Any] | None = None, + done: np.ndarray | None = None, + info: dict[str, Any] | None = None, ) -> tuple[np.ndarray, np.ndarray | None, np.ndarray]: """Patch transition next obs with final_observation without mutating actor inputs.""" - if not isinstance(info, dict): - return next_obs, next_privileged, np.zeros((next_obs.shape[0],), dtype=bool) - - final_mask = np.asarray(info.get("_final_observation"), dtype=bool) - final_obs_dict = info.get("final_observation") - if ( - final_mask.shape != (next_obs.shape[0],) - or not np.any(final_mask) - or not isinstance(final_obs_dict, dict) - ): + terminal_contract = resolve_terminal_observation_contract( + next_obs_batch_size=next_obs.shape[0], + final_observation=final_observation, + done=done, + info=info, + ) + if not np.any(terminal_contract.terminal_mask) or terminal_contract.terminal_obs is None: return next_obs, next_privileged, np.zeros((next_obs.shape[0],), dtype=bool) - final_obs, final_privileged = split_obs_dict(final_obs_dict) - storage_next_obs = next_obs.copy() - storage_next_obs[final_mask] = np.asarray(final_obs, dtype=next_obs.dtype)[final_mask] + transition_next_obs = next_obs.copy() + transition_next_obs[terminal_contract.terminal_mask] = np.asarray( + terminal_contract.terminal_obs, dtype=next_obs.dtype + )[terminal_contract.terminal_mask] - storage_next_privileged = next_privileged - if next_privileged is not None and final_privileged is not None: - storage_next_privileged = next_privileged.copy() - storage_next_privileged[final_mask] = np.asarray( - final_privileged, dtype=next_privileged.dtype - )[final_mask] + transition_next_privileged = next_privileged + if next_privileged is not None and terminal_contract.terminal_privileged is not None: + transition_next_privileged = next_privileged.copy() + transition_next_privileged[terminal_contract.terminal_mask] = np.asarray( + terminal_contract.terminal_privileged, dtype=next_privileged.dtype + )[terminal_contract.terminal_mask] - return storage_next_obs, storage_next_privileged, final_mask + return transition_next_obs, transition_next_privileged, terminal_contract.terminal_mask def resolve_transition_bootstrap_contract( next_obs: np.ndarray, next_privileged: np.ndarray | None, - info: dict[str, Any] | None, + info: dict[str, Any] | None = None, + final_observation: dict[str, Any] | None = None, + done: np.ndarray | None = None, truncated: np.ndarray | None = None, ) -> TransitionBootstrapContract: """Resolve actor/storage observations and timeout bootstrap masks for a step.""" - storage_next_obs, storage_next_privileged, final_mask = patch_transition_next_obs( - next_obs, next_privileged, info + transition_next_obs, transition_next_privileged, terminal_mask = patch_transition_next_obs( + next_obs, + next_privileged, + final_observation=final_observation, + done=done, + info=info, ) - timeout_final_mask = final_mask + timeout_terminal_mask = terminal_mask if truncated is not None: - timeout_final_mask = np.logical_and(final_mask, np.asarray(truncated, dtype=bool).ravel()) + timeout_terminal_mask = np.logical_and( + terminal_mask, np.asarray(truncated, dtype=bool).ravel() + ) return TransitionBootstrapContract( actor_next_obs=next_obs, actor_next_privileged=next_privileged, - storage_next_obs=storage_next_obs, - storage_next_privileged=storage_next_privileged, - final_mask=final_mask, - timeout_final_mask=timeout_final_mask, + transition_next_obs=transition_next_obs, + transition_next_privileged=transition_next_privileged, + terminal_mask=terminal_mask, + timeout_terminal_mask=timeout_terminal_mask, + ) + + +def resolve_terminal_observation_contract( + next_obs_batch_size: int, + final_observation: dict[str, Any] | None = None, + done: np.ndarray | None = None, + info: dict[str, Any] | None = None, + truncated: np.ndarray | None = None, +) -> TerminalObservationContract: + """Resolve terminal observation facts without constructing patched next obs.""" + terminal_mask = _resolve_terminal_mask(next_obs_batch_size, done, info) + resolved_final_observation = _resolve_final_observation(final_observation, info) + + terminal_obs: np.ndarray | None = None + terminal_privileged: np.ndarray | None = None + if np.any(terminal_mask) and isinstance(resolved_final_observation, dict): + terminal_obs, terminal_privileged = split_obs_dict(resolved_final_observation) + + timeout_terminal_mask = terminal_mask + if truncated is not None: + timeout_terminal_mask = np.logical_and( + terminal_mask, np.asarray(truncated, dtype=bool).ravel() + ) + + return TerminalObservationContract( + terminal_obs=terminal_obs, + terminal_privileged=terminal_privileged, + terminal_mask=terminal_mask, + timeout_terminal_mask=timeout_terminal_mask, ) + + +def _resolve_final_observation( + final_observation: dict[str, Any] | None, + info: dict[str, Any] | None, +) -> dict[str, Any] | None: + if isinstance(final_observation, dict): + return final_observation + if isinstance(info, dict): + final_obs = info.get("final_observation") + if isinstance(final_obs, dict): + return final_obs + return None + + +def _resolve_terminal_mask( + next_obs_batch_size: int, + done: np.ndarray | None, + info: dict[str, Any] | None, +) -> np.ndarray: + if done is not None: + done_mask = np.asarray(done, dtype=bool).ravel() + if done_mask.shape == (next_obs_batch_size,): + return done_mask + return np.zeros((next_obs_batch_size,), dtype=bool) + if isinstance(info, dict): + terminal_mask = np.asarray(info.get("_final_observation"), dtype=bool) + if terminal_mask.shape == (next_obs_batch_size,): + return terminal_mask + return np.zeros((next_obs_batch_size,), dtype=bool) diff --git a/src/unilab/utils/rsl_rl_vec_env_wrapper.py b/src/unilab/utils/rsl_rl_vec_env_wrapper.py index 4e720f5f7..609c0a726 100644 --- a/src/unilab/utils/rsl_rl_vec_env_wrapper.py +++ b/src/unilab/utils/rsl_rl_vec_env_wrapper.py @@ -107,15 +107,15 @@ def step(self, actions): if len(done_indices) > 0: if hasattr(state, "truncated"): infos["time_outs"] = to_torch(state.truncated, self.device).bool() + final_observation = getattr(state, "final_observation", None) if ( - torch.any(infos["time_outs"]) + final_observation is None and hasattr(state, "info") and isinstance(state.info, dict) - and "final_observation" in state.info ): - infos["time_out_bootstrap_obs"] = self._obs_to_tensordict( - state.info["final_observation"] - ) + final_observation = state.info.get("final_observation") + if torch.any(infos["time_outs"]) and isinstance(final_observation, dict): + infos["time_out_bootstrap_obs"] = self._obs_to_tensordict(final_observation) self.episode_returns[done_indices] = 0 self.episode_lengths[done_indices] = 0 diff --git a/tests/base/test_np_env.py b/tests/base/test_np_env.py index cc8b65067..796465bcc 100644 --- a/tests/base/test_np_env.py +++ b/tests/base/test_np_env.py @@ -341,6 +341,9 @@ def test_final_observation_is_dict(self): env = _TerminatingStubEnv(num_envs=2, terminate_indices=[0]) env.init_state() env.step(np.zeros((2, 3))) + assert isinstance(env.state.final_observation, dict) + assert "obs" in env.state.final_observation + assert "privileged" in env.state.final_observation assert isinstance(env.state.info["final_observation"], dict) assert "obs" in env.state.info["final_observation"] assert "privileged" in env.state.info["final_observation"] @@ -351,6 +354,7 @@ def test_final_observation_captures_pre_reset_obs(self): env.step(np.zeros((2, 3))) # update_state fills actor=1.0 before reset replaces it with 0.0 # final_observation should capture the pre-reset obs (1.0) + np.testing.assert_array_equal(env.state.final_observation["obs"][0], 1.0) np.testing.assert_array_equal(env.state.info["final_observation"]["obs"][0], 1.0) def test_final_observation_mask(self): @@ -363,11 +367,13 @@ def test_final_observation_mask(self): def test_final_observation_mask_clears_on_first_non_terminal_step(self): env = _StubNpEnv(num_envs=3) env.init_state() - np.testing.assert_array_equal(env.state.info["_final_observation"], [True, True, True]) + assert env.state.final_observation is None + np.testing.assert_array_equal(env.state.info["_final_observation"], [False, False, False]) state = env.step(np.zeros((3, 3))) np.testing.assert_array_equal(state.done, [False, False, False]) + assert state.final_observation is None np.testing.assert_array_equal(state.info["_final_observation"], [False, False, False]) def test_final_observation_mask_clears_after_terminal_step(self): @@ -384,6 +390,7 @@ def test_final_observation_mask_clears_after_terminal_step(self): non_terminal_state = env.step(np.zeros((3, 3))) np.testing.assert_array_equal(non_terminal_state.done, [False, False, False]) + assert non_terminal_state.final_observation is None np.testing.assert_array_equal( non_terminal_state.info["_final_observation"], [False, False, False] ) diff --git a/tests/ipc/test_replay_buffer.py b/tests/ipc/test_replay_buffer.py index 8aacf1164..800b86527 100644 --- a/tests/ipc/test_replay_buffer.py +++ b/tests/ipc/test_replay_buffer.py @@ -144,6 +144,34 @@ def test_sample_values_match_added_data(): assert batch["rewards"].max() <= 15.0 +def test_add_patches_terminal_next_obs_without_prebuilding_full_transition_copy(): + buf = _make_buf(capacity=8) + obs = torch.zeros(4, _OBS_DIM) + act = torch.zeros(4, _ACTION_DIM) + rew = torch.zeros(4) + next_obs = torch.arange(4 * _OBS_DIM, dtype=torch.float32).reshape(4, _OBS_DIM) + done = torch.tensor([0.0, 1.0, 0.0, 1.0]) + trunc = torch.tensor([0.0, 1.0, 0.0, 0.0]) + terminal_mask = torch.tensor([False, True, False, True]) + terminal_next_obs = torch.full((4, _OBS_DIM), 99.0) + + buf.add( + obs, + act, + rew, + next_obs, + done, + trunc, + terminal_mask=terminal_mask, + terminal_next_obs=terminal_next_obs, + ) + + stored_next_obs = buf._storage[:4, buf._nobs_sl] + expected = next_obs.clone() + expected[terminal_mask] = terminal_next_obs[terminal_mask] + assert torch.allclose(stored_next_obs, expected) + + # --------------------------------------------------------------------------- # Multiprocess test # --------------------------------------------------------------------------- diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 8df129897..947383059 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -833,6 +833,9 @@ def value(self, obs): (), { "truncated": np.array([True, False]), + "final_observation": { + "obs": np.array([[3.0, 4.0], [9.0, 9.0]], dtype=np.float32), + }, "info": { "final_observation": { "obs": np.array([[3.0, 4.0], [9.0, 9.0]], dtype=np.float32), @@ -1048,6 +1051,9 @@ def step(self, actions): "reward": np.array([1.0], dtype=np.float32), "done": np.array([True]), "truncated": np.array([True]), + "final_observation": { + "obs": np.array([[7.0, 8.0, 9.0]], dtype=np.float32), + }, "info": { "final_observation": {"obs": np.array([[7.0, 8.0, 9.0]], dtype=np.float32)} }, diff --git a/tests/utils/test_final_observation.py b/tests/utils/test_final_observation.py index 03a216084..1bae31357 100644 --- a/tests/utils/test_final_observation.py +++ b/tests/utils/test_final_observation.py @@ -4,6 +4,7 @@ from unilab.utils.final_observation import ( patch_transition_next_obs, + resolve_terminal_observation_contract, resolve_transition_bootstrap_contract, ) @@ -11,21 +12,21 @@ def test_patch_transition_next_obs_uses_final_observation_without_mutating_actor_obs(): next_obs = np.array([[1.0, 1.0], [2.0, 2.0]], dtype=np.float32) next_privileged = np.array([[10.0], [20.0]], dtype=np.float32) - info = { - "_final_observation": np.array([True, False]), - "final_observation": { - "obs": np.array([[7.0, 8.0], [9.0, 9.0]], dtype=np.float32), - "privileged": np.array([[70.0], [90.0]], dtype=np.float32), - }, + final_observation = { + "obs": np.array([[7.0, 8.0], [9.0, 9.0]], dtype=np.float32), + "privileged": np.array([[70.0], [90.0]], dtype=np.float32), } - storage_next_obs, storage_next_privileged, final_mask = patch_transition_next_obs( - next_obs, next_privileged, info + transition_next_obs, transition_next_privileged, terminal_mask = patch_transition_next_obs( + next_obs, + next_privileged, + final_observation=final_observation, + done=np.array([True, False]), ) - np.testing.assert_array_equal(final_mask, np.array([True, False])) - np.testing.assert_array_equal(storage_next_obs, np.array([[7.0, 8.0], [2.0, 2.0]])) - np.testing.assert_array_equal(storage_next_privileged, np.array([[70.0], [20.0]])) + np.testing.assert_array_equal(terminal_mask, np.array([True, False])) + np.testing.assert_array_equal(transition_next_obs, np.array([[7.0, 8.0], [2.0, 2.0]])) + np.testing.assert_array_equal(transition_next_privileged, np.array([[70.0], [20.0]])) np.testing.assert_array_equal(next_obs, np.array([[1.0, 1.0], [2.0, 2.0]])) np.testing.assert_array_equal(next_privileged, np.array([[10.0], [20.0]])) @@ -34,16 +35,37 @@ def test_resolve_transition_bootstrap_contract_separates_actor_and_storage_paths contract = resolve_transition_bootstrap_contract( next_obs=np.array([[1.0, 1.0], [2.0, 2.0]], dtype=np.float32), next_privileged=None, - info={ - "_final_observation": np.array([False, True]), - "final_observation": {"obs": np.array([[5.0, 5.0], [8.0, 9.0]], dtype=np.float32)}, - }, + final_observation={"obs": np.array([[5.0, 5.0], [8.0, 9.0]], dtype=np.float32)}, + done=np.array([False, True]), truncated=np.array([False, True]), ) np.testing.assert_array_equal(contract.actor_next_obs, np.array([[1.0, 1.0], [2.0, 2.0]])) np.testing.assert_array_equal( - contract.storage_next_obs, np.array([[1.0, 1.0], [8.0, 9.0]], dtype=np.float32) + contract.transition_next_obs, np.array([[1.0, 1.0], [8.0, 9.0]], dtype=np.float32) + ) + np.testing.assert_array_equal(contract.terminal_mask, np.array([False, True])) + np.testing.assert_array_equal(contract.timeout_terminal_mask, np.array([False, True])) + + +def test_resolve_terminal_observation_contract_returns_terminal_rows_without_copying_next_obs(): + terminal_contract = resolve_terminal_observation_contract( + next_obs_batch_size=2, + final_observation={ + "obs": np.array([[5.0, 5.0], [8.0, 9.0]], dtype=np.float32), + "privileged": np.array([[1.0], [7.0]], dtype=np.float32), + }, + done=np.array([False, True]), + truncated=np.array([False, True]), + ) + + np.testing.assert_array_equal(terminal_contract.terminal_mask, np.array([False, True])) + np.testing.assert_array_equal(terminal_contract.timeout_terminal_mask, np.array([False, True])) + np.testing.assert_array_equal( + terminal_contract.terminal_obs, + np.array([[5.0, 5.0], [8.0, 9.0]], dtype=np.float32), + ) + np.testing.assert_array_equal( + terminal_contract.terminal_privileged, + np.array([[1.0], [7.0]], dtype=np.float32), ) - np.testing.assert_array_equal(contract.final_mask, np.array([False, True])) - np.testing.assert_array_equal(contract.timeout_final_mask, np.array([False, True])) From 23ca79df1decd211f95f94a56a88844d2cd095f3 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Sun, 12 Apr 2026 23:14:07 +0800 Subject: [PATCH 3/5] fix: degrade mujoco imports for non-slow tests --- scripts/play_interactive.py | 26 ++++++++- src/unilab/utils/render_many.py | 55 ++++++++++++------- tests/base/test_reward_override.py | 6 ++ tests/config/test_reward_injection.py | 5 ++ tests/integration/test_appo_rsl_reward.py | 6 ++ .../test_reward_injection_integration.py | 8 +++ 6 files changed, 85 insertions(+), 21 deletions(-) diff --git a/scripts/play_interactive.py b/scripts/play_interactive.py index b9bea0a88..c0e5c154c 100644 --- a/scripts/play_interactive.py +++ b/scripts/play_interactive.py @@ -18,13 +18,27 @@ import argparse import sys import time +import types from pathlib import Path -import mujoco -import mujoco.viewer import numpy as np import torch +try: + import mujoco + import mujoco.viewer +except ImportError as exc: # pragma: no cover - environment dependent + mujoco = types.SimpleNamespace( # type: ignore[assignment] + viewer=types.SimpleNamespace(launch_passive=None), + MjData=None, + mj_setState=None, + mj_forward=None, + mjtState=types.SimpleNamespace(mjSTATE_FULLPHYSICS=0), + ) + _MUJOCO_IMPORT_ERROR: ImportError | None = exc +else: + _MUJOCO_IMPORT_ERROR = None + ROOT_DIR = Path(__file__).parent.parent sys.path.append(str(ROOT_DIR)) @@ -50,6 +64,13 @@ from tensordict import TensorDict +def _require_mujoco() -> None: + if mujoco is None: # pragma: no cover - environment dependent + raise ImportError( + "MuJoCo is required for interactive play but could not be imported." + ) from _MUJOCO_IMPORT_ERROR + + def _infer_checkpoint_actor_input_dim(ckpt_path: str) -> int | None: loaded = torch.load(ckpt_path, map_location="cpu", weights_only=True) state_dict = loaded.get("actor_state_dict") @@ -387,6 +408,7 @@ def _render_reward_debug_targets( def play_interactive(args): + _require_mujoco() if torch.cuda.is_available(): device = "cuda" elif torch.backends.mps.is_available(): diff --git a/src/unilab/utils/render_many.py b/src/unilab/utils/render_many.py index f604b0001..bdb437be4 100644 --- a/src/unilab/utils/render_many.py +++ b/src/unilab/utils/render_many.py @@ -39,9 +39,24 @@ def _resolve_gl_backend() -> str: # Must be set *before* importing mujoco (it reads the var at import time) os.environ["MUJOCO_GL"] = _resolve_gl_backend() -import mujoco # noqa: E402 import numpy as np +try: # noqa: E402 + import mujoco # type: ignore[reportMissingImports] +except ImportError as exc: # pragma: no cover - environment dependent + mujoco = None # type: ignore[assignment] + _MUJOCO_IMPORT_ERROR: ImportError | None = exc +else: + _MUJOCO_IMPORT_ERROR = None + + +def _require_mujoco(): + if mujoco is None: # pragma: no cover - environment dependent + raise ImportError( + "MuJoCo is required for render_many but could not be imported." + ) from _MUJOCO_IMPORT_ERROR + return mujoco + def get_grid_offsets(num_envs, spacing=1.0): rows = int(math.ceil(math.sqrt(num_envs))) @@ -69,12 +84,13 @@ def init_worker(model_path, shape): """Initialize MuJoCo context for worker process.""" import atexit - _worker_ctx["model"] = mujoco.MjModel.from_xml_path(model_path) + mj = _require_mujoco() + _worker_ctx["model"] = mj.MjModel.from_xml_path(model_path) _worker_ctx["model"].vis.global_.offwidth = 3840 _worker_ctx["model"].vis.global_.offheight = 2160 - _worker_ctx["data"] = mujoco.MjData(_worker_ctx["model"]) - _worker_ctx["renderer"] = mujoco.Renderer(_worker_ctx["model"], height=shape[1], width=shape[0]) + _worker_ctx["data"] = mj.MjData(_worker_ctx["model"]) + _worker_ctx["renderer"] = mj.Renderer(_worker_ctx["model"], height=shape[1], width=shape[0]) atexit.register(_close_worker) @@ -83,6 +99,7 @@ def render_frame_job(args): Worker function to render a single frame. args: (state_batch, offsets, transparent, cam_distance, cam_elevation, cam_azimuth) """ + mj = _require_mujoco() state_batch, offsets, transparent, cam_distance, cam_elevation, cam_azimuth = args model = _worker_ctx["model"] @@ -90,10 +107,10 @@ def render_frame_job(args): renderer = _worker_ctx["renderer"] # Visual options - vopt = mujoco.MjvOption() - vopt.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = transparent - pert = mujoco.MjvPerturb() - catmask = mujoco.mjtCatBit.mjCAT_DYNAMIC + vopt = mj.MjvOption() + vopt.flags[mj.mjtVisFlag.mjVIS_TRANSPARENT] = transparent + pert = mj.MjvPerturb() + catmask = mj.mjtCatBit.mjCAT_DYNAMIC # Helper to set state def set_state(d, s, offset=None): @@ -127,7 +144,7 @@ def set_state(d, s, offset=None): apply_root_offset = True # 2. Box offset - box_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "box") + box_id = mj.mj_name2id(model, mj.mjtObj.mjOBJ_BODY, "box") if box_id >= 0: jnt_adr = model.body_jntadr[box_id] if jnt_adr >= 0: @@ -136,15 +153,15 @@ def set_state(d, s, offset=None): d.qpos[qpos_adr + 1] += offset[1] # 3. Target offset (target_x, target_y) - target_x = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "target_x") + target_x = mj.mj_name2id(model, mj.mjtObj.mjOBJ_JOINT, "target_x") if target_x >= 0: d.qpos[model.jnt_qposadr[target_x]] += offset[0] - target_y = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "target_y") + target_y = mj.mj_name2id(model, mj.mjtObj.mjOBJ_JOINT, "target_y") if target_y >= 0: d.qpos[model.jnt_qposadr[target_y]] += offset[1] - mujoco.mj_forward(model, d) + mj.mj_forward(model, d) # Post-process: Shift all geometries if robot root wasn't moved if apply_root_offset and offset is not None: @@ -160,8 +177,8 @@ def set_state(d, s, offset=None): # Let's iterate bodies. # Simple heuristic: Shift everything except Box and Target? - box_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "box") - target_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "mocap_target") + box_body_id = mj.mj_name2id(model, mj.mjtObj.mjOBJ_BODY, "box") + target_body_id = mj.mj_name2id(model, mj.mjtObj.mjOBJ_BODY, "mocap_target") # Also target might be just a body named "mocap_target" @@ -174,7 +191,7 @@ def set_state(d, s, offset=None): # Everything else (Robot Base, Robot Links, Decoration) should be shifted. is_box_or_target = (body_id == box_body_id) or (body_id == target_body_id) - is_plane = model.geom_type[i] == mujoco.mjtGeom.mjGEOM_PLANE + is_plane = model.geom_type[i] == mj.mjtGeom.mjGEOM_PLANE if not is_box_or_target and not is_plane: d.geom_xpos[i, 0] += offset[0] @@ -196,7 +213,7 @@ def set_state(d, s, offset=None): set_state(data, state_batch[0], offsets[0] if offsets is not None else None) # Init Camera - cam = mujoco.MjvCamera() + cam = mj.MjvCamera() if offsets is not None: center_x = np.mean(offsets[:, 0]) center_y = np.mean(offsets[:, 1]) @@ -204,16 +221,16 @@ def set_state(d, s, offset=None): cam.distance = cam_distance cam.elevation = cam_elevation cam.azimuth = cam_azimuth - cam.type = mujoco.mjtCamera.mjCAMERA_FREE + cam.type = mj.mjtCamera.mjCAMERA_FREE else: - cam.type = mujoco.mjtCamera.mjCAMERA_FREE + cam.type = mj.mjtCamera.mjCAMERA_FREE renderer.update_scene(data, camera=cam, scene_option=vopt) # 2. Add other robots for i in range(1, num_envs): set_state(data, state_batch[i], offsets[i] if offsets is not None else None) - mujoco.mjv_addGeoms(model, data, vopt, pert, catmask, renderer.scene) + mj.mjv_addGeoms(model, data, vopt, pert, catmask, renderer.scene) return renderer.render() diff --git a/tests/base/test_reward_override.py b/tests/base/test_reward_override.py index e77e49a65..e5451376d 100644 --- a/tests/base/test_reward_override.py +++ b/tests/base/test_reward_override.py @@ -6,8 +6,13 @@ from unilab.utils.algo_utils import ensure_registries +def _require_mujoco() -> None: + pytest.importorskip("mujoco", exc_type=ImportError) + + def test_reward_override_go1(): """Test Go1 reward config override.""" + _require_mujoco() ensure_registries() from unilab.envs.locomotion.go1.joystick import RewardConfig @@ -31,6 +36,7 @@ def test_reward_override_go1(): def test_reward_override_g1(): """Test G1 reward config override.""" + _require_mujoco() ensure_registries() from unilab.envs.locomotion.g1.joystick_sac import RewardConfigSAC diff --git a/tests/config/test_reward_injection.py b/tests/config/test_reward_injection.py index 57d8f6223..9cf94a3cd 100644 --- a/tests/config/test_reward_injection.py +++ b/tests/config/test_reward_injection.py @@ -7,6 +7,10 @@ from omegaconf import OmegaConf +def _require_mujoco() -> None: + pytest.importorskip("mujoco", exc_type=ImportError) + + def test_reward_config_loading_g1(): """Test G1 SAC reward config loads correctly.""" with initialize(config_path="../../conf/offpolicy", version_base="1.3"): @@ -44,6 +48,7 @@ def test_resolve_reward_dict_reads_task_reward(): def test_reward_config_conversion(): """Test reward config converts to dataclasses via registry.""" + _require_mujoco() from unilab.base import registry from unilab.utils.algo_utils import ensure_registries diff --git a/tests/integration/test_appo_rsl_reward.py b/tests/integration/test_appo_rsl_reward.py index 53a1e2ebd..65ebb2008 100644 --- a/tests/integration/test_appo_rsl_reward.py +++ b/tests/integration/test_appo_rsl_reward.py @@ -3,8 +3,13 @@ import pytest +def _require_mujoco() -> None: + pytest.importorskip("mujoco", exc_type=ImportError) + + def test_appo_reward_override(): """Test APPO with reward override.""" + _require_mujoco() from unilab.base import registry from unilab.utils.algo_utils import ensure_registries @@ -29,6 +34,7 @@ def test_appo_reward_override(): def test_rsl_rl_reward_override(): """Test RSL-RL with reward override.""" + _require_mujoco() from unilab.base import registry from unilab.utils.algo_utils import ensure_registries diff --git a/tests/integration/test_reward_injection_integration.py b/tests/integration/test_reward_injection_integration.py index 042033eee..5fac964f5 100644 --- a/tests/integration/test_reward_injection_integration.py +++ b/tests/integration/test_reward_injection_integration.py @@ -6,9 +6,14 @@ import pytest +def _require_mujoco() -> None: + pytest.importorskip("mujoco", exc_type=ImportError) + + @pytest.mark.slow def test_reward_injection_in_training(): """Test reward config is properly injected during training.""" + _require_mujoco() from hydra import compose, initialize from scripts.train_offpolicy import build_runner @@ -40,6 +45,7 @@ def test_reward_injection_in_training(): def test_reward_override_propagation(): """Test reward override propagates through multiprocess collector.""" + _require_mujoco() from unilab.base import registry from unilab.envs.locomotion.go1.joystick import RewardConfig from unilab.utils.algo_utils import ensure_registries @@ -88,6 +94,7 @@ def test_reward_override_propagation(): def test_backward_compatibility_no_reward_config(): """Test env requires reward config - should fail without it.""" + _require_mujoco() from unilab.base import registry from unilab.utils.algo_utils import ensure_registries @@ -104,6 +111,7 @@ def test_backward_compatibility_no_reward_config(): def test_zero_scale_skips_computation(): """Test that reward functions with scale=0 are skipped.""" + _require_mujoco() from unilab.base import registry from unilab.envs.locomotion.go1.joystick import RewardConfig from unilab.utils.algo_utils import ensure_registries From 115ac18310612c2397e88ed3e1dbd202fd7bc895 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Sun, 12 Apr 2026 23:30:53 +0800 Subject: [PATCH 4/5] Revert "fix: degrade mujoco imports for non-slow tests" This reverts commit 23ca79df1decd211f95f94a56a88844d2cd095f3. --- scripts/play_interactive.py | 26 +-------- src/unilab/utils/render_many.py | 55 +++++++------------ tests/base/test_reward_override.py | 6 -- tests/config/test_reward_injection.py | 5 -- tests/integration/test_appo_rsl_reward.py | 6 -- .../test_reward_injection_integration.py | 8 --- 6 files changed, 21 insertions(+), 85 deletions(-) diff --git a/scripts/play_interactive.py b/scripts/play_interactive.py index c0e5c154c..b9bea0a88 100644 --- a/scripts/play_interactive.py +++ b/scripts/play_interactive.py @@ -18,27 +18,13 @@ import argparse import sys import time -import types from pathlib import Path +import mujoco +import mujoco.viewer import numpy as np import torch -try: - import mujoco - import mujoco.viewer -except ImportError as exc: # pragma: no cover - environment dependent - mujoco = types.SimpleNamespace( # type: ignore[assignment] - viewer=types.SimpleNamespace(launch_passive=None), - MjData=None, - mj_setState=None, - mj_forward=None, - mjtState=types.SimpleNamespace(mjSTATE_FULLPHYSICS=0), - ) - _MUJOCO_IMPORT_ERROR: ImportError | None = exc -else: - _MUJOCO_IMPORT_ERROR = None - ROOT_DIR = Path(__file__).parent.parent sys.path.append(str(ROOT_DIR)) @@ -64,13 +50,6 @@ from tensordict import TensorDict -def _require_mujoco() -> None: - if mujoco is None: # pragma: no cover - environment dependent - raise ImportError( - "MuJoCo is required for interactive play but could not be imported." - ) from _MUJOCO_IMPORT_ERROR - - def _infer_checkpoint_actor_input_dim(ckpt_path: str) -> int | None: loaded = torch.load(ckpt_path, map_location="cpu", weights_only=True) state_dict = loaded.get("actor_state_dict") @@ -408,7 +387,6 @@ def _render_reward_debug_targets( def play_interactive(args): - _require_mujoco() if torch.cuda.is_available(): device = "cuda" elif torch.backends.mps.is_available(): diff --git a/src/unilab/utils/render_many.py b/src/unilab/utils/render_many.py index bdb437be4..f604b0001 100644 --- a/src/unilab/utils/render_many.py +++ b/src/unilab/utils/render_many.py @@ -39,24 +39,9 @@ def _resolve_gl_backend() -> str: # Must be set *before* importing mujoco (it reads the var at import time) os.environ["MUJOCO_GL"] = _resolve_gl_backend() +import mujoco # noqa: E402 import numpy as np -try: # noqa: E402 - import mujoco # type: ignore[reportMissingImports] -except ImportError as exc: # pragma: no cover - environment dependent - mujoco = None # type: ignore[assignment] - _MUJOCO_IMPORT_ERROR: ImportError | None = exc -else: - _MUJOCO_IMPORT_ERROR = None - - -def _require_mujoco(): - if mujoco is None: # pragma: no cover - environment dependent - raise ImportError( - "MuJoCo is required for render_many but could not be imported." - ) from _MUJOCO_IMPORT_ERROR - return mujoco - def get_grid_offsets(num_envs, spacing=1.0): rows = int(math.ceil(math.sqrt(num_envs))) @@ -84,13 +69,12 @@ def init_worker(model_path, shape): """Initialize MuJoCo context for worker process.""" import atexit - mj = _require_mujoco() - _worker_ctx["model"] = mj.MjModel.from_xml_path(model_path) + _worker_ctx["model"] = mujoco.MjModel.from_xml_path(model_path) _worker_ctx["model"].vis.global_.offwidth = 3840 _worker_ctx["model"].vis.global_.offheight = 2160 - _worker_ctx["data"] = mj.MjData(_worker_ctx["model"]) - _worker_ctx["renderer"] = mj.Renderer(_worker_ctx["model"], height=shape[1], width=shape[0]) + _worker_ctx["data"] = mujoco.MjData(_worker_ctx["model"]) + _worker_ctx["renderer"] = mujoco.Renderer(_worker_ctx["model"], height=shape[1], width=shape[0]) atexit.register(_close_worker) @@ -99,7 +83,6 @@ def render_frame_job(args): Worker function to render a single frame. args: (state_batch, offsets, transparent, cam_distance, cam_elevation, cam_azimuth) """ - mj = _require_mujoco() state_batch, offsets, transparent, cam_distance, cam_elevation, cam_azimuth = args model = _worker_ctx["model"] @@ -107,10 +90,10 @@ def render_frame_job(args): renderer = _worker_ctx["renderer"] # Visual options - vopt = mj.MjvOption() - vopt.flags[mj.mjtVisFlag.mjVIS_TRANSPARENT] = transparent - pert = mj.MjvPerturb() - catmask = mj.mjtCatBit.mjCAT_DYNAMIC + vopt = mujoco.MjvOption() + vopt.flags[mujoco.mjtVisFlag.mjVIS_TRANSPARENT] = transparent + pert = mujoco.MjvPerturb() + catmask = mujoco.mjtCatBit.mjCAT_DYNAMIC # Helper to set state def set_state(d, s, offset=None): @@ -144,7 +127,7 @@ def set_state(d, s, offset=None): apply_root_offset = True # 2. Box offset - box_id = mj.mj_name2id(model, mj.mjtObj.mjOBJ_BODY, "box") + box_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "box") if box_id >= 0: jnt_adr = model.body_jntadr[box_id] if jnt_adr >= 0: @@ -153,15 +136,15 @@ def set_state(d, s, offset=None): d.qpos[qpos_adr + 1] += offset[1] # 3. Target offset (target_x, target_y) - target_x = mj.mj_name2id(model, mj.mjtObj.mjOBJ_JOINT, "target_x") + target_x = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "target_x") if target_x >= 0: d.qpos[model.jnt_qposadr[target_x]] += offset[0] - target_y = mj.mj_name2id(model, mj.mjtObj.mjOBJ_JOINT, "target_y") + target_y = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "target_y") if target_y >= 0: d.qpos[model.jnt_qposadr[target_y]] += offset[1] - mj.mj_forward(model, d) + mujoco.mj_forward(model, d) # Post-process: Shift all geometries if robot root wasn't moved if apply_root_offset and offset is not None: @@ -177,8 +160,8 @@ def set_state(d, s, offset=None): # Let's iterate bodies. # Simple heuristic: Shift everything except Box and Target? - box_body_id = mj.mj_name2id(model, mj.mjtObj.mjOBJ_BODY, "box") - target_body_id = mj.mj_name2id(model, mj.mjtObj.mjOBJ_BODY, "mocap_target") + box_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "box") + target_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "mocap_target") # Also target might be just a body named "mocap_target" @@ -191,7 +174,7 @@ def set_state(d, s, offset=None): # Everything else (Robot Base, Robot Links, Decoration) should be shifted. is_box_or_target = (body_id == box_body_id) or (body_id == target_body_id) - is_plane = model.geom_type[i] == mj.mjtGeom.mjGEOM_PLANE + is_plane = model.geom_type[i] == mujoco.mjtGeom.mjGEOM_PLANE if not is_box_or_target and not is_plane: d.geom_xpos[i, 0] += offset[0] @@ -213,7 +196,7 @@ def set_state(d, s, offset=None): set_state(data, state_batch[0], offsets[0] if offsets is not None else None) # Init Camera - cam = mj.MjvCamera() + cam = mujoco.MjvCamera() if offsets is not None: center_x = np.mean(offsets[:, 0]) center_y = np.mean(offsets[:, 1]) @@ -221,16 +204,16 @@ def set_state(d, s, offset=None): cam.distance = cam_distance cam.elevation = cam_elevation cam.azimuth = cam_azimuth - cam.type = mj.mjtCamera.mjCAMERA_FREE + cam.type = mujoco.mjtCamera.mjCAMERA_FREE else: - cam.type = mj.mjtCamera.mjCAMERA_FREE + cam.type = mujoco.mjtCamera.mjCAMERA_FREE renderer.update_scene(data, camera=cam, scene_option=vopt) # 2. Add other robots for i in range(1, num_envs): set_state(data, state_batch[i], offsets[i] if offsets is not None else None) - mj.mjv_addGeoms(model, data, vopt, pert, catmask, renderer.scene) + mujoco.mjv_addGeoms(model, data, vopt, pert, catmask, renderer.scene) return renderer.render() diff --git a/tests/base/test_reward_override.py b/tests/base/test_reward_override.py index e5451376d..e77e49a65 100644 --- a/tests/base/test_reward_override.py +++ b/tests/base/test_reward_override.py @@ -6,13 +6,8 @@ from unilab.utils.algo_utils import ensure_registries -def _require_mujoco() -> None: - pytest.importorskip("mujoco", exc_type=ImportError) - - def test_reward_override_go1(): """Test Go1 reward config override.""" - _require_mujoco() ensure_registries() from unilab.envs.locomotion.go1.joystick import RewardConfig @@ -36,7 +31,6 @@ def test_reward_override_go1(): def test_reward_override_g1(): """Test G1 reward config override.""" - _require_mujoco() ensure_registries() from unilab.envs.locomotion.g1.joystick_sac import RewardConfigSAC diff --git a/tests/config/test_reward_injection.py b/tests/config/test_reward_injection.py index 9cf94a3cd..57d8f6223 100644 --- a/tests/config/test_reward_injection.py +++ b/tests/config/test_reward_injection.py @@ -7,10 +7,6 @@ from omegaconf import OmegaConf -def _require_mujoco() -> None: - pytest.importorskip("mujoco", exc_type=ImportError) - - def test_reward_config_loading_g1(): """Test G1 SAC reward config loads correctly.""" with initialize(config_path="../../conf/offpolicy", version_base="1.3"): @@ -48,7 +44,6 @@ def test_resolve_reward_dict_reads_task_reward(): def test_reward_config_conversion(): """Test reward config converts to dataclasses via registry.""" - _require_mujoco() from unilab.base import registry from unilab.utils.algo_utils import ensure_registries diff --git a/tests/integration/test_appo_rsl_reward.py b/tests/integration/test_appo_rsl_reward.py index 65ebb2008..53a1e2ebd 100644 --- a/tests/integration/test_appo_rsl_reward.py +++ b/tests/integration/test_appo_rsl_reward.py @@ -3,13 +3,8 @@ import pytest -def _require_mujoco() -> None: - pytest.importorskip("mujoco", exc_type=ImportError) - - def test_appo_reward_override(): """Test APPO with reward override.""" - _require_mujoco() from unilab.base import registry from unilab.utils.algo_utils import ensure_registries @@ -34,7 +29,6 @@ def test_appo_reward_override(): def test_rsl_rl_reward_override(): """Test RSL-RL with reward override.""" - _require_mujoco() from unilab.base import registry from unilab.utils.algo_utils import ensure_registries diff --git a/tests/integration/test_reward_injection_integration.py b/tests/integration/test_reward_injection_integration.py index 5fac964f5..042033eee 100644 --- a/tests/integration/test_reward_injection_integration.py +++ b/tests/integration/test_reward_injection_integration.py @@ -6,14 +6,9 @@ import pytest -def _require_mujoco() -> None: - pytest.importorskip("mujoco", exc_type=ImportError) - - @pytest.mark.slow def test_reward_injection_in_training(): """Test reward config is properly injected during training.""" - _require_mujoco() from hydra import compose, initialize from scripts.train_offpolicy import build_runner @@ -45,7 +40,6 @@ def test_reward_injection_in_training(): def test_reward_override_propagation(): """Test reward override propagates through multiprocess collector.""" - _require_mujoco() from unilab.base import registry from unilab.envs.locomotion.go1.joystick import RewardConfig from unilab.utils.algo_utils import ensure_registries @@ -94,7 +88,6 @@ def test_reward_override_propagation(): def test_backward_compatibility_no_reward_config(): """Test env requires reward config - should fail without it.""" - _require_mujoco() from unilab.base import registry from unilab.utils.algo_utils import ensure_registries @@ -111,7 +104,6 @@ def test_backward_compatibility_no_reward_config(): def test_zero_scale_skips_computation(): """Test that reward functions with scale=0 are skipped.""" - _require_mujoco() from unilab.base import registry from unilab.envs.locomotion.go1.joystick import RewardConfig from unilab.utils.algo_utils import ensure_registries From 8f792f4f7e0389ab9d625dfef31e75cc3e208cf0 Mon Sep 17 00:00:00 2001 From: tatp-yf Date: Sun, 12 Apr 2026 21:34:43 +0800 Subject: [PATCH 5/5] feat: add flashsac offpolicy integration --- conf/offpolicy/algo/flashsac.yaml | 43 ++ .../task/flashsac/g1_sac/motrix.yaml | 28 ++ .../task/flashsac/g1_sac/mujoco.yaml | 28 ++ scripts/train_offpolicy.py | 72 +++- src/unilab/algos/torch/flash_sac/__init__.py | 12 + src/unilab/algos/torch/flash_sac/layers.py | 277 ++++++++++++ src/unilab/algos/torch/flash_sac/learner.py | 407 ++++++++++++++++++ src/unilab/algos/torch/flash_sac/network.py | 193 +++++++++ src/unilab/algos/torch/flash_sac/runner.py | 129 ++++++ src/unilab/algos/torch/flash_sac/update.py | 85 ++++ src/unilab/algos/torch/offpolicy/runner.py | 9 +- src/unilab/algos/torch/offpolicy/worker.py | 12 +- src/unilab/config/structured_configs.py | 52 +++ src/unilab/ipc/async_runner.py | 5 +- src/unilab/utils/algo_utils.py | 23 +- tests/algos/test_flash_sac_learner.py | 77 ++++ tests/config/test_locomotion_params.py | 28 ++ tests/scripts/test_train_scripts.py | 19 + tests/utils/test_algo_utils.py | 15 + 19 files changed, 1507 insertions(+), 7 deletions(-) create mode 100644 conf/offpolicy/algo/flashsac.yaml create mode 100644 conf/offpolicy/task/flashsac/g1_sac/motrix.yaml create mode 100644 conf/offpolicy/task/flashsac/g1_sac/mujoco.yaml create mode 100644 src/unilab/algos/torch/flash_sac/__init__.py create mode 100644 src/unilab/algos/torch/flash_sac/layers.py create mode 100644 src/unilab/algos/torch/flash_sac/learner.py create mode 100644 src/unilab/algos/torch/flash_sac/network.py create mode 100644 src/unilab/algos/torch/flash_sac/runner.py create mode 100644 src/unilab/algos/torch/flash_sac/update.py create mode 100644 tests/algos/test_flash_sac_learner.py diff --git a/conf/offpolicy/algo/flashsac.yaml b/conf/offpolicy/algo/flashsac.yaml new file mode 100644 index 000000000..06c16f793 --- /dev/null +++ b/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 diff --git a/conf/offpolicy/task/flashsac/g1_sac/motrix.yaml b/conf/offpolicy/task/flashsac/g1_sac/motrix.yaml new file mode 100644 index 000000000..d80004b77 --- /dev/null +++ b/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] diff --git a/conf/offpolicy/task/flashsac/g1_sac/mujoco.yaml b/conf/offpolicy/task/flashsac/g1_sac/mujoco.yaml new file mode 100644 index 000000000..7ffb242f7 --- /dev/null +++ b/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] diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index 9941e3df7..f9629ad25 100644 --- a/scripts/train_offpolicy.py +++ b/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) diff --git a/src/unilab/algos/torch/flash_sac/__init__.py b/src/unilab/algos/torch/flash_sac/__init__.py new file mode 100644 index 000000000..2219b855b --- /dev/null +++ b/src/unilab/algos/torch/flash_sac/__init__.py @@ -0,0 +1,12 @@ +"""FlashSAC algorithm package.""" + +from unilab.algos.torch.flash_sac.learner import FlashSACLearner +from unilab.algos.torch.flash_sac.network import FlashSACActor, FlashSACDoubleCritic +from unilab.algos.torch.flash_sac.runner import FlashSACRunner + +__all__ = [ + "FlashSACActor", + "FlashSACDoubleCritic", + "FlashSACLearner", + "FlashSACRunner", +] diff --git a/src/unilab/algos/torch/flash_sac/layers.py b/src/unilab/algos/torch/flash_sac/layers.py new file mode 100644 index 000000000..bd7adb882 --- /dev/null +++ b/src/unilab/algos/torch/flash_sac/layers.py @@ -0,0 +1,277 @@ +"""FlashSAC layers and lightweight normalization helpers.""" + +from __future__ import annotations + +import math +from typing import cast + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def safe_tanh_log_det_jacobian(x: torch.Tensor) -> torch.Tensor: + """Stable log|det J_tanh(x)| term.""" + return cast(torch.Tensor, 2.0 * (math.log(2.0) - x - F.softplus(-2.0 * x))) + + +class UnitLinear(nn.Module): + """Linear layer with post-step weight normalization.""" + + def __init__(self, input_dim: int, output_dim: int): + super().__init__() + self.w = nn.Linear(input_dim, output_dim, bias=False) + nn.init.orthogonal_(self.w.weight, gain=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return cast(torch.Tensor, self.w(x)) + + def normalize_parameters(self) -> None: + with torch.no_grad(): + self.w.weight.copy_(F.normalize(self.w.weight, dim=-1, eps=1e-8)) + + +class UnitBatchNorm(nn.Module): + """BatchNorm variant with normalized affine parameters.""" + + running_mean: torch.Tensor + running_var: torch.Tensor + + def __init__(self, input_dim: int, momentum: float = 0.01, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter(torch.ones(input_dim)) + self.bias = nn.Parameter(torch.zeros(input_dim)) + self.register_buffer("running_mean", torch.zeros(input_dim)) + self.register_buffer("running_var", torch.ones(input_dim)) + self.momentum = momentum + self.eps = eps + + def forward(self, x: torch.Tensor, training: bool) -> torch.Tensor: + return F.batch_norm( + x, + self.running_mean, + self.running_var, + self.weight, + self.bias, + training=training, + momentum=self.momentum, + eps=self.eps, + ) + + def normalize_parameters(self) -> None: + with torch.no_grad(): + sqsum = torch.sum(self.weight * self.weight + self.bias * self.bias, dim=-1) + norm_factor = math.sqrt(float(self.weight.shape[-1])) * torch.rsqrt(sqsum + 1e-8) + self.weight.mul_(norm_factor) + self.bias.mul_(norm_factor) + + +class UnitRMSNorm(nn.Module): + """RMSNorm with unit-length scale vector.""" + + def __init__(self, input_dim: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(input_dim)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + rms = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps) + return (x / rms) * self.weight + + def normalize_parameters(self) -> None: + with torch.no_grad(): + sqsum = torch.sum(self.weight * self.weight, dim=-1) + norm_factor = math.sqrt(float(self.weight.shape[-1])) * torch.rsqrt(sqsum + 1e-8) + self.weight.mul_(norm_factor) + + +class FlashSACEmbedder(nn.Module): + def __init__(self, input_dim: int, hidden_dim: int): + super().__init__() + self.norm = UnitBatchNorm(input_dim) + self.w = UnitLinear(input_dim, hidden_dim) + + def forward(self, x: torch.Tensor, training: bool) -> torch.Tensor: + return cast(torch.Tensor, self.w(self.norm(x, training=training))) + + +class FlashSACBlock(nn.Module): + def __init__(self, hidden_dim: int, expansion: int = 4): + super().__init__() + self.w1 = UnitLinear(hidden_dim, hidden_dim * expansion) + self.norm1 = UnitBatchNorm(hidden_dim * expansion) + self.w2 = UnitLinear(hidden_dim * expansion, hidden_dim) + self.norm2 = UnitBatchNorm(hidden_dim) + + def forward(self, x: torch.Tensor, training: bool) -> torch.Tensor: + residual = x + x = self.w1(x) + x = self.norm1(x, training=training) + x = F.relu(x) + x = self.w2(x) + x = self.norm2(x, training=training) + x = F.relu(x) + return x + residual + + +class NormalTanhPolicy(nn.Module): + def __init__( + self, + hidden_dim: int, + action_dim: int, + log_std_min: float = -10.0, + log_std_max: float = 2.0, + ): + super().__init__() + self.mean_w = UnitLinear(hidden_dim, action_dim) + self.mean_bias = nn.Parameter(torch.zeros(action_dim)) + self.std_w = UnitLinear(hidden_dim, action_dim) + self.std_bias = nn.Parameter(torch.zeros(action_dim)) + self.log_std_min = log_std_min + self.log_std_max = log_std_max + + def get_mean_and_std(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + mean = F.linear(x, self.mean_w.w.weight, self.mean_bias) + raw_log_std = F.linear(x, self.std_w.w.weight, self.std_bias) + log_std = self.log_std_min + (self.log_std_max - self.log_std_min) * 0.5 * ( + 1.0 + torch.tanh(raw_log_std) + ) + std = torch.exp(log_std) + return mean, std + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + mean, std = self.get_mean_and_std(x) + dist = torch.distributions.Normal(mean, std) + raw_action = dist.rsample() + tanh_action = torch.tanh(raw_action) + log_prob = dist.log_prob(raw_action) + log_prob = log_prob - safe_tanh_log_det_jacobian(raw_action) + log_prob = log_prob.sum(dim=-1) + return tanh_action, {"log_prob": log_prob, "mean": mean, "std": std} + + +class EnsembleUnitLinear(nn.Module): + def __init__(self, num_ensemble: int, input_dim: int, output_dim: int): + super().__init__() + self.weight = nn.Parameter(torch.empty(num_ensemble, output_dim, input_dim)) + for idx in range(num_ensemble): + nn.init.orthogonal_(self.weight[idx], gain=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.einsum("nbi,noi->nbo", x, self.weight) + + def normalize_parameters(self) -> None: + with torch.no_grad(): + self.weight.copy_(F.normalize(self.weight, dim=-1, eps=1e-8)) + + +class EnsembleUnitBatchNorm(nn.Module): + running_mean: torch.Tensor + running_var: torch.Tensor + + def __init__( + self, num_ensemble: int, input_dim: int, momentum: float = 0.01, eps: float = 1e-5 + ): + super().__init__() + self.weight = nn.Parameter(torch.ones(num_ensemble, input_dim)) + self.bias = nn.Parameter(torch.zeros(num_ensemble, input_dim)) + self.register_buffer("running_mean", torch.zeros(num_ensemble, input_dim)) + self.register_buffer("running_var", torch.ones(num_ensemble, input_dim)) + self.momentum = momentum + self.eps = eps + + def forward(self, x: torch.Tensor, training: bool) -> torch.Tensor: + if training: + mean = x.mean(dim=1, keepdim=True) + var = x.var(dim=1, correction=0, keepdim=True) + with torch.no_grad(): + batch_size = max(x.shape[1], 1) + correction = batch_size / max(batch_size - 1, 1) + self.running_mean.lerp_(mean.squeeze(1).float(), self.momentum) + self.running_var.lerp_((var.squeeze(1) * correction).float(), self.momentum) + normed = (x - mean) * torch.rsqrt(var + self.eps) + else: + normed = (x - self.running_mean.unsqueeze(1)) * torch.rsqrt( + self.running_var.unsqueeze(1) + self.eps + ) + return normed * self.weight.unsqueeze(1) + self.bias.unsqueeze(1) + + def normalize_parameters(self) -> None: + with torch.no_grad(): + sqsum = torch.sum( + self.weight * self.weight + self.bias * self.bias, dim=-1, keepdim=True + ) + norm_factor = math.sqrt(float(self.weight.shape[-1])) * torch.rsqrt(sqsum + 1e-8) + self.weight.mul_(norm_factor) + self.bias.mul_(norm_factor) + + +class EnsembleUnitRMSNorm(nn.Module): + def __init__(self, num_ensemble: int, input_dim: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(num_ensemble, input_dim)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + rms = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + self.eps) + return (x / rms) * self.weight.unsqueeze(1) + + def normalize_parameters(self) -> None: + with torch.no_grad(): + sqsum = torch.sum(self.weight * self.weight, dim=-1, keepdim=True) + norm_factor = math.sqrt(float(self.weight.shape[-1])) * torch.rsqrt(sqsum + 1e-8) + self.weight.mul_(norm_factor) + + +class EnsembleFlashSACEmbedder(nn.Module): + def __init__(self, num_ensemble: int, input_dim: int, hidden_dim: int): + super().__init__() + self.norm = EnsembleUnitBatchNorm(num_ensemble, input_dim) + self.w = EnsembleUnitLinear(num_ensemble, input_dim, hidden_dim) + + def forward(self, x: torch.Tensor, training: bool) -> torch.Tensor: + return cast(torch.Tensor, self.w(self.norm(x, training=training))) + + +class EnsembleFlashSACBlock(nn.Module): + def __init__(self, num_ensemble: int, hidden_dim: int, expansion: int = 4): + super().__init__() + self.w1 = EnsembleUnitLinear(num_ensemble, hidden_dim, hidden_dim * expansion) + self.norm1 = EnsembleUnitBatchNorm(num_ensemble, hidden_dim * expansion) + self.w2 = EnsembleUnitLinear(num_ensemble, hidden_dim * expansion, hidden_dim) + self.norm2 = EnsembleUnitBatchNorm(num_ensemble, hidden_dim) + + def forward(self, x: torch.Tensor, training: bool) -> torch.Tensor: + residual = x + x = self.w1(x) + x = self.norm1(x, training=training) + x = F.relu(x) + x = self.w2(x) + x = self.norm2(x, training=training) + x = F.relu(x) + return x + residual + + +class EnsembleCategoricalValue(nn.Module): + support: torch.Tensor + + def __init__( + self, + num_ensemble: int, + hidden_dim: int, + num_bins: int, + min_v: float, + max_v: float, + ): + super().__init__() + self.logit_w = EnsembleUnitLinear(num_ensemble, hidden_dim, num_bins) + self.logit_bias = nn.Parameter(torch.zeros(num_ensemble, num_bins)) + self.register_buffer("support", torch.linspace(min_v, max_v, num_bins)) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + logits = self.logit_w(x) + self.logit_bias.unsqueeze(1) + log_probs = F.log_softmax(logits, dim=-1) + probs = log_probs.exp() + support = self.support.view(1, 1, -1) + values = cast(torch.Tensor, torch.sum(probs * support, dim=-1)) + return values, {"log_prob": log_probs} diff --git a/src/unilab/algos/torch/flash_sac/learner.py b/src/unilab/algos/torch/flash_sac/learner.py new file mode 100644 index 000000000..36cdbfc66 --- /dev/null +++ b/src/unilab/algos/torch/flash_sac/learner.py @@ -0,0 +1,407 @@ +"""FlashSAC learner adapted to UniLab's off-policy contract.""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass +from typing import Any, cast + +import torch +import torch.nn as nn +import torch.optim as optim + +from unilab.algos.torch.common.normalization import EmpiricalNormalization +from unilab.algos.torch.flash_sac.network import ( + FlashSACActor, + FlashSACDoubleCritic, + FlashSACTemperature, +) +from unilab.algos.torch.flash_sac.update import ( + build_lr_lambda, + compute_categorical_td_target, + resolve_target_entropy, + select_min_q_log_probs, +) + + +@dataclass +class RunningMeanStd: + mean: torch.Tensor + var: torch.Tensor + count: torch.Tensor + + @classmethod + def create(cls, device: torch.device) -> "RunningMeanStd": + return cls( + mean=torch.zeros(1, device=device, dtype=torch.float32), + var=torch.ones(1, device=device, dtype=torch.float32), + count=torch.tensor(1e-4, device=device, dtype=torch.float32), + ) + + def update(self, x: torch.Tensor) -> None: + x = x.reshape(-1).to(dtype=torch.float32) + if x.numel() == 0: + return + batch_mean = x.mean() + batch_var = x.var(unbiased=False) + batch_count = torch.tensor(float(x.numel()), device=x.device, dtype=torch.float32) + + delta = batch_mean - self.mean + total_count = self.count + batch_count + new_mean = self.mean + delta * batch_count / total_count + m_a = self.var * self.count + m_b = batch_var * batch_count + correction = delta.pow(2) * self.count * batch_count / total_count + new_var = (m_a + m_b + correction) / total_count + + self.mean = new_mean + self.var = new_var + self.count = total_count + + def state_dict(self) -> dict[str, torch.Tensor]: + return {"mean": self.mean, "var": self.var, "count": self.count} + + def load_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None: + self.mean = state_dict["mean"] + self.var = state_dict["var"] + self.count = state_dict["count"] + + +class RewardNormalizer: + """Adaptive reward scaling with running reward statistics.""" + + def __init__(self, g_max: float, device: torch.device, eps: float = 1e-8): + self.g_max = g_max + self.eps = eps + self.device = device + self.rms = RunningMeanStd.create(device) + self.reward_abs_max = torch.tensor(0.0, device=device, dtype=torch.float32) + + def update(self, rewards: torch.Tensor) -> None: + rewards = rewards.reshape(-1).to(device=self.device, dtype=torch.float32) + if rewards.numel() == 0: + return + self.rms.update(rewards) + self.reward_abs_max = torch.maximum(self.reward_abs_max, rewards.abs().max()) + + def normalize(self, rewards: torch.Tensor) -> torch.Tensor: + denominator = torch.maximum( + torch.sqrt(self.rms.var + self.eps), + self.reward_abs_max / max(self.g_max, self.eps), + ) + return rewards / denominator + + def state_dict(self) -> dict[str, Any]: + return { + "rms": self.rms.state_dict(), + "reward_abs_max": self.reward_abs_max, + } + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + self.rms.load_state_dict(state_dict["rms"]) + self.reward_abs_max = state_dict["reward_abs_max"] + + +class FlashSACLearner: + def __init__( + self, + obs_dim: int, + action_dim: int, + privileged_dim: int = 0, + device: str = "cpu", + gamma: float = 0.99, + tau: float = 0.01, + actor_lr: float = 3e-4, + critic_lr: float = 3e-4, + actor_hidden_dim: int = 128, + critic_hidden_dim: int = 256, + actor_num_blocks: int = 2, + critic_num_blocks: int = 2, + num_atoms: int = 101, + critic_min_v: float = -5.0, + critic_max_v: float = 5.0, + temp_initial_value: float = 0.01, + temp_target_sigma: float = 0.15, + temp_target_entropy: float | None = None, + actor_bc_alpha: float = 0.0, + actor_noise_zeta_mu: float = 2.0, + actor_noise_zeta_max: int = 16, + learning_rate_init: float = 3e-4, + learning_rate_peak: float = 3e-4, + learning_rate_end: float = 1.5e-4, + learning_rate_warmup_steps: int = 0, + learning_rate_decay_steps: int = 500000, + normalize_reward: bool = True, + normalized_g_max: float = 5.0, + n_step: int = 1, + obs_normalization: bool = False, + use_amp: bool = False, + use_compile: bool = False, + ): + self.device = torch.device(device) + self.gamma = gamma + self.tau = tau + self.n_step = n_step + self.actor_bc_alpha = actor_bc_alpha + self.obs_dim = obs_dim + self.critic_obs_dim = obs_dim + privileged_dim + self.action_dim = action_dim + self.privileged_dim = privileged_dim + self.update_count = 0 + self.use_amp = bool(use_amp and self.device.type == "cuda") + self.use_compile = bool( + use_compile and hasattr(torch, "compile") and self.device.type == "cuda" + ) + + self.actor = FlashSACActor( + num_blocks=actor_num_blocks, + input_dim=obs_dim, + hidden_dim=actor_hidden_dim, + action_dim=action_dim, + noise_zeta_mu=actor_noise_zeta_mu, + noise_zeta_max=actor_noise_zeta_max, + device=self.device, + ) + self.critic = FlashSACDoubleCritic( + num_blocks=critic_num_blocks, + input_dim=self.critic_obs_dim + action_dim, + hidden_dim=critic_hidden_dim, + num_bins=num_atoms, + min_v=critic_min_v, + max_v=critic_max_v, + device=self.device, + ) + self.target_critic = copy.deepcopy(self.critic).to(self.device) + self.target_critic.eval() + self.temperature = FlashSACTemperature(temp_initial_value).to(self.device) + + self.target_entropy = resolve_target_entropy( + action_dim=action_dim, + target_sigma=temp_target_sigma, + target_entropy=temp_target_entropy, + ) + + self.obs_normalizer: EmpiricalNormalization | nn.Identity + if obs_normalization: + self.obs_normalizer = EmpiricalNormalization(shape=obs_dim, device=self.device) + else: + self.obs_normalizer = nn.Identity() + + self.reward_normalizer = ( + RewardNormalizer(g_max=normalized_g_max, device=self.device) + if normalize_reward + else None + ) + + self.scaler: Any | None = getattr(torch.amp, "GradScaler")("cuda") if self.use_amp else None + lr_peak = learning_rate_peak if learning_rate_peak > 0 else actor_lr + fused = self.device.type == "cuda" + self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=lr_peak, fused=fused) + self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=lr_peak, fused=fused) + self.temperature_optimizer = optim.Adam( + self.temperature.parameters(), lr=lr_peak, fused=fused + ) + + scheduler_fn = build_lr_lambda( + init_lr=learning_rate_init, + peak_lr=lr_peak, + end_lr=learning_rate_end, + warmup_steps=learning_rate_warmup_steps, + decay_steps=learning_rate_decay_steps, + ) + self.actor_scheduler = optim.lr_scheduler.LambdaLR(self.actor_optimizer, scheduler_fn) + self.critic_scheduler = optim.lr_scheduler.LambdaLR(self.critic_optimizer, scheduler_fn) + self.temperature_scheduler = optim.lr_scheduler.LambdaLR( + self.temperature_optimizer, scheduler_fn + ) + + if self.use_compile: + self.actor.get_mean_and_std = torch.compile( # type: ignore[method-assign] + self.actor.get_mean_and_std, mode="reduce-overhead" + ) + + def _maybe_normalize_obs(self, obs: torch.Tensor, *, update: bool) -> torch.Tensor: + if isinstance(self.obs_normalizer, nn.Identity): + return obs + return cast(torch.Tensor, self.obs_normalizer(obs, update=update)) + + def _build_critic_obs( + self, + obs: torch.Tensor, + privileged: torch.Tensor | None, + ) -> torch.Tensor: + if privileged is None: + return obs + return torch.cat([obs, privileged], dim=-1) + + def _autocast(self): + return torch.autocast(device_type="cuda", dtype=torch.float16, enabled=self.use_amp) + + @staticmethod + def _set_requires_grad(module: nn.Module, requires_grad: bool) -> None: + for param in module.parameters(): + param.requires_grad_(requires_grad) + + def update_critic(self, batch: dict[str, torch.Tensor]) -> dict[str, float]: + obs = batch["obs"].to(self.device) + actions = batch["actions"].to(self.device) + rewards = batch["rewards"].to(self.device) + next_obs = batch["next_obs"].to(self.device) + terminated = batch["dones"].to(self.device) + truncated = batch.get("truncated", torch.zeros_like(terminated)).to(self.device) + privileged = batch.get("privileged") + next_privileged = batch.get("next_privileged") + privileged = privileged.to(self.device) if privileged is not None else None + next_privileged = next_privileged.to(self.device) if next_privileged is not None else None + + obs = self._maybe_normalize_obs(obs, update=True) + next_obs = self._maybe_normalize_obs(next_obs, update=False) + critic_obs = self._build_critic_obs(obs, privileged) + critic_next_obs = self._build_critic_obs(next_obs, next_privileged) + + if self.reward_normalizer is not None: + self.reward_normalizer.update(rewards) + rewards = self.reward_normalizer.normalize(rewards) + + gamma = self.gamma**self.n_step + + with torch.no_grad(): + with self._autocast(): + next_actions, actor_info = self.actor(next_obs, training=False) + entropy_bonus = -self.temperature().detach() * actor_info["log_prob"] + next_q_values, next_q_info = self.target_critic( + critic_next_obs, next_actions, training=False + ) + next_q_log_probs = select_min_q_log_probs(next_q_values, next_q_info["log_prob"]) + support = cast(torch.Tensor, self.target_critic.predictor.support) + target_probs = compute_categorical_td_target( + support=support, + target_log_probs=next_q_log_probs, + reward=rewards, + terminated=terminated, + truncated=truncated, + actor_entropy=entropy_bonus, + gamma=gamma, + ) + + with self._autocast(): + _, pred_info = self.critic(critic_obs, actions, training=True) + critic_loss = -(target_probs.unsqueeze(0) * pred_info["log_prob"]).sum(dim=-1).mean() + + self.critic_optimizer.zero_grad(set_to_none=True) + if self.scaler is not None: + self.scaler.scale(critic_loss).backward() + self.scaler.step(self.critic_optimizer) + self.scaler.update() + else: + critic_loss.backward() + self.critic_optimizer.step() + self.critic_scheduler.step() + self.critic.normalize_parameters() + + return { + "critic_loss": float(critic_loss.detach().cpu()), + "reward_scale_std": float( + torch.sqrt(self.reward_normalizer.rms.var).detach().cpu() + if self.reward_normalizer is not None + else torch.tensor(1.0) + ), + } + + def update_actor(self, batch: dict[str, torch.Tensor]) -> dict[str, float]: + obs = batch["obs"].to(self.device) + expert_actions = batch["actions"].to(self.device) + privileged = batch.get("privileged") + privileged = privileged.to(self.device) if privileged is not None else None + + obs = self._maybe_normalize_obs(obs, update=False) + critic_obs = self._build_critic_obs(obs, privileged) + + with self._autocast(): + actions, actor_info = self.actor(obs, training=True) + log_probs = actor_info["log_prob"] + + self._set_requires_grad(self.critic, False) + q_values, _ = self.critic(critic_obs, actions, training=False) + self._set_requires_grad(self.critic, True) + min_q = torch.min(q_values[0], q_values[1]) + + actor_loss = (self.temperature().detach() * log_probs - min_q).mean() + if self.actor_bc_alpha > 0: + bc_loss = torch.mean((actions - expert_actions) ** 2) + actor_loss = ( + actor_loss + self.actor_bc_alpha * min_q.abs().mean().detach() * bc_loss + ) + + self.actor_optimizer.zero_grad(set_to_none=True) + if self.scaler is not None: + self.scaler.scale(actor_loss).backward() + self.scaler.step(self.actor_optimizer) + self.scaler.update() + else: + actor_loss.backward() + self.actor_optimizer.step() + self.actor_scheduler.step() + self.actor.normalize_parameters() + + entropy = -log_probs.detach().mean() + temp_value = self.temperature() + temp_loss = temp_value * (entropy - self.target_entropy) + self.temperature_optimizer.zero_grad(set_to_none=True) + temp_loss.backward() + self.temperature_optimizer.step() + self.temperature_scheduler.step() + + return { + "actor_loss": float(actor_loss.detach().cpu()), + "actor_entropy": float(entropy.detach().cpu()), + "temperature": float(temp_value.detach().cpu()), + "temperature_loss": float(temp_loss.detach().cpu()), + } + + def soft_update_target(self) -> None: + with torch.no_grad(): + for target_param, param in zip( + self.target_critic.parameters(), self.critic.parameters() + ): + target_param.data.mul_(1.0 - self.tau).add_(param.data, alpha=self.tau) + + def get_state_dict(self) -> dict[str, Any]: + return { + "actor": self.actor.state_dict(), + "critic": self.critic.state_dict(), + "target_critic": self.target_critic.state_dict(), + "temperature": self.temperature.state_dict(), + "actor_optimizer": self.actor_optimizer.state_dict(), + "critic_optimizer": self.critic_optimizer.state_dict(), + "temperature_optimizer": self.temperature_optimizer.state_dict(), + "actor_scheduler": self.actor_scheduler.state_dict(), + "critic_scheduler": self.critic_scheduler.state_dict(), + "temperature_scheduler": self.temperature_scheduler.state_dict(), + "obs_normalizer": ( + self.obs_normalizer.state_dict() + if hasattr(self.obs_normalizer, "state_dict") + else None + ), + "reward_normalizer": ( + self.reward_normalizer.state_dict() if self.reward_normalizer is not None else None + ), + "update_count": self.update_count, + } + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + self.actor.load_state_dict(state_dict["actor"]) + self.critic.load_state_dict(state_dict["critic"]) + self.target_critic.load_state_dict(state_dict["target_critic"]) + self.temperature.load_state_dict(state_dict["temperature"]) + self.actor_optimizer.load_state_dict(state_dict["actor_optimizer"]) + self.critic_optimizer.load_state_dict(state_dict["critic_optimizer"]) + self.temperature_optimizer.load_state_dict(state_dict["temperature_optimizer"]) + self.actor_scheduler.load_state_dict(state_dict["actor_scheduler"]) + self.critic_scheduler.load_state_dict(state_dict["critic_scheduler"]) + self.temperature_scheduler.load_state_dict(state_dict["temperature_scheduler"]) + if state_dict.get("obs_normalizer") and hasattr(self.obs_normalizer, "load_state_dict"): + self.obs_normalizer.load_state_dict(state_dict["obs_normalizer"]) + if self.reward_normalizer is not None and state_dict.get("reward_normalizer"): + self.reward_normalizer.load_state_dict(state_dict["reward_normalizer"]) + self.update_count = int(state_dict.get("update_count", 0)) diff --git a/src/unilab/algos/torch/flash_sac/network.py b/src/unilab/algos/torch/flash_sac/network.py new file mode 100644 index 000000000..7d93a2548 --- /dev/null +++ b/src/unilab/algos/torch/flash_sac/network.py @@ -0,0 +1,193 @@ +"""FlashSAC actor, critic, and temperature modules.""" + +from __future__ import annotations + +import math +from typing import cast + +import torch +import torch.nn as nn + +from unilab.algos.torch.flash_sac.layers import ( + EnsembleCategoricalValue, + EnsembleFlashSACBlock, + EnsembleFlashSACEmbedder, + EnsembleUnitRMSNorm, + FlashSACBlock, + FlashSACEmbedder, + NormalTanhPolicy, + UnitRMSNorm, +) + + +def _normalize_module_tree(module: nn.Module) -> None: + with torch.no_grad(): + for child in module.modules(): + if child is module: + continue + normalize = getattr(child, "normalize_parameters", None) + if callable(normalize): + normalize() + + +class FlashSACActor(nn.Module): + zeta_cdf: torch.Tensor + _noise: torch.Tensor + _repeat_count: torch.Tensor + _repeat_target: torch.Tensor + + def __init__( + self, + num_blocks: int, + input_dim: int, + hidden_dim: int, + action_dim: int, + noise_zeta_mu: float = 2.0, + noise_zeta_max: int = 16, + device: str | torch.device = "cpu", + ): + super().__init__() + self.embedder = FlashSACEmbedder(input_dim=input_dim, hidden_dim=hidden_dim) + self.encoder = nn.ModuleList([FlashSACBlock(hidden_dim) for _ in range(num_blocks)]) + self.post_norm = UnitRMSNorm(hidden_dim) + self.predictor = NormalTanhPolicy(hidden_dim=hidden_dim, action_dim=action_dim) + self.noise_zeta_mu = noise_zeta_mu + self.noise_zeta_max = noise_zeta_max + ns = torch.arange(1, noise_zeta_max + 1, dtype=torch.float32) + pmf = ns.pow(-noise_zeta_mu) + self.register_buffer("zeta_cdf", torch.cumsum(pmf / pmf.sum(), dim=0)) + self.register_buffer("_noise", torch.zeros(0), persistent=False) + self.register_buffer("_repeat_count", torch.zeros(0, dtype=torch.int32), persistent=False) + self.register_buffer("_repeat_target", torch.zeros(0, dtype=torch.int32), persistent=False) + self.to(device) + self.normalize_parameters() + + def normalize_parameters(self) -> None: + _normalize_module_tree(self) + + def _encode(self, observations: torch.Tensor, training: bool) -> torch.Tensor: + x = self.embedder(observations, training=training) + for block in self.encoder: + x = block(x, training=training) + return cast(torch.Tensor, self.post_norm(x)) + + def get_mean_and_std( + self, observations: torch.Tensor, training: bool + ) -> tuple[torch.Tensor, torch.Tensor]: + encoded = self._encode(observations, training=training) + return self.predictor.get_mean_and_std(encoded) + + def forward( + self, observations: torch.Tensor, training: bool + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + encoded = self._encode(observations, training=training) + return cast(tuple[torch.Tensor, dict[str, torch.Tensor]], self.predictor(encoded)) + + def _ensure_exploration_state( + self, batch_size: int, action_dim: int, device: torch.device, dtype: torch.dtype + ) -> None: + if ( + self._noise.numel() == batch_size * action_dim + and self._noise.device == device + and self._noise.dtype == dtype + ): + return + self._noise = torch.zeros(batch_size, action_dim, device=device, dtype=dtype) + self._repeat_count = torch.zeros(batch_size, device=device, dtype=torch.int32) + self._repeat_target = torch.zeros(batch_size, device=device, dtype=torch.int32) + + def _sample_repeat_targets(self, batch_size: int, device: torch.device) -> torch.Tensor: + draws = torch.rand(batch_size, device=device) + cdf = self.zeta_cdf.to(device) + return cast(torch.Tensor, torch.searchsorted(cdf, draws).to(torch.int32) + 1) + + @torch.no_grad() + def explore( + self, + obs: torch.Tensor, + dones: torch.Tensor | None = None, + deterministic: bool = False, + ) -> torch.Tensor: + if isinstance(dones, bool): + deterministic = dones + dones = None + mean, std = self.get_mean_and_std(obs, training=False) + if deterministic: + return torch.tanh(mean) + + batch_size, action_dim = mean.shape + self._ensure_exploration_state(batch_size, action_dim, mean.device, mean.dtype) + + if dones is None: + done_mask = torch.zeros(batch_size, device=mean.device, dtype=torch.bool) + else: + done_mask = dones.to(device=mean.device).reshape(-1) > 0.5 + + reinit = done_mask | (self._repeat_count <= 0) | (self._repeat_count >= self._repeat_target) + if torch.any(reinit): + new_noise = torch.randn_like(mean) + new_target = self._sample_repeat_targets(batch_size, mean.device) + self._noise = torch.where(reinit.unsqueeze(-1), new_noise, self._noise) + self._repeat_target = torch.where(reinit, new_target, self._repeat_target) + self._repeat_count = torch.where( + reinit, torch.zeros_like(self._repeat_count), self._repeat_count + ) + + actions = torch.tanh(mean + std * self._noise) + self._repeat_count = self._repeat_count + 1 + return actions + + +class FlashSACDoubleCritic(nn.Module): + def __init__( + self, + num_blocks: int, + input_dim: int, + hidden_dim: int, + num_bins: int, + min_v: float, + max_v: float, + num_qs: int = 2, + device: str | torch.device = "cpu", + ): + super().__init__() + self.embedder = EnsembleFlashSACEmbedder(num_qs, input_dim, hidden_dim) + self.encoder = nn.ModuleList( + [EnsembleFlashSACBlock(num_qs, hidden_dim) for _ in range(num_blocks)] + ) + self.post_norm = EnsembleUnitRMSNorm(num_qs, hidden_dim) + self.predictor = EnsembleCategoricalValue( + num_ensemble=num_qs, + hidden_dim=hidden_dim, + num_bins=num_bins, + min_v=min_v, + max_v=max_v, + ) + self.to(device) + self.normalize_parameters() + + def normalize_parameters(self) -> None: + _normalize_module_tree(self) + + def forward( + self, + observations: torch.Tensor, + actions: torch.Tensor, + training: bool, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + x = torch.cat((observations, actions), dim=-1) + x = x.unsqueeze(0).expand(self.predictor.logit_w.weight.shape[0], -1, -1) + x = self.embedder(x, training=training) + for block in self.encoder: + x = block(x, training=training) + x = self.post_norm(x) + return cast(tuple[torch.Tensor, dict[str, torch.Tensor]], self.predictor(x)) + + +class FlashSACTemperature(nn.Module): + def __init__(self, initial_value: float = 0.01): + super().__init__() + self.log_temp = nn.Parameter(torch.tensor([math.log(initial_value)], dtype=torch.float32)) + + def forward(self) -> torch.Tensor: + return torch.exp(self.log_temp) diff --git a/src/unilab/algos/torch/flash_sac/runner.py b/src/unilab/algos/torch/flash_sac/runner.py new file mode 100644 index 000000000..45986fdda --- /dev/null +++ b/src/unilab/algos/torch/flash_sac/runner.py @@ -0,0 +1,129 @@ +"""FlashSAC runner using the shared off-policy runtime.""" + +from __future__ import annotations + +from typing import Any + +from unilab.algos.torch.flash_sac.learner import FlashSACLearner +from unilab.algos.torch.offpolicy.runner import OffPolicyRunner +from unilab.utils.device_utils import get_default_device + + +class FlashSACRunner(OffPolicyRunner): + def __init__( + self, + env_name: str, + env_cfg_override: dict[str, Any] | None = None, + device: str | None = None, + num_envs: int = 2048, + replay_buffer_n: int = 512, + batch_size: int = 2048, + warmup_steps: int = 10000, + updates_per_step: int = 1, + policy_frequency: int = 2, + sync_collection: bool = True, + env_steps_per_sync: int = 1, + gamma: float = 0.99, + tau: float = 0.01, + actor_lr: float = 3e-4, + critic_lr: float = 3e-4, + obs_normalization: bool = False, + actor_hidden_dim: int = 128, + critic_hidden_dim: int = 256, + num_atoms: int = 101, + use_amp: bool = False, + sim_backend: str = "mujoco", + actor_num_blocks: int = 2, + critic_num_blocks: int = 2, + actor_bc_alpha: float = 0.0, + actor_noise_zeta_mu: float = 2.0, + actor_noise_zeta_max: int = 16, + critic_min_v: float = -5.0, + critic_max_v: float = 5.0, + target_sigma: float = 0.15, + target_entropy: float | None = None, + temp_initial_value: float = 0.01, + learning_rate_init: float = 3e-4, + learning_rate_peak: float = 3e-4, + learning_rate_end: float = 1.5e-4, + learning_rate_warmup_steps: int = 0, + learning_rate_decay_steps: int = 500000, + normalize_reward: bool = True, + normalized_g_max: float = 5.0, + n_step: int = 1, + use_compile: bool = False, + ): + from unilab.base import registry + from unilab.utils.algo_utils import ensure_registries + from unilab.utils.obs_utils import get_obs_dims + + ensure_registries() + env: Any = registry.make( + env_name, num_envs=1, sim_backend=sim_backend, env_cfg_override=env_cfg_override + ) + obs_dim, privileged_dim = get_obs_dims(env.obs_groups_spec) + action_shape = env.action_space.shape + assert action_shape is not None + action_dim = int(action_shape[0]) + env.close() + + runtime_device = device or get_default_device() + learner = FlashSACLearner( + obs_dim=obs_dim, + action_dim=action_dim, + privileged_dim=privileged_dim, + device=runtime_device, + gamma=gamma, + tau=tau, + actor_lr=actor_lr, + critic_lr=critic_lr, + actor_hidden_dim=actor_hidden_dim, + critic_hidden_dim=critic_hidden_dim, + actor_num_blocks=actor_num_blocks, + critic_num_blocks=critic_num_blocks, + num_atoms=num_atoms, + critic_min_v=critic_min_v, + critic_max_v=critic_max_v, + temp_initial_value=temp_initial_value, + temp_target_sigma=target_sigma, + temp_target_entropy=target_entropy, + actor_bc_alpha=actor_bc_alpha, + actor_noise_zeta_mu=actor_noise_zeta_mu, + actor_noise_zeta_max=actor_noise_zeta_max, + learning_rate_init=learning_rate_init, + learning_rate_peak=learning_rate_peak, + learning_rate_end=learning_rate_end, + learning_rate_warmup_steps=learning_rate_warmup_steps, + learning_rate_decay_steps=learning_rate_decay_steps, + normalize_reward=normalize_reward, + normalized_g_max=normalized_g_max, + n_step=n_step, + obs_normalization=obs_normalization, + use_amp=use_amp, + use_compile=use_compile, + ) + + super().__init__( + learner=learner, + env_name=env_name, + algo_type="flashsac", + num_envs=num_envs, + replay_buffer_n=replay_buffer_n, + batch_size=batch_size, + warmup_steps=warmup_steps, + updates_per_step=updates_per_step, + policy_frequency=policy_frequency, + sync_collection=sync_collection, + env_steps_per_sync=env_steps_per_sync, + device=runtime_device, + actor_hidden_dim=actor_hidden_dim, + use_layer_norm=False, + obs_normalization=obs_normalization, + sim_backend=sim_backend, + env_cfg_override=env_cfg_override, + actor_kwargs={ + "actor_num_blocks": actor_num_blocks, + "actor_noise_zeta_mu": actor_noise_zeta_mu, + "actor_noise_zeta_max": actor_noise_zeta_max, + }, + ) diff --git a/src/unilab/algos/torch/flash_sac/update.py b/src/unilab/algos/torch/flash_sac/update.py new file mode 100644 index 000000000..4ad6010e7 --- /dev/null +++ b/src/unilab/algos/torch/flash_sac/update.py @@ -0,0 +1,85 @@ +"""FlashSAC update helpers.""" + +from __future__ import annotations + +import math + +import torch + + +def build_lr_lambda( + init_lr: float, + peak_lr: float, + end_lr: float, + warmup_steps: int, + decay_steps: int, +): + init_ratio = init_lr / peak_lr if peak_lr > 0 else 1.0 + end_ratio = end_lr / peak_lr if peak_lr > 0 else 1.0 + warmup_steps = max(int(warmup_steps), 0) + decay_steps = max(int(decay_steps), warmup_steps + 1) + + def schedule(step: int) -> float: + if warmup_steps > 0 and step < warmup_steps: + progress = float(step) / float(max(warmup_steps, 1)) + return init_ratio + (1.0 - init_ratio) * progress + progress = min(max((step - warmup_steps) / float(decay_steps - warmup_steps), 0.0), 1.0) + cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) + return end_ratio + (1.0 - end_ratio) * cosine + + return schedule + + +def select_min_q_log_probs( + next_q_values: torch.Tensor, + next_q_log_probs: torch.Tensor, +) -> torch.Tensor: + num_bins = next_q_log_probs.shape[-1] + min_indices = next_q_values.argmin(dim=0) + gather_index = min_indices[None, :, None].expand(1, -1, num_bins) + return torch.gather(next_q_log_probs, dim=0, index=gather_index)[0] + + +def compute_categorical_td_target( + support: torch.Tensor, + target_log_probs: torch.Tensor, + reward: torch.Tensor, + terminated: torch.Tensor, + truncated: torch.Tensor, + actor_entropy: torch.Tensor, + gamma: float, +) -> torch.Tensor: + batch_size, num_bins = target_log_probs.shape + support = support.view(1, -1) + reward = reward.view(-1, 1) + terminated = terminated.view(-1, 1) + truncated = truncated.view(-1, 1) + actor_entropy = actor_entropy.view(-1, 1) + + bootstrap = torch.clamp(1.0 - terminated + truncated, 0.0, 1.0) + target_bin_values = reward + bootstrap * gamma * (support - actor_entropy) + target_bin_values = torch.clamp(target_bin_values, float(support.min()), float(support.max())) + + bin_width = float(support[0, 1] - support[0, 0]) + offsets = (target_bin_values - float(support.min())) / max(bin_width, 1e-8) + lower = torch.floor(offsets).long().clamp(0, num_bins - 1) + upper = torch.ceil(offsets).long().clamp(0, num_bins - 1) + frac = offsets - lower.float() + + probs = target_log_probs.exp() + target_probs = torch.zeros(batch_size, num_bins, dtype=probs.dtype, device=probs.device) + target_probs.scatter_add_(1, lower, probs * (1.0 - frac)) + target_probs.scatter_add_(1, upper, probs * frac) + return target_probs + + +def resolve_target_entropy( + action_dim: int, + target_sigma: float, + target_entropy: float | None, +) -> float: + if target_entropy is not None: + return float(target_entropy) + sigma = max(float(target_sigma), 1e-6) + per_dim_entropy = 0.5 * math.log(2.0 * math.pi * math.e * sigma * sigma) + return float(action_dim) * per_dim_entropy diff --git a/src/unilab/algos/torch/offpolicy/runner.py b/src/unilab/algos/torch/offpolicy/runner.py index c23c97d6f..73d3b2145 100644 --- a/src/unilab/algos/torch/offpolicy/runner.py +++ b/src/unilab/algos/torch/offpolicy/runner.py @@ -23,7 +23,7 @@ def __init__( self, learner, env_name: str, - algo_type: str, # "sac" or "td3" + algo_type: str, # "sac", "td3", or "flashsac" num_envs: int = 4096, replay_buffer_n: int = 1024, batch_size: int = 8192, @@ -38,6 +38,7 @@ def __init__( obs_normalization: bool = False, sim_backend: str = "mujoco", env_cfg_override: dict | None = None, + actor_kwargs: dict | None = None, ): super().__init__( env_name=env_name, @@ -62,6 +63,7 @@ def __init__( self.actor_hidden_dim = actor_hidden_dim self.use_layer_norm = use_layer_norm self.obs_normalization = obs_normalization + self.actor_kwargs = actor_kwargs or {} self.obs_dim, self.action_dim, self.privileged_dim = get_env_dims( self.env_name, sim_backend, env_cfg_override @@ -145,6 +147,7 @@ def learn( "env_cfg_override": self.env_cfg_override, "obs_dim": self.obs_dim, "action_dim": self.action_dim, + "actor_kwargs": self.actor_kwargs, } self._start_collector( target_fn=off_policy_collector_fn, @@ -157,7 +160,9 @@ def learn( # Setup logger logger = OffPolicyLogger( - algo_name=f"Fast{self.algo_type.upper()}", + algo_name=( + "FlashSAC" if self.algo_type == "flashsac" else f"Fast{self.algo_type.upper()}" + ), max_iterations=max_iterations, num_envs=self.num_envs, env_name=self.env_name, diff --git a/src/unilab/algos/torch/offpolicy/worker.py b/src/unilab/algos/torch/offpolicy/worker.py index 3b1437798..778c29b63 100644 --- a/src/unilab/algos/torch/offpolicy/worker.py +++ b/src/unilab/algos/torch/offpolicy/worker.py @@ -62,6 +62,7 @@ def off_policy_collector_fn( env_cfg_override: dict | None = None, obs_dim: int | None = None, action_dim: int | None = None, + actor_kwargs: dict | None = None, **kwargs, ): """Entry point for the off-policy collector subprocess.""" @@ -93,6 +94,7 @@ def off_policy_collector_fn( env_cfg_override=env_cfg_override, obs_dim=obs_dim, action_dim=action_dim, + actor_kwargs=actor_kwargs, ) except Exception as e: print(f"[Collector] Exception: {e}", file=sys.stderr, flush=True) @@ -127,6 +129,7 @@ def _run_collector( env_cfg_override, obs_dim, action_dim, + actor_kwargs, ): from unilab.base import registry from unilab.ipc import SharedWeightSync @@ -152,7 +155,14 @@ def _run_collector( action_dim=action_dim, ) actor = build_actor( - algo_type, obs_dim, action_dim, actor_hidden_dim, use_layer_norm, "cpu", num_envs + algo_type, + obs_dim, + action_dim, + actor_hidden_dim, + use_layer_norm, + "cpu", + num_envs, + **(actor_kwargs or {}), ) actor.eval() diff --git a/src/unilab/config/structured_configs.py b/src/unilab/config/structured_configs.py index d8511640b..afd3fb9a3 100644 --- a/src/unilab/config/structured_configs.py +++ b/src/unilab/config/structured_configs.py @@ -96,6 +96,58 @@ class TD3Config(BaseConfig): algo_params: TD3AlgoParams = field(default_factory=TD3AlgoParams) +# ── Off-policy: FlashSAC ───────────────────────────────────────────────────── + + +@dataclass +class FlashSACAlgoParams: + normalize_reward: bool = True + normalized_g_max: float = 5.0 + actor_num_blocks: int = 2 + critic_num_blocks: int = 2 + actor_bc_alpha: float = 0.0 + actor_noise_zeta_mu: float = 2.0 + actor_noise_zeta_max: int = 16 + critic_min_v: float = -5.0 + critic_max_v: float = 5.0 + temp_initial_value: float = 0.01 + temp_target_sigma: float = 0.15 + temp_target_entropy: float | None = None + learning_rate_init: float = 3e-4 + learning_rate_peak: float = 3e-4 + learning_rate_end: float = 1.5e-4 + learning_rate_warmup_steps: int = 0 + learning_rate_decay_steps: int = 500000 + n_step: int = 1 + use_compile: bool = False + + +@dataclass +class FlashSACConfig(BaseConfig): + algo: str = "flashsac" + algo_log_name: str = "flash_sac" + seed: int = 1 + num_envs: int = 2048 + batch_size: int = 2048 + replay_buffer_n: int = 512 + updates_per_step: int = 1 + warmup_steps: int = 10000 + policy_frequency: int = 2 + env_steps_per_sync: int = 1 + max_iterations: int = 5000 + save_interval: int = 1000 + gamma: float = 0.99 + tau: float = 0.01 + actor_lr: float = 3e-4 + critic_lr: float = 3e-4 + actor_hidden_dim: int = 128 + critic_hidden_dim: int = 256 + num_atoms: int = 101 + obs_normalization: bool = False + use_layer_norm: bool = False + algo_params: FlashSACAlgoParams = field(default_factory=FlashSACAlgoParams) + + # ── APPO ───────────────────────────────────────────────────────────────────── diff --git a/src/unilab/ipc/async_runner.py b/src/unilab/ipc/async_runner.py index e3c91ff5f..3352e5007 100644 --- a/src/unilab/ipc/async_runner.py +++ b/src/unilab/ipc/async_runner.py @@ -58,7 +58,10 @@ def learn( ) -> None: ... def _start_collector(self, target_fn: Callable, kwargs: dict) -> None: - self._collector_process = _SPAWN_CTX.Process(target=target_fn, kwargs=kwargs, daemon=True) + collector_kwargs = {**self.extra_kwargs, **kwargs} + self._collector_process = _SPAWN_CTX.Process( + target=target_fn, kwargs=collector_kwargs, daemon=True + ) self._collector_process.start() def close(self) -> None: diff --git a/src/unilab/utils/algo_utils.py b/src/unilab/utils/algo_utils.py index a4895365b..767f02f3a 100644 --- a/src/unilab/utils/algo_utils.py +++ b/src/unilab/utils/algo_utils.py @@ -84,7 +84,16 @@ def ensure_registries( def build_actor( - algo_type, obs_dim, action_dim, actor_hidden_dim, use_layer_norm, device, num_envs=1 + algo_type, + obs_dim, + action_dim, + actor_hidden_dim, + use_layer_norm, + device, + num_envs=1, + actor_num_blocks: int = 2, + actor_noise_zeta_mu: float = 2.0, + actor_noise_zeta_max: int = 16, ): """Build the correct actor model based on algorithm type.""" if algo_type == "sac": @@ -110,5 +119,17 @@ def build_actor( log_std_max=0.0, device=device, ) + elif algo_type == "flashsac": + from unilab.algos.torch.flash_sac.network import FlashSACActor + + return FlashSACActor( + num_blocks=actor_num_blocks, + input_dim=obs_dim, + hidden_dim=actor_hidden_dim, + action_dim=action_dim, + noise_zeta_mu=actor_noise_zeta_mu, + noise_zeta_max=actor_noise_zeta_max, + device=device, + ) else: raise ValueError(f"Unknown algo_type: {algo_type}") diff --git a/tests/algos/test_flash_sac_learner.py b/tests/algos/test_flash_sac_learner.py new file mode 100644 index 000000000..392952526 --- /dev/null +++ b/tests/algos/test_flash_sac_learner.py @@ -0,0 +1,77 @@ +"""Unit tests for FlashSAC learner and actor interfaces.""" + +from __future__ import annotations + +import torch + +from unilab.algos.torch.flash_sac.learner import FlashSACLearner + + +def _make_batch(batch_size: int = 32) -> dict[str, torch.Tensor]: + obs = torch.randn(batch_size, 98) + privileged = torch.randn(batch_size, 3) + actions = torch.tanh(torch.randn(batch_size, 29)) + rewards = torch.randn(batch_size) + next_obs = torch.randn(batch_size, 98) + next_privileged = torch.randn(batch_size, 3) + dones = torch.zeros(batch_size) + truncated = torch.zeros(batch_size) + return { + "obs": obs, + "privileged": privileged, + "actions": actions, + "rewards": rewards, + "next_obs": next_obs, + "next_privileged": next_privileged, + "dones": dones, + "truncated": truncated, + } + + +def test_flashsac_learner_exposes_expected_dims(): + learner = FlashSACLearner(obs_dim=98, action_dim=29, privileged_dim=3, device="cpu") + + assert learner.obs_dim == 98 + assert learner.critic_obs_dim == 101 + assert learner.action_dim == 29 + + +def test_flashsac_actor_explore_and_forward_shapes(): + learner = FlashSACLearner(obs_dim=98, action_dim=29, privileged_dim=3, device="cpu") + obs = torch.randn(4, 98) + + actions = learner.actor.explore(obs, deterministic=False) + deterministic_actions = learner.actor.explore(obs, deterministic=True) + sampled_actions, info = learner.actor(obs, training=True) + + assert actions.shape == (4, 29) + assert deterministic_actions.shape == (4, 29) + assert sampled_actions.shape == (4, 29) + assert info["log_prob"].shape == (4,) + + +def test_flashsac_update_steps_run_on_cpu(): + learner = FlashSACLearner(obs_dim=98, action_dim=29, privileged_dim=3, device="cpu") + batch = _make_batch() + + critic_metrics = learner.update_critic(batch) + actor_metrics = learner.update_actor(batch) + learner.soft_update_target() + + assert "critic_loss" in critic_metrics + assert "reward_scale_std" in critic_metrics + assert "actor_loss" in actor_metrics + assert "temperature" in actor_metrics + + +def test_flashsac_state_dict_round_trip(): + learner = FlashSACLearner(obs_dim=98, action_dim=29, privileged_dim=3, device="cpu") + batch = _make_batch() + learner.update_critic(batch) + learner.update_actor(batch) + state_dict = learner.get_state_dict() + + restored = FlashSACLearner(obs_dim=98, action_dim=29, privileged_dim=3, device="cpu") + restored.load_state_dict(state_dict) + + assert restored.get_state_dict()["update_count"] == learner.get_state_dict()["update_count"] diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index ae2e2cfa6..9d61a386b 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -36,6 +36,17 @@ def test_td3_config_defaults(): assert cfg.algo_params.weight_decay == 0.1 +def test_flashsac_config_defaults(): + from unilab.config.structured_configs import FlashSACAlgoParams, FlashSACConfig + + cfg = FlashSACConfig() + assert cfg.algo == "flashsac" + assert cfg.batch_size == 2048 + assert cfg.obs_normalization is False + assert isinstance(cfg.algo_params, FlashSACAlgoParams) + assert cfg.algo_params.normalize_reward is True + + def test_ppo_config_defaults(): from unilab.config.structured_configs import PPOConfig @@ -124,6 +135,23 @@ def test_offpolicy_go2_task_overrides(): assert cfg.training.task_name == "Go2JoystickFlatTerrain" +def test_offpolicy_flashsac_g1_task_overrides(): + from hydra import compose, initialize_config_dir + from hydra.core.global_hydra import GlobalHydra + + GlobalHydra.instance().clear() + with initialize_config_dir(config_dir=str(CONF_DIR / "offpolicy"), version_base="1.3"): + cfg = compose( + "config", + overrides=["algo=flashsac", "task=flashsac/g1_sac/mujoco"], + ) + assert cfg.algo.algo == "flashsac" + assert cfg.training.task_name == "G1WalkTaskMjSAC" + assert cfg.training.sim_backend == "mujoco" + assert cfg.algo.algo_params.actor_num_blocks == 2 + assert cfg.algo.algo_params.normalize_reward is True + + # --------------------------------------------------------------------------- # Hydra YAML loading — appo # --------------------------------------------------------------------------- diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index 947383059..11778caef 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -1113,6 +1113,25 @@ def test_offpolicy_td3_hydra_default_algo_log_name(): assert cfg.algo.load_run == "-1" +def test_offpolicy_flashsac_hydra_algo_log_name(): + cfg = _offpolicy_cfg(["algo=flashsac", "task=flashsac/g1_sac/mujoco"]) + assert cfg.algo.algo_log_name == "flash_sac" + assert cfg.algo.load_run == "-1" + + +def test_offpolicy_flashsac_rejects_multi_gpu(): + cfg = _offpolicy_cfg( + [ + "algo=flashsac", + "task=flashsac/g1_sac/mujoco", + "training.num_gpus=2", + ] + ) + + with pytest.raises(ValueError, match="FlashSAC does not support training.num_gpus > 1"): + _offpolicy().build_runner("flashsac", cfg) + + def test_train_rsl_rl_get_log_root_uses_algo_log_name(monkeypatch: pytest.MonkeyPatch): """Verify _get_log_root uses algo.algo_log_name (issue #168).""" mod = _train_rsl_rl(monkeypatch) diff --git a/tests/utils/test_algo_utils.py b/tests/utils/test_algo_utils.py index ad40021c4..f0110b8aa 100644 --- a/tests/utils/test_algo_utils.py +++ b/tests/utils/test_algo_utils.py @@ -186,6 +186,21 @@ def test_builds_td3_actor_for_multiple_envs(self) -> None: ) assert hasattr(actor, "forward") + def test_builds_flashsac_actor(self) -> None: + actor = build_actor( + algo_type="flashsac", + obs_dim=98, + action_dim=29, + actor_hidden_dim=128, + use_layer_norm=False, + device="cpu", + actor_num_blocks=2, + actor_noise_zeta_mu=2.0, + actor_noise_zeta_max=16, + ) + assert hasattr(actor, "forward") + assert hasattr(actor, "explore") + def test_raises_for_unknown_algo_type(self) -> None: with pytest.raises(ValueError, match="Unknown algo_type"): build_actor(