From a21e240112af22e533301bf663db83f0a52cabd2 Mon Sep 17 00:00:00 2001 From: mio Date: Sun, 24 May 2026 05:00:37 +0800 Subject: [PATCH 1/3] perf: optimize fastsac compile --- conf/offpolicy/algo/sac.yaml | 1 + scripts/train_offpolicy.py | 1 + src/unilab/algos/torch/fast_sac/learner.py | 129 ++++++++++++------ src/unilab/structured_configs.py | 1 + tests/algos/test_fast_sac_compile.py | 44 ++++++ .../test_offpolicy_double_buffer_runner.py | 55 ++++++++ tests/config/test_locomotion_params.py | 2 + 7 files changed, 193 insertions(+), 40 deletions(-) create mode 100644 tests/algos/test_fast_sac_compile.py diff --git a/conf/offpolicy/algo/sac.yaml b/conf/offpolicy/algo/sac.yaml index bd7cee8bd..7e679207d 100644 --- a/conf/offpolicy/algo/sac.yaml +++ b/conf/offpolicy/algo/sac.yaml @@ -27,3 +27,4 @@ algo_params: alpha_init: 0.01 target_entropy_ratio: 0.0 max_grad_norm: 0.0 + use_compile: false diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index fc0138037..399268725 100644 --- a/scripts/train_offpolicy.py +++ b/scripts/train_offpolicy.py @@ -159,6 +159,7 @@ def build_runner(algo_name: str, cfg: DictConfig): use_layer_norm=cfg.algo.use_layer_norm, max_grad_norm=cfg.algo.algo_params.max_grad_norm, use_amp=cfg.training.use_amp, + use_compile=cfg.algo.algo_params.use_compile, use_symmetry=cfg.algo.use_symmetry, symmetry_augmentation=_symmetry_aug, critic_obs_dim=_critic_dim, diff --git a/src/unilab/algos/torch/fast_sac/learner.py b/src/unilab/algos/torch/fast_sac/learner.py index c0434328f..203b7a87f 100644 --- a/src/unilab/algos/torch/fast_sac/learner.py +++ b/src/unilab/algos/torch/fast_sac/learner.py @@ -393,6 +393,7 @@ def __init__( use_autotune: bool = True, use_symmetry: bool = False, use_amp: bool = False, + use_compile: bool = False, symmetry_augmentation: SymmetryAugmentation | None = None, world_size: int = 1, ): @@ -402,6 +403,7 @@ def __init__( self.max_grad_norm = max_grad_norm self.use_autotune = use_autotune self.use_amp = use_amp and device in ("cuda", "xpu") + self.use_compile = bool(use_compile) # CUDA uses fp16 + GradScaler; XPU uses bf16 (fp32 range, no scaler needed). self._amp_dtype = torch.bfloat16 if device == "xpu" else torch.float16 self.world_size = world_size @@ -491,6 +493,23 @@ def __init__( "FastSACLearner use_symmetry=True requires a symmetry_augmentation contract" ) self.use_symmetry = use_symmetry + if self.use_compile: + self._compile_training_methods() + + def _compile_training_methods(self) -> None: + compile_fn = getattr(torch, "compile", None) + if compile_fn is None: + raise ValueError("FastSAC compile requires torch.compile support") + if torch.device(self.device).type != "cuda": + raise ValueError("FastSAC compile currently requires a CUDA device") + + compile_kwargs = {"options": {"triton.cudagraphs": False}} + self._critic_loss_tensors = compile_fn( # type: ignore[method-assign] + self._critic_loss_tensors, **compile_kwargs + ) + self._actor_loss_tensors = compile_fn( # type: ignore[method-assign] + self._actor_loss_tensors, **compile_kwargs + ) def _autocast(self): return torch.amp.autocast( # pyright: ignore[reportPrivateImportUsage] @@ -518,6 +537,63 @@ def _reduce_gradients(self, model: nn.Module) -> None: p.grad.copy_(flat[offset : offset + n].view_as(p.grad)) offset += n + def _critic_loss_tensors( + self, + critic_obs: torch.Tensor, + actions: torch.Tensor, + rewards: torch.Tensor, + next_obs: torch.Tensor, + critic_next_obs: torch.Tensor, + dones: torch.Tensor, + truncated: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + bootstrap = torch.clamp(1.0 - dones.float() + truncated.float(), 0.0, 1.0) + discount = torch.full_like(dones, self.gamma) + + with torch.no_grad(): + with self._autocast(): + next_actions, next_log_probs, _ = self.actor.get_actions_and_log_probs(next_obs) + adjusted_rewards = ( + rewards - discount * bootstrap * self.log_alpha.exp() * next_log_probs + ) + + with self._autocast(): + target_distributions = self.qnet_target.projection( + critic_next_obs, next_actions, adjusted_rewards, bootstrap, discount + ) + target_values = self.qnet_target.get_value(target_distributions) + target_q_max = target_values.max() + target_q_min = target_values.min() + + with self._autocast(): + q_outputs = self.qnet(critic_obs, actions) + critic_log_probs = F.log_softmax(q_outputs, dim=-1).clamp(min=-30.0) + critic_losses = -torch.sum(target_distributions * critic_log_probs, dim=-1) + qf_loss = critic_losses.mean(dim=1).sum(dim=0) + + return qf_loss, target_q_max, target_q_min, next_log_probs.detach() + + def _actor_loss_tensors( + self, + obs: torch.Tensor, + critic_obs: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + with self._autocast(): + actions, log_probs, log_std = self.actor.get_actions_and_log_probs(obs) + + with torch.no_grad(): + action_std = log_std.exp().mean() + policy_entropy = -log_probs.mean() + + with self._autocast(): + q_outputs = self.qnet(critic_obs, actions) + q_probs = F.softmax(q_outputs, dim=-1) + q_values = self.qnet.get_value(q_probs) + qf_value = q_values.mean(dim=0) + actor_loss = (self.log_alpha.exp().detach() * log_probs - qf_value).mean() + + return actor_loss, policy_entropy, action_std + def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: """One critic update step.""" obs = batch["obs"] @@ -555,28 +631,15 @@ def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: dones = dones.repeat(2) truncated = truncated.repeat(2) - bootstrap = torch.clamp(1.0 - dones.float() + truncated.float(), 0.0, 1.0) - discount = torch.full_like(dones, self.gamma) - - with torch.no_grad(): - with self._autocast(): - next_actions, next_log_probs, _ = self.actor.get_actions_and_log_probs(next_obs) - adjusted_rewards = ( - rewards - discount * bootstrap * self.log_alpha.exp() * next_log_probs - ) - - with self._autocast(): - target_distributions = self.qnet_target.projection( - critic_next_obs, next_actions, adjusted_rewards, bootstrap, discount - ) - target_values = self.qnet_target.get_value(target_distributions) - - # Critic loss: cross-entropy with projected distributions - with self._autocast(): - q_outputs = self.qnet(critic_obs, actions) - critic_log_probs = F.log_softmax(q_outputs, dim=-1).clamp(min=-30.0) - critic_losses = -torch.sum(target_distributions * critic_log_probs, dim=-1) - qf_loss = critic_losses.mean(dim=1).sum(dim=0) + qf_loss, target_q_max, target_q_min, next_log_probs = self._critic_loss_tensors( + critic_obs, + actions, + rewards, + next_obs, + critic_next_obs, + dones, + truncated, + ) # Skip if NaN if torch.isfinite(qf_loss): @@ -610,10 +673,8 @@ def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: alpha_loss = torch.tensor(0.0, device=self.device) if self.use_autotune: self.alpha_optimizer.zero_grad(set_to_none=True) - # using next_log_probs like holosoma - # holosoma: alpha_loss = (-self.log_alpha.exp() * (next_state_log_probs.detach() + self.target_entropy)).mean() alpha_loss = ( - -self.log_alpha.exp() * (next_log_probs.detach() + self.target_entropy) + -self.log_alpha.exp() * (next_log_probs + self.target_entropy) ).mean() if torch.isfinite(alpha_loss): alpha_loss.backward() @@ -625,8 +686,8 @@ def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: return { "qf_loss": qf_loss.item(), "critic_grad_norm": critic_grad_norm.item(), - "target_q_max": target_values.max().item(), - "target_q_min": target_values.min().item(), + "target_q_max": target_q_max.item(), + "target_q_min": target_q_min.item(), "alpha_loss": alpha_loss.item(), "alpha": self.log_alpha.exp().item(), } @@ -645,19 +706,7 @@ def update_actor(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: dim=0, ) - with self._autocast(): - actions, log_probs, log_std = self.actor.get_actions_and_log_probs(obs) - - with torch.no_grad(): - action_std = log_std.exp().mean() - policy_entropy = -log_probs.mean() - - with self._autocast(): - q_outputs = self.qnet(critic_obs, actions) - q_probs = F.softmax(q_outputs, dim=-1) - q_values = self.qnet.get_value(q_probs) - qf_value = q_values.mean(dim=0) - actor_loss = (self.log_alpha.exp().detach() * log_probs - qf_value).mean() + actor_loss, policy_entropy, action_std = self._actor_loss_tensors(obs, critic_obs) # Skip if NaN if torch.isfinite(actor_loss): diff --git a/src/unilab/structured_configs.py b/src/unilab/structured_configs.py index dc4b3cfe9..dfe2387aa 100644 --- a/src/unilab/structured_configs.py +++ b/src/unilab/structured_configs.py @@ -25,6 +25,7 @@ class SACAlgoParams: alpha_init: float = 0.01 target_entropy_ratio: float = 0.0 max_grad_norm: float = 0.0 + use_compile: bool = False @dataclass diff --git a/tests/algos/test_fast_sac_compile.py b/tests/algos/test_fast_sac_compile.py new file mode 100644 index 000000000..00d24b299 --- /dev/null +++ b/tests/algos/test_fast_sac_compile.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import torch + +from unilab.algos.torch.fast_sac.learner import FastSACLearner + + +def test_fast_sac_compile_targets_training_hot_paths(monkeypatch) -> None: + calls: list[tuple[str, dict[str, Any]]] = [] + + def fake_compile(fn: Callable, **kwargs): + calls.append((fn.__qualname__, kwargs)) + return fn + + learner = FastSACLearner( + obs_dim=4, + action_dim=2, + critic_obs_dim=5, + device="cpu", + actor_hidden_dim=8, + critic_hidden_dim=8, + num_atoms=3, + num_q_networks=2, + use_layer_norm=False, + use_autotune=False, + ) + learner.device = "cuda" + monkeypatch.setattr(torch, "compile", fake_compile) + + learner._compile_training_methods() + + assert calls == [ + ( + "FastSACLearner._critic_loss_tensors", + {"options": {"triton.cudagraphs": False}}, + ), + ( + "FastSACLearner._actor_loss_tensors", + {"options": {"triton.cudagraphs": False}}, + ), + ] diff --git a/tests/algos/test_offpolicy_double_buffer_runner.py b/tests/algos/test_offpolicy_double_buffer_runner.py index 77be28bba..0cf05069d 100644 --- a/tests/algos/test_offpolicy_double_buffer_runner.py +++ b/tests/algos/test_offpolicy_double_buffer_runner.py @@ -181,6 +181,61 @@ def __init__(self, *args, **kwargs): assert isinstance(runner, _FakeRunner) assert runner.kwargs["device"] == device assert runner.kwargs["replay_prefetch_mode"] == "one_tick" + assert runner.kwargs["learner"].kwargs["use_compile"] is False + + +def test_sac_compile_override_is_passed_to_learner(monkeypatch: pytest.MonkeyPatch): + import gymnasium as gym + + mod = _offpolicy() + cfg = _offpolicy_cfg( + [ + "algo=sac", + "training.device=cuda", + "algo.use_symmetry=false", + "algo.algo_params.use_compile=true", + ] + ) + + class _FakeEnv: + obs_groups_spec = {"obs": 4, "critic": 6} + action_space = gym.spaces.Box(-1.0, 1.0, shape=(2,)) + + def build_symmetry_augmentation(self, device=None): + return None + + def close(self): + pass + + class _FakeLearner: + class actor: + @staticmethod + def state_dict(): + return {"w": MagicMock(shape=(4,))} + + update_count = 0 + + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + class _FakeRunner: + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + monkeypatch.setattr(mod, "ensure_registries", lambda: None) + monkeypatch.setattr(mod, "create_env", lambda *args, **kwargs: _FakeEnv()) + + import unilab.algos.torch.fast_sac.learner as learner_mod + + monkeypatch.setattr(learner_mod, "FastSACLearner", _FakeLearner) + + import unilab.algos.torch.offpolicy.double_buffer_runner as db_mod + + monkeypatch.setattr(db_mod, "DoubleBufferOffPolicyRunner", _FakeRunner) + + runner = mod.build_runner("sac", cfg) + + assert runner.kwargs["learner"].kwargs["use_compile"] is True def test_sac_async_collection_rejects_cpu_pinned_double_buffer(): diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index a0f6d6d94..7c1580175 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -57,6 +57,7 @@ def test_sac_config_defaults(): assert cfg.use_symmetry is False assert isinstance(cfg.algo_params, SACAlgoParams) assert cfg.algo_params.alpha_init == 0.01 + assert cfg.algo_params.use_compile is False def test_td3_config_defaults(): @@ -142,6 +143,7 @@ def test_offpolicy_sac_g1_task_overrides(): assert cfg.algo.max_iterations == 5000 assert cfg.algo.use_symmetry is True assert cfg.algo.algo_params.target_entropy_ratio == pytest.approx(0.0) + assert cfg.algo.algo_params.use_compile is False assert cfg.training.task_name == "G1WalkFlat" assert cfg.env.control_config.action_scale == pytest.approx(1.0) From 810b77efde26f41efb88c8c6900480479d6a5840 Mon Sep 17 00:00:00 2001 From: mio Date: Sun, 24 May 2026 12:03:57 +0800 Subject: [PATCH 2/3] perf: fast sac amp_dtype configure chain --- conf/offpolicy/algo/sac.yaml | 1 + scripts/train_offpolicy.py | 1 + src/unilab/algos/torch/fast_sac/learner.py | 35 +++++++++++++++++----- src/unilab/algos/torch/fast_sac/runner.py | 2 ++ src/unilab/structured_configs.py | 1 + tests/algos/test_fast_sac_compile.py | 16 ++++++++++ 6 files changed, 48 insertions(+), 8 deletions(-) diff --git a/conf/offpolicy/algo/sac.yaml b/conf/offpolicy/algo/sac.yaml index 7e679207d..cf334e5e2 100644 --- a/conf/offpolicy/algo/sac.yaml +++ b/conf/offpolicy/algo/sac.yaml @@ -27,4 +27,5 @@ algo_params: alpha_init: 0.01 target_entropy_ratio: 0.0 max_grad_norm: 0.0 + amp_dtype: auto use_compile: false diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index 399268725..53389086a 100644 --- a/scripts/train_offpolicy.py +++ b/scripts/train_offpolicy.py @@ -159,6 +159,7 @@ def build_runner(algo_name: str, cfg: DictConfig): use_layer_norm=cfg.algo.use_layer_norm, max_grad_norm=cfg.algo.algo_params.max_grad_norm, use_amp=cfg.training.use_amp, + amp_dtype=cfg.algo.algo_params.amp_dtype, use_compile=cfg.algo.algo_params.use_compile, use_symmetry=cfg.algo.use_symmetry, symmetry_augmentation=_symmetry_aug, diff --git a/src/unilab/algos/torch/fast_sac/learner.py b/src/unilab/algos/torch/fast_sac/learner.py index 203b7a87f..ff2af9214 100644 --- a/src/unilab/algos/torch/fast_sac/learner.py +++ b/src/unilab/algos/torch/fast_sac/learner.py @@ -393,19 +393,21 @@ def __init__( use_autotune: bool = True, use_symmetry: bool = False, use_amp: bool = False, + amp_dtype: str = "auto", use_compile: bool = False, symmetry_augmentation: SymmetryAugmentation | None = None, world_size: int = 1, ): self.device = device + self._device_type = torch.device(device).type self.gamma = gamma self.tau = tau self.max_grad_norm = max_grad_norm self.use_autotune = use_autotune - self.use_amp = use_amp and device in ("cuda", "xpu") + self.use_amp = bool(use_amp) and self._device_type in ("cuda", "xpu") self.use_compile = bool(use_compile) - # CUDA uses fp16 + GradScaler; XPU uses bf16 (fp32 range, no scaler needed). - self._amp_dtype = torch.bfloat16 if device == "xpu" else torch.float16 + self.amp_dtype = amp_dtype + self._amp_dtype = self._resolve_amp_dtype(amp_dtype, self._device_type) self.world_size = world_size self.critic_obs_dim = critic_obs_dim @@ -483,7 +485,7 @@ def __init__( # AMP scaler for mixed precision (fp16 only; bf16 has fp32 range and skips scaler) self.scaler = ( torch.amp.GradScaler("cuda") # pyright: ignore[reportPrivateImportUsage] - if self.use_amp and device == "cuda" + if self._should_use_grad_scaler(self.use_amp, self._device_type, self._amp_dtype) else None ) @@ -496,6 +498,25 @@ def __init__( if self.use_compile: self._compile_training_methods() + @staticmethod + def _resolve_amp_dtype(amp_dtype: str, device_type: str) -> torch.dtype: + normalized = amp_dtype.lower() + if normalized == "auto": + return torch.bfloat16 if device_type == "xpu" else torch.float16 + if normalized == "fp16": + return torch.float16 + if normalized == "bf16": + return torch.bfloat16 + raise ValueError("FastSAC amp_dtype must be one of: auto, fp16, bf16") + + @staticmethod + def _should_use_grad_scaler( + use_amp: bool, + device_type: str, + amp_dtype: torch.dtype, + ) -> bool: + return bool(use_amp) and device_type == "cuda" and amp_dtype == torch.float16 + def _compile_training_methods(self) -> None: compile_fn = getattr(torch, "compile", None) if compile_fn is None: @@ -513,7 +534,7 @@ def _compile_training_methods(self) -> None: def _autocast(self): return torch.amp.autocast( # pyright: ignore[reportPrivateImportUsage] - self.device, dtype=self._amp_dtype, enabled=self.use_amp + self._device_type, dtype=self._amp_dtype, enabled=self.use_amp ) def _reduce_gradients(self, model: nn.Module) -> None: @@ -673,9 +694,7 @@ def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: alpha_loss = torch.tensor(0.0, device=self.device) if self.use_autotune: self.alpha_optimizer.zero_grad(set_to_none=True) - alpha_loss = ( - -self.log_alpha.exp() * (next_log_probs + self.target_entropy) - ).mean() + alpha_loss = (-self.log_alpha.exp() * (next_log_probs + self.target_entropy)).mean() if torch.isfinite(alpha_loss): alpha_loss.backward() if self.world_size > 1 and self.log_alpha.grad is not None: diff --git a/src/unilab/algos/torch/fast_sac/runner.py b/src/unilab/algos/torch/fast_sac/runner.py index 5fe70d0dc..7c1d8a0ec 100644 --- a/src/unilab/algos/torch/fast_sac/runner.py +++ b/src/unilab/algos/torch/fast_sac/runner.py @@ -38,6 +38,7 @@ def __init__( use_layer_norm: bool = True, max_grad_norm: float = 0.0, use_amp: bool = False, + amp_dtype: str = "auto", sim_backend: str = "mujoco", use_symmetry: bool = False, world_size: int = 1, @@ -90,6 +91,7 @@ def __init__( use_layer_norm=use_layer_norm, max_grad_norm=max_grad_norm, use_amp=use_amp, + amp_dtype=amp_dtype, use_symmetry=use_symmetry, symmetry_augmentation=symmetry_augmentation, world_size=getattr(self, "world_size", world_size), diff --git a/src/unilab/structured_configs.py b/src/unilab/structured_configs.py index dfe2387aa..44813efc2 100644 --- a/src/unilab/structured_configs.py +++ b/src/unilab/structured_configs.py @@ -25,6 +25,7 @@ class SACAlgoParams: alpha_init: float = 0.01 target_entropy_ratio: float = 0.0 max_grad_norm: float = 0.0 + amp_dtype: str = "auto" use_compile: bool = False diff --git a/tests/algos/test_fast_sac_compile.py b/tests/algos/test_fast_sac_compile.py index 00d24b299..acaf20512 100644 --- a/tests/algos/test_fast_sac_compile.py +++ b/tests/algos/test_fast_sac_compile.py @@ -3,6 +3,7 @@ from collections.abc import Callable from typing import Any +import pytest import torch from unilab.algos.torch.fast_sac.learner import FastSACLearner @@ -42,3 +43,18 @@ def fake_compile(fn: Callable, **kwargs): {"options": {"triton.cudagraphs": False}}, ), ] + + +def test_fast_sac_amp_dtype_resolution_and_scaler_rules() -> None: + assert FastSACLearner._resolve_amp_dtype("auto", "cuda") is torch.float16 + assert FastSACLearner._resolve_amp_dtype("auto", "xpu") is torch.bfloat16 + assert FastSACLearner._resolve_amp_dtype("fp16", "cuda") is torch.float16 + assert FastSACLearner._resolve_amp_dtype("bf16", "cuda") is torch.bfloat16 + + assert FastSACLearner._should_use_grad_scaler(True, "cuda", torch.float16) + assert not FastSACLearner._should_use_grad_scaler(True, "cuda", torch.bfloat16) + assert not FastSACLearner._should_use_grad_scaler(True, "xpu", torch.bfloat16) + assert not FastSACLearner._should_use_grad_scaler(False, "cuda", torch.float16) + + with pytest.raises(ValueError, match="amp_dtype"): + FastSACLearner._resolve_amp_dtype("tf32", "cuda") From 2420695195457fb15f3e2669cc476461657b762c Mon Sep 17 00:00:00 2001 From: mio Date: Sun, 24 May 2026 12:34:43 +0800 Subject: [PATCH 3/3] perf: flashsac add compile --- conf/offpolicy/algo/flashsac.yaml | 1 + .../algos/torch/flash_sac/double_buffer.py | 5 +- src/unilab/algos/torch/flash_sac/learner.py | 135 ++++++++++++++---- src/unilab/algos/torch/flash_sac/runner.py | 2 + src/unilab/structured_configs.py | 1 + tests/algos/test_flash_sac_learner.py | 48 +++++++ .../test_offpolicy_double_buffer_runner.py | 66 ++++++++- tests/config/test_locomotion_params.py | 2 + 8 files changed, 223 insertions(+), 37 deletions(-) diff --git a/conf/offpolicy/algo/flashsac.yaml b/conf/offpolicy/algo/flashsac.yaml index 222bdf62c..d293d4dd1 100644 --- a/conf/offpolicy/algo/flashsac.yaml +++ b/conf/offpolicy/algo/flashsac.yaml @@ -40,4 +40,5 @@ algo_params: learning_rate_warmup_steps: 0 learning_rate_decay_steps: 500000 n_step: 1 + amp_dtype: auto use_compile: false diff --git a/src/unilab/algos/torch/flash_sac/double_buffer.py b/src/unilab/algos/torch/flash_sac/double_buffer.py index 921e241a7..ee628e860 100644 --- a/src/unilab/algos/torch/flash_sac/double_buffer.py +++ b/src/unilab/algos/torch/flash_sac/double_buffer.py @@ -30,10 +30,6 @@ def _validate_flashsac_double_buffer_runtime( ) if cfg.algo.algo_params.n_step != 1: raise ValueError("FlashSAC-B initially supports n_step=1 only") - if cfg.algo.algo_params.use_compile: - raise ValueError("FlashSAC-B initially requires algo.algo_params.use_compile=false") - if cfg.training.use_amp: - raise ValueError("FlashSAC-B initially requires training.use_amp=false") def build_flashsac_double_buffer_runner( @@ -96,6 +92,7 @@ def build_flashsac_double_buffer_runner( n_step=cfg.algo.algo_params.n_step, obs_normalization=cfg.algo.obs_normalization, use_amp=cfg.training.use_amp, + amp_dtype=cfg.algo.algo_params.amp_dtype, use_compile=cfg.algo.algo_params.use_compile, ) diff --git a/src/unilab/algos/torch/flash_sac/learner.py b/src/unilab/algos/torch/flash_sac/learner.py index 370233d7a..90525d41a 100644 --- a/src/unilab/algos/torch/flash_sac/learner.py +++ b/src/unilab/algos/torch/flash_sac/learner.py @@ -18,7 +18,6 @@ ) from unilab.algos.torch.flash_sac.update import ( build_lr_lambda, - compute_categorical_td_target, resolve_target_entropy, select_min_q_log_probs, ) @@ -167,6 +166,7 @@ def __init__( n_step: int = 1, obs_normalization: bool = False, use_amp: bool = False, + amp_dtype: str = "auto", use_compile: bool = False, ): self.device = torch.device(device) @@ -179,8 +179,8 @@ def __init__( self.action_dim = action_dim self.update_count = 0 self.use_amp = bool(use_amp and self.device.type in ("cuda", "xpu")) - # CUDA uses fp16 + GradScaler; XPU uses bf16 (fp32 range, no scaler needed). - self._amp_dtype = torch.bfloat16 if self.device.type == "xpu" else torch.float16 + self.amp_dtype = amp_dtype + self._amp_dtype = self._resolve_amp_dtype(amp_dtype, self.device.type) self.use_compile = bool( use_compile and hasattr(torch, "compile") and self.device.type == "cuda" ) @@ -228,7 +228,7 @@ def __init__( # GradScaler is only needed for fp16 (cuda); bf16 on xpu doesn't need it. self.scaler: Any | None = ( getattr(torch.amp, "GradScaler")("cuda") - if self.use_amp and self.device.type == "cuda" + if self._should_use_grad_scaler(self.use_amp, self.device.type, self._amp_dtype) else None ) lr_peak = learning_rate_peak if learning_rate_peak > 0 else actor_lr @@ -253,9 +253,42 @@ def __init__( ) 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" - ) + self._compile_training_methods() + + def _compile_training_methods(self) -> None: + compile_fn = getattr(torch, "compile", None) + if compile_fn is None or self.device.type != "cuda": + return + + compile_kwargs = {"options": {"triton.cudagraphs": False}} + self.actor.get_mean_and_std = compile_fn( # type: ignore[method-assign] + self.actor.get_mean_and_std, **compile_kwargs + ) + self._critic_loss_tensors = compile_fn( # type: ignore[method-assign] + self._critic_loss_tensors, **compile_kwargs + ) + self._actor_loss_tensors = compile_fn( # type: ignore[method-assign] + self._actor_loss_tensors, **compile_kwargs + ) + + @staticmethod + def _resolve_amp_dtype(amp_dtype: str, device_type: str) -> torch.dtype: + normalized = amp_dtype.lower() + if normalized == "auto": + return torch.bfloat16 if device_type == "xpu" else torch.float16 + if normalized == "fp16": + return torch.float16 + if normalized == "bf16": + return torch.bfloat16 + raise ValueError("FlashSAC amp_dtype must be one of: auto, fp16, bf16") + + @staticmethod + def _should_use_grad_scaler( + use_amp: bool, + device_type: str, + amp_dtype: torch.dtype, + ) -> bool: + return bool(use_amp) and device_type == "cuda" and amp_dtype == torch.float16 def _maybe_normalize_obs(self, obs: torch.Tensor, *, update: bool) -> torch.Tensor: if isinstance(self.obs_normalizer, nn.Identity): @@ -281,6 +314,60 @@ def _set_requires_grad(module: nn.Module, requires_grad: bool) -> None: for param in module.parameters(): param.requires_grad_(requires_grad) + def _critic_loss_tensors( + self, + next_q_values: torch.Tensor, + next_q_log_probs_full: torch.Tensor, + support: torch.Tensor, + rewards: torch.Tensor, + dones: torch.Tensor, + truncated: torch.Tensor, + actor_entropy: torch.Tensor, + pred_log_probs: torch.Tensor, + gamma: float, + ) -> torch.Tensor: + next_q_log_probs = select_min_q_log_probs(next_q_values, next_q_log_probs_full) + batch_size, num_bins = next_q_log_probs.shape + support_view = support.view(1, -1) + rewards = rewards.view(-1, 1) + dones = dones.view(-1, 1) + truncated = truncated.view(-1, 1) + actor_entropy = actor_entropy.view(-1, 1) + + bootstrap = torch.clamp(1.0 - dones + truncated, 0.0, 1.0) + support_min = support_view.min() + support_max = support_view.max() + target_bin_values = rewards + bootstrap * gamma * (support_view - actor_entropy) + target_bin_values = torch.clamp(target_bin_values, support_min, support_max) + + bin_width = torch.clamp(support_view[0, 1] - support_view[0, 0], min=1e-8) + offsets = (target_bin_values - support_min) / bin_width + 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 = next_q_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 cast(torch.Tensor, -(target_probs.unsqueeze(0) * pred_log_probs).sum(dim=-1).mean()) + + def _actor_loss_tensors( + self, + log_probs: torch.Tensor, + q_values: torch.Tensor, + actions: torch.Tensor, + expert_actions: torch.Tensor, + temp_value: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + min_q = torch.min(q_values[0], q_values[1]) + actor_loss = (temp_value.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 + entropy = -log_probs.detach().mean() + return actor_loss, entropy + def update_critic(self, batch: dict[str, torch.Tensor]) -> dict[str, float]: obs = batch["obs"].to(self.device) actions = batch["actions"].to(self.device) @@ -309,22 +396,22 @@ def update_critic(self, batch: dict[str, torch.Tensor]) -> dict[str, float]: qs_all, q_info_all = self.target_critic(obs_all, act_all, training=True) next_q_values = qs_all.chunk(2, dim=1)[1] next_q_log_probs_full = q_info_all["log_prob"].chunk(2, dim=1)[1] - next_q_log_probs = select_min_q_log_probs(next_q_values, next_q_log_probs_full) 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, - dones=dones, - truncated=truncated, - actor_entropy=actor_entropy, - gamma=gamma, - ) with self._autocast(): _, pred_info_all = self.critic(obs_all, act_all, training=True) pred_log_probs = pred_info_all["log_prob"].chunk(2, dim=1)[0] - critic_loss = -(target_probs.unsqueeze(0) * pred_log_probs).sum(dim=-1).mean() + critic_loss = self._critic_loss_tensors( + next_q_values, + next_q_log_probs_full, + support, + rewards, + dones, + truncated, + actor_entropy, + pred_log_probs, + gamma, + ) self.critic_optimizer.zero_grad(set_to_none=True) if self.scaler is not None: @@ -365,14 +452,9 @@ def update_actor(self, batch: dict[str, torch.Tensor]) -> dict[str, float]: 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 - ) + actor_loss, entropy = self._actor_loss_tensors( + log_probs, q_values, actions, expert_actions, self.temperature() + ) self.actor_optimizer.zero_grad(set_to_none=True) if self.scaler is not None: @@ -385,7 +467,6 @@ def update_actor(self, batch: dict[str, torch.Tensor]) -> dict[str, float]: 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) diff --git a/src/unilab/algos/torch/flash_sac/runner.py b/src/unilab/algos/torch/flash_sac/runner.py index 49006b4be..165257edf 100644 --- a/src/unilab/algos/torch/flash_sac/runner.py +++ b/src/unilab/algos/torch/flash_sac/runner.py @@ -51,6 +51,7 @@ def __init__( normalize_reward: bool = True, normalized_g_max: float = 5.0, n_step: int = 1, + amp_dtype: str = "auto", use_compile: bool = False, seed: int | None = None, trace_enabled: bool = False, @@ -107,6 +108,7 @@ def __init__( n_step=n_step, obs_normalization=obs_normalization, use_amp=use_amp, + amp_dtype=amp_dtype, use_compile=use_compile, ) diff --git a/src/unilab/structured_configs.py b/src/unilab/structured_configs.py index 44813efc2..754d83c71 100644 --- a/src/unilab/structured_configs.py +++ b/src/unilab/structured_configs.py @@ -121,6 +121,7 @@ class FlashSACAlgoParams: learning_rate_warmup_steps: int = 0 learning_rate_decay_steps: int = 500000 n_step: int = 1 + amp_dtype: str = "auto" use_compile: bool = False diff --git a/tests/algos/test_flash_sac_learner.py b/tests/algos/test_flash_sac_learner.py index c46772199..f5992fede 100644 --- a/tests/algos/test_flash_sac_learner.py +++ b/tests/algos/test_flash_sac_learner.py @@ -2,6 +2,10 @@ from __future__ import annotations +from collections.abc import Callable +from typing import Any + +import pytest import torch from unilab.algos.torch.flash_sac.learner import FlashSACLearner, RewardNormalizer @@ -37,6 +41,50 @@ def test_flashsac_learner_exposes_expected_dims(): assert learner.action_dim == 29 +def test_flashsac_compile_targets_training_hot_paths(monkeypatch) -> None: + calls: list[tuple[str, dict[str, Any]]] = [] + + def fake_compile(fn: Callable, **kwargs): + calls.append((fn.__qualname__, kwargs)) + return fn + + learner = FlashSACLearner(obs_dim=98, action_dim=29, critic_obs_dim=101, device="cpu") + learner.device = torch.device("cuda") + monkeypatch.setattr(torch, "compile", fake_compile) + + learner._compile_training_methods() + + assert calls == [ + ( + "FlashSACActor.get_mean_and_std", + {"options": {"triton.cudagraphs": False}}, + ), + ( + "FlashSACLearner._critic_loss_tensors", + {"options": {"triton.cudagraphs": False}}, + ), + ( + "FlashSACLearner._actor_loss_tensors", + {"options": {"triton.cudagraphs": False}}, + ), + ] + + +def test_flashsac_amp_dtype_resolution_and_scaler_rules() -> None: + assert FlashSACLearner._resolve_amp_dtype("auto", "cuda") is torch.float16 + assert FlashSACLearner._resolve_amp_dtype("auto", "xpu") is torch.bfloat16 + assert FlashSACLearner._resolve_amp_dtype("fp16", "cuda") is torch.float16 + assert FlashSACLearner._resolve_amp_dtype("bf16", "cuda") is torch.bfloat16 + + assert FlashSACLearner._should_use_grad_scaler(True, "cuda", torch.float16) + assert not FlashSACLearner._should_use_grad_scaler(True, "cuda", torch.bfloat16) + assert not FlashSACLearner._should_use_grad_scaler(True, "xpu", torch.bfloat16) + assert not FlashSACLearner._should_use_grad_scaler(False, "cuda", torch.float16) + + with pytest.raises(ValueError, match="amp_dtype"): + FlashSACLearner._resolve_amp_dtype("tf32", "cuda") + + def test_flashsac_actor_explore_and_forward_shapes(): learner = FlashSACLearner(obs_dim=98, action_dim=29, critic_obs_dim=101, device="cpu") obs = torch.randn(4, 98) diff --git a/tests/algos/test_offpolicy_double_buffer_runner.py b/tests/algos/test_offpolicy_double_buffer_runner.py index 0cf05069d..adc5ab171 100644 --- a/tests/algos/test_offpolicy_double_buffer_runner.py +++ b/tests/algos/test_offpolicy_double_buffer_runner.py @@ -328,7 +328,11 @@ def test_flashsac_double_buffer_n_step_rejected(): _offpolicy().build_runner("flashsac", cfg) -def test_flashsac_double_buffer_compile_rejected(): +def test_flashsac_double_buffer_compile_passed_to_learner(monkeypatch: pytest.MonkeyPatch): + import gymnasium as gym + + import unilab.algos.torch.flash_sac.double_buffer as flash_db_mod + cfg = _offpolicy_cfg( [ "algo=flashsac", @@ -336,20 +340,70 @@ def test_flashsac_double_buffer_compile_rejected(): "algo.algo_params.use_compile=true", ] ) - with pytest.raises(ValueError, match="use_compile=false"): - _offpolicy().build_runner("flashsac", cfg) + + class _FakeEnv: + obs_groups_spec = {"obs": 4, "critic": 6} + action_space = gym.spaces.Box(-1.0, 1.0, shape=(2,)) + + def close(self): + pass + + class _FakeLearner: + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + class _FakeRunner: + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + monkeypatch.setattr(flash_db_mod, "ensure_registries", lambda: None) + monkeypatch.setattr(flash_db_mod, "create_env", lambda *args, **kwargs: _FakeEnv()) + monkeypatch.setattr(flash_db_mod, "FlashSACLearner", _FakeLearner) + monkeypatch.setattr(flash_db_mod, "DoubleBufferOffPolicyRunner", _FakeRunner) + + runner = _offpolicy().build_runner("flashsac", cfg) + + assert runner.kwargs["learner"].kwargs["use_compile"] is True -def test_flashsac_double_buffer_amp_rejected(): +def test_flashsac_double_buffer_amp_passed_to_learner(monkeypatch: pytest.MonkeyPatch): + import gymnasium as gym + + import unilab.algos.torch.flash_sac.double_buffer as flash_db_mod + cfg = _offpolicy_cfg( [ "algo=flashsac", "training.device=cuda", "training.use_amp=true", + "algo.algo_params.amp_dtype=bf16", ] ) - with pytest.raises(ValueError, match="training.use_amp=false"): - _offpolicy().build_runner("flashsac", cfg) + + class _FakeEnv: + obs_groups_spec = {"obs": 4, "critic": 6} + action_space = gym.spaces.Box(-1.0, 1.0, shape=(2,)) + + def close(self): + pass + + class _FakeLearner: + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + class _FakeRunner: + def __init__(self, *args, **kwargs): + self.kwargs = kwargs + + monkeypatch.setattr(flash_db_mod, "ensure_registries", lambda: None) + monkeypatch.setattr(flash_db_mod, "create_env", lambda *args, **kwargs: _FakeEnv()) + monkeypatch.setattr(flash_db_mod, "FlashSACLearner", _FakeLearner) + monkeypatch.setattr(flash_db_mod, "DoubleBufferOffPolicyRunner", _FakeRunner) + + runner = _offpolicy().build_runner("flashsac", cfg) + + assert runner.kwargs["learner"].kwargs["use_amp"] is True + assert runner.kwargs["learner"].kwargs["amp_dtype"] == "bf16" @pytest.mark.parametrize( diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index 7c1580175..08df6ca1f 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -82,6 +82,7 @@ def test_flashsac_config_defaults(): assert cfg.obs_normalization is False assert isinstance(cfg.algo_params, FlashSACAlgoParams) assert cfg.algo_params.normalize_reward is True + assert cfg.algo_params.amp_dtype == "auto" assert cfg.algo_params.use_compile is False @@ -194,6 +195,7 @@ def test_offpolicy_flashsac_g1_task_overrides(): assert cfg.training.sim_backend == "mujoco" assert cfg.algo.algo_params.actor_num_blocks == 2 assert cfg.algo.algo_params.normalize_reward is True + assert cfg.algo.algo_params.amp_dtype == "auto" def test_offpolicy_flashsac_go2_task_overrides():