diff --git a/conf/offpolicy/algo/flashsac.yaml b/conf/offpolicy/algo/flashsac.yaml index 06c16f793..0de53ad67 100644 --- a/conf/offpolicy/algo/flashsac.yaml +++ b/conf/offpolicy/algo/flashsac.yaml @@ -3,16 +3,16 @@ algo: flashsac algo_log_name: flash_sac load_run: "-1" seed: 1 -num_envs: 2048 +num_envs: 1024 batch_size: 2048 replay_buffer_n: 512 -updates_per_step: 1 -warmup_steps: 10000 +updates_per_step: 2 +warmup_steps: 100000 policy_frequency: 2 env_steps_per_sync: 1 max_iterations: 5000 save_interval: 1000 -gamma: 0.99 +gamma: 0.97 tau: 0.01 actor_lr: 3.0e-4 critic_lr: 3.0e-4 diff --git a/conf/offpolicy/task/flashsac/g1_joystick/mujoco.yaml b/conf/offpolicy/task/flashsac/g1_joystick/mujoco.yaml new file mode 100644 index 000000000..3800d7d82 --- /dev/null +++ b/conf/offpolicy/task/flashsac/g1_joystick/mujoco.yaml @@ -0,0 +1,39 @@ +# @package _global_ +training: + task_name: G1JoystickFlatTerrain + sim_backend: mujoco + use_amp: true +algo: + num_envs: 1024 + replay_buffer_n: 9766 + warmup_steps: 100000 + updates_per_step: 2 + max_iterations: 48829 + save_interval: 5000 + gamma: 0.97 +env: + control_config: + action_scale: 0.5 + commands: + vel_limit: + - [-1.0, -0.5, -1.0] + - [1.0, 0.5, 1.0] +reward: + scales: + tracking_lin_vel: 1.0 + tracking_ang_vel: 0.75 + feet_phase: 1.0 + lin_vel_z: 0.0 + ang_vel_xy: -0.15 + base_height: 0.0 + orientation: -2.0 + action_rate: 0.0 + pose: -0.1 + tracking_sigma: 0.25 + gait_frequency: 1.5 + feet_phase_swing_height: 0.15 + feet_phase_tracking_sigma: 0.01 + base_height_target: 0.754 + min_base_height: 0.55 + max_tilt_deg: 25.0 + 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/docs/zh_CN/02-simulation-backends.md b/docs/zh_CN/02-simulation-backends.md index 23f274a74..3b0c96ff7 100644 --- a/docs/zh_CN/02-simulation-backends.md +++ b/docs/zh_CN/02-simulation-backends.md @@ -45,23 +45,25 @@ uv run python scripts/generate_support_matrix.py --write | PPO (torch) | `g1_joystick` (G1 joystick) | Tested | Tested | | PPO (torch) | `g1_motion_tracking` (G1 motion tracking) | Tested | Tested | | PPO (torch) | `g1_flip_tracking` (G1 flip tracking) | Tested | Tested | -| PPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | - | +| PPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | Tested | +| PPO (torch) | `allegro_inhand_grasp` (allegro inhand grasp) | Tested | Tested | | PPO (mlx) | `go1_joystick` (Go1 joystick) | Tested | Tested | | PPO (mlx) | `go2_joystick` (Go2 joystick) | Tested | Tested | | PPO (mlx) | `g1_joystick` (G1 joystick) | Tested | Tested | | PPO (mlx) | `g1_motion_tracking` (G1 motion tracking) | Configured | Configured | | PPO (mlx) | `g1_flip_tracking` (G1 flip tracking) | Configured | Configured | -| PPO (mlx) | `allegro_inhand` (Allegro in-hand) | Configured | - | +| PPO (mlx) | `allegro_inhand` (Allegro in-hand) | Configured | Configured | +| PPO (mlx) | `allegro_inhand_grasp` (allegro inhand grasp) | Configured | Configured | | APPO (torch) | `go1_joystick` (Go1 joystick) | Tested | Registered | | APPO (torch) | `go2_joystick` (Go2 joystick) | Tested | Registered | | APPO (torch) | `g1_joystick` (G1 joystick) | Tested | Registered | | APPO (torch) | `g1_motion_tracking` (G1 motion tracking) | Tested | Tested | | APPO (torch) | `g1_flip_tracking` (G1 flip tracking) | Tested | Tested | -| APPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | - | +| APPO (torch) | `allegro_inhand` (Allegro in-hand) | Tested | Registered | | SAC (torch) | `go1_joystick` (Go1 joystick) | Tested | Tested | | SAC (torch) | `go2_joystick` (Go2 joystick) | Tested | Tested | | SAC (torch) | `g1_sac` (G1 SAC locomotion) | Tested | Tested | -| SAC (torch) | `allegro_sac` (Allegro SAC in-hand) | Tested | - | +| SAC (torch) | `allegro_sac` (Allegro SAC in-hand) | - | - | | TD3 (torch) | `go1_joystick` (Go1 joystick) | Tested | Tested | | TD3 (torch) | `go2_joystick` (Go2 joystick) | Tested | Tested | diff --git a/docs/zh_CN/07-domain-randomization.md b/docs/zh_CN/07-domain-randomization.md index ae7e0b52b..71fce0f73 100644 --- a/docs/zh_CN/07-domain-randomization.md +++ b/docs/zh_CN/07-domain-randomization.md @@ -25,7 +25,7 @@ | `G1WalkTaskMjSAC` | 是 | 是：复用 [`G1JoystickDomainRandomizationProvider`](../../src/unilab/envs/locomotion/g1/joystick.py) | 任务状态采样 + common payload | push | [`g1/joystick_sac.py`](../../src/unilab/envs/locomotion/g1/joystick_sac.py) | | `G1MotionTracking` | 是 | 是：`Domain_Rand + Provider + ResetPlan` | 大量任务特有 reset 采样 + common payload | push | [`motion_tracking/g1/tracking.py`](../../src/unilab/envs/motion_tracking/g1/tracking.py) | | `AllegroInhandRotation` | 是 | 是：`DomainRandConfig + Provider + ResetPlan` | 纯任务特有 reset 采样，`randomization=None` | 无 | [`inhand_rot_allegro/rotation.py`](../../src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py) | -| `AllegroInhandRotationSac` | 是 | 是：`DomainRandConfigSac + Provider + ResetPlan` | 纯任务特有 reset 采样，`randomization=None` | 无 | [`inhand_rot_allegro/rotation_sac.py`](../../src/unilab/envs/manipulation/inhand_rot_allegro/rotation_sac.py) | +| `AllegroInhandRotationSac` | 是 | 是：`DomainRandConfigSac + Provider + ResetPlan` | 纯任务特有 reset 采样，`randomization=None` | 无 | [`inhand_rot_allegro/README.md`](../../src/unilab/envs/manipulation/inhand_rot_allegro/README.md) | ## 任务域随机化清单 @@ -100,7 +100,7 @@ backend capability 当前是： - [`go1/joystick.py`](../../src/unilab/envs/locomotion/go1/joystick.py) - [`go2/joystick.py`](../../src/unilab/envs/locomotion/go2/joystick.py) - [`inhand_rot_allegro/rotation.py`](../../src/unilab/envs/manipulation/inhand_rot_allegro/rotation.py) -- [`inhand_rot_allegro/rotation_sac.py`](../../src/unilab/envs/manipulation/inhand_rot_allegro/rotation_sac.py) +- [`inhand_rot_allegro/README.md`](../../src/unilab/envs/manipulation/inhand_rot_allegro/README.md) 因此从“当前任务 DR 状态”角度，`kp/kd` 还不能算各任务已经统一接入的域随机项。 diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index f9629ad25..3294e6ff1 100644 --- a/scripts/train_offpolicy.py +++ b/scripts/train_offpolicy.py @@ -16,6 +16,7 @@ from unilab.training import ( BackendAdapter, + assert_offpolicy_task_choice_matches_algo, create_env, ensure_registries, get_log_root, @@ -76,6 +77,7 @@ def extract_play_obs(obs_dict): def build_offpolicy_env_cfg_override(algo_name: str, cfg: DictConfig) -> dict[str, Any] | None: + assert_offpolicy_task_choice_matches_algo(cfg, algo_name=algo_name) return cast( dict[str, Any] | None, BackendAdapter(cfg, root_dir=ROOT_DIR, algo_name=algo_name).build_task_env_cfg_override(), @@ -444,6 +446,7 @@ def main(cfg: DictConfig) -> None: algo_name = cfg.algo.algo task_name = cfg.training.task_name + assert_offpolicy_task_choice_matches_algo(cfg, algo_name=algo_name) if cfg.training.log_dir is None: timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") diff --git a/src/unilab/algos/torch/flash_sac/learner.py b/src/unilab/algos/torch/flash_sac/learner.py index 36cdbfc66..edf3d4bdb 100644 --- a/src/unilab/algos/torch/flash_sac/learner.py +++ b/src/unilab/algos/torch/flash_sac/learner.py @@ -68,38 +68,72 @@ def load_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None: class RewardNormalizer: - """Adaptive reward scaling with running reward statistics.""" + """Adaptive reward scaling with running discounted-return statistics.""" - def __init__(self, g_max: float, device: torch.device, eps: float = 1e-8): + def __init__( + self, + gamma: float, + g_max: float, + device: torch.device, + eps: float = 1e-8, + ): + self.gamma = gamma 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) + self.g_r = torch.zeros(0, device=device, dtype=torch.float32) + self.g_r_max = torch.tensor(0.0, device=device, dtype=torch.float32) + + def _ensure_g_r_shape(self, num_envs: int) -> None: + if self.g_r.shape == (num_envs,): + return + self.g_r = torch.zeros(num_envs, device=self.device, dtype=torch.float32) - def update(self, rewards: torch.Tensor) -> None: - rewards = rewards.reshape(-1).to(device=self.device, dtype=torch.float32) + def update_from_transitions( + self, + rewards: torch.Tensor, + terminated: torch.Tensor, + truncated: torch.Tensor, + ) -> None: + rewards = rewards.to(device=self.device, dtype=torch.float32) + terminated = terminated.to(device=self.device, dtype=torch.float32) + truncated = truncated.to(device=self.device, dtype=torch.float32) + + if rewards.ndim == 1: + rewards = rewards.unsqueeze(0) + terminated = terminated.unsqueeze(0) + truncated = truncated.unsqueeze(0) if rewards.numel() == 0: return - self.rms.update(rewards) - self.reward_abs_max = torch.maximum(self.reward_abs_max, rewards.abs().max()) + + num_envs = int(rewards.shape[-1]) + self._ensure_g_r_shape(num_envs) + done = torch.clamp(terminated + truncated, min=0.0, max=1.0) + + for step in range(rewards.shape[0]): + self.g_r = self.gamma * (1.0 - done[step]) * self.g_r + rewards[step] + self.g_r_max = torch.maximum(self.g_r_max, self.g_r.abs().max()) + self.rms.update(self.g_r) 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), + self.g_r_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, + "g_r": self.g_r, + "g_r_max": self.g_r_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"] + self.g_r = state_dict["g_r"] + self.g_r_max = state_dict["g_r_max"] class FlashSACLearner: @@ -188,7 +222,7 @@ def __init__( self.obs_normalizer = nn.Identity() self.reward_normalizer = ( - RewardNormalizer(g_max=normalized_g_max, device=self.device) + RewardNormalizer(gamma=self.gamma, g_max=normalized_g_max, device=self.device) if normalize_reward else None ) @@ -237,6 +271,16 @@ def _build_critic_obs( def _autocast(self): return torch.autocast(device_type="cuda", dtype=torch.float16, enabled=self.use_amp) + def update_reward_stats( + self, + rewards: torch.Tensor, + terminated: torch.Tensor, + truncated: torch.Tensor, + ) -> None: + if self.reward_normalizer is None: + return + self.reward_normalizer.update_from_transitions(rewards, terminated, truncated) + @staticmethod def _set_requires_grad(module: nn.Module, requires_grad: bool) -> None: for param in module.parameters(): @@ -260,7 +304,6 @@ def update_critic(self, batch: dict[str, torch.Tensor]) -> dict[str, float]: 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 diff --git a/src/unilab/algos/torch/offpolicy/runner.py b/src/unilab/algos/torch/offpolicy/runner.py index 73d3b2145..8f0ff3fb7 100644 --- a/src/unilab/algos/torch/offpolicy/runner.py +++ b/src/unilab/algos/torch/offpolicy/runner.py @@ -5,6 +5,7 @@ import sys import time from collections import deque +from typing import cast import torch @@ -78,6 +79,57 @@ def _build_learner(self): def _collector_fn(self, stop_event, **kwargs): off_policy_collector_fn(stop_event=stop_event, **kwargs) + @staticmethod + def _read_recent_replay_field( + replay_buffer, field_name: str, start_ptr: int, count: int + ) -> torch.Tensor: + idx = start_ptr % replay_buffer.capacity + + if hasattr(replay_buffer, field_name): + source = getattr(replay_buffer, field_name) + else: + packed_key = { + "rewards": "_rew_col", + "dones": "_done_col", + "truncated": "_trunc_col", + }[field_name] + source = replay_buffer._storage[:, getattr(replay_buffer, packed_key)] + + if idx + count <= replay_buffer.capacity: + return cast(torch.Tensor, source[idx : idx + count].clone()) + + split = replay_buffer.capacity - idx + return cast(torch.Tensor, torch.cat([source[idx:], source[: count - split]], dim=0).clone()) + + def _update_reward_stats_from_replay(self, replay_buffer, start_ptr: int, end_ptr: int) -> int: + if not hasattr(self.learner, "update_reward_stats"): + return end_ptr + if getattr(self.learner, "reward_normalizer", None) is None: + return end_ptr + + count = end_ptr - start_ptr + if count <= 0: + return end_ptr + if count > replay_buffer.capacity: + count = replay_buffer.capacity + start_ptr = end_ptr - count + if count % self.num_envs != 0: + count -= count % self.num_envs + start_ptr = end_ptr - count + if count <= 0: + return end_ptr + + rewards = self._read_recent_replay_field(replay_buffer, "rewards", start_ptr, count) + dones = self._read_recent_replay_field(replay_buffer, "dones", start_ptr, count) + truncated = self._read_recent_replay_field(replay_buffer, "truncated", start_ptr, count) + num_steps = count // self.num_envs + self.learner.update_reward_stats( + rewards.view(num_steps, self.num_envs), + dones.view(num_steps, self.num_envs), + truncated.view(num_steps, self.num_envs), + ) + return end_ptr + def learn( self, max_iterations: int = 1500, @@ -180,6 +232,7 @@ def learn( latest_reward_components: dict[str, float] = {} last_buf_log = 0 write_read_ema = 0.0 + reward_stats_ptr = 0 # Training loop for iteration in range(1, max_iterations + 1): @@ -245,6 +298,11 @@ def learn( wait_time = time.time() - wait_start self._drain_metrics(metrics_queue, reward_history, latest_reward_components, logger) collect_time = time.time() - iter_start + reward_stats_ptr = self._update_reward_stats_from_replay( + replay_buffer, + reward_stats_ptr, + int(replay_buffer.ptr[0]), + ) train_start = time.time() from collections import defaultdict diff --git a/src/unilab/algos/torch/offpolicy/worker.py b/src/unilab/algos/torch/offpolicy/worker.py index 778c29b63..027f1d94e 100644 --- a/src/unilab/algos/torch/offpolicy/worker.py +++ b/src/unilab/algos/torch/offpolicy/worker.py @@ -6,6 +6,7 @@ import queue import sys +from typing import cast import numpy as np import torch @@ -39,6 +40,21 @@ def resolve_collector_actor_dims( return obs_dim, action_dim +def sample_offpolicy_actions( + actor, + algo_type: str, + obs_torch: torch.Tensor, + prev_dones_torch: torch.Tensor, +) -> torch.Tensor: + """Sample collector actions using the algorithm's exploration policy.""" + if algo_type in ("sac", "td3", "flashsac"): + return cast( + torch.Tensor, + actor.explore(obs_torch, dones=prev_dones_torch, deterministic=False), + ) + raise ValueError(f"Unsupported off-policy algo_type for collector action sampling: {algo_type}") + + def off_policy_collector_fn( stop_event, env_name: str, @@ -245,11 +261,13 @@ def _run_collector( else: _t_infer = _time.perf_counter() obs_torch = torch.from_numpy(obs_np_input) - if algo_type in ("sac", "td3"): - dones_torch = torch.from_numpy(prev_dones_np) - actions_torch = actor.explore(obs_torch, dones=dones_torch, deterministic=False) - else: - actions_torch = torch.zeros((num_envs, action_dim)) + dones_torch = torch.from_numpy(prev_dones_np) + actions_torch = sample_offpolicy_actions( + actor=actor, + algo_type=algo_type, + obs_torch=obs_torch, + prev_dones_torch=dones_torch, + ) actions_np = actions_torch.numpy() timing_accum_ms["mlp_infer_ms"] += (_time.perf_counter() - _t_infer) * 1000 diff --git a/src/unilab/config/structured_configs.py b/src/unilab/config/structured_configs.py index afd3fb9a3..4bfbf03a5 100644 --- a/src/unilab/config/structured_configs.py +++ b/src/unilab/config/structured_configs.py @@ -127,16 +127,16 @@ class FlashSACConfig(BaseConfig): algo: str = "flashsac" algo_log_name: str = "flash_sac" seed: int = 1 - num_envs: int = 2048 + num_envs: int = 1024 batch_size: int = 2048 replay_buffer_n: int = 512 - updates_per_step: int = 1 - warmup_steps: int = 10000 + updates_per_step: int = 2 + warmup_steps: int = 100000 policy_frequency: int = 2 env_steps_per_sync: int = 1 max_iterations: int = 5000 save_interval: int = 1000 - gamma: float = 0.99 + gamma: float = 0.97 tau: float = 0.01 actor_lr: float = 3e-4 critic_lr: float = 3e-4 diff --git a/src/unilab/training/__init__.py b/src/unilab/training/__init__.py index 5e7945c6c..1726b7fcc 100644 --- a/src/unilab/training/__init__.py +++ b/src/unilab/training/__init__.py @@ -2,9 +2,11 @@ from unilab.training.backend_adapter import BackendAdapter from unilab.training.common import ( + assert_offpolicy_task_choice_matches_algo, create_env, ensure_registries, get_entrypoint_log_root, + get_hydra_runtime_choice, get_latest_checkpoint, get_latest_run, get_log_root, @@ -17,9 +19,11 @@ __all__ = [ "BackendAdapter", + "assert_offpolicy_task_choice_matches_algo", "create_env", "ensure_registries", "get_entrypoint_log_root", + "get_hydra_runtime_choice", "get_latest_checkpoint", "get_latest_run", "get_log_root", diff --git a/src/unilab/training/common.py b/src/unilab/training/common.py index b3e98307e..55087b13b 100644 --- a/src/unilab/training/common.py +++ b/src/unilab/training/common.py @@ -8,6 +8,7 @@ from typing import Any, Callable, TypeVar import numpy as np +from hydra.core.hydra_config import HydraConfig from omegaconf import DictConfig, OmegaConf from unilab.utils.algo_utils import ensure_registries as _ensure_registries @@ -20,6 +21,51 @@ def ensure_registries() -> None: _ensure_registries() +def get_hydra_runtime_choice(cfg: DictConfig, group: str) -> str | None: + """Return a selected Hydra config-group choice when runtime metadata is available.""" + cfg_choice = OmegaConf.select(cfg, f"hydra.runtime.choices.{group}") + if cfg_choice is not None: + return str(cfg_choice) + + if not HydraConfig.initialized(): + return None + + try: + runtime_choice = HydraConfig.get().runtime.choices.get(group) + except Exception: + return None + return str(runtime_choice) if runtime_choice is not None else None + + +def assert_offpolicy_task_choice_matches_algo( + cfg: DictConfig, + *, + algo_name: str | None = None, +) -> None: + """Reject offpolicy configs whose task owner path does not match the selected algo.""" + cfg_algo_name = str(OmegaConf.select(cfg, "algo.algo")) + if algo_name is not None and cfg_algo_name != algo_name: + raise ValueError( + f"Off-policy algo argument {algo_name!r} is inconsistent with cfg.algo.algo={cfg_algo_name!r}" + ) + + selected_algo = algo_name or cfg_algo_name + task_choice = get_hydra_runtime_choice(cfg, "task") + if task_choice is None: + return + + task_algo, sep, _ = task_choice.partition("/") + if not sep: + raise ValueError( + f"Off-policy task choice must use task=//; got task={task_choice}" + ) + if task_algo != selected_algo: + raise ValueError( + f"Off-policy algo/task mismatch: algo={selected_algo} is inconsistent with task={task_choice}. " + "Use task=// with the same algo prefix." + ) + + def get_log_root(root_dir: str | Path, cfg: DictConfig) -> Path: """Resolve the algorithm log root, honoring optional training.log_root overrides.""" configured_root = OmegaConf.select(cfg, "training.log_root") diff --git a/tests/algos/test_flash_sac_learner.py b/tests/algos/test_flash_sac_learner.py index 392952526..1be258b5e 100644 --- a/tests/algos/test_flash_sac_learner.py +++ b/tests/algos/test_flash_sac_learner.py @@ -4,7 +4,7 @@ import torch -from unilab.algos.torch.flash_sac.learner import FlashSACLearner +from unilab.algos.torch.flash_sac.learner import FlashSACLearner, RewardNormalizer def _make_batch(batch_size: int = 32) -> dict[str, torch.Tensor]: @@ -75,3 +75,31 @@ def test_flashsac_state_dict_round_trip(): restored.load_state_dict(state_dict) assert restored.get_state_dict()["update_count"] == learner.get_state_dict()["update_count"] + + +def test_reward_normalizer_tracks_discounted_returns() -> None: + normalizer = RewardNormalizer(gamma=0.5, g_max=5.0, device=torch.device("cpu")) + + normalizer.update_from_transitions( + rewards=torch.tensor([[2.0, 1.0], [4.0, 3.0]]), + terminated=torch.tensor([[0.0, 1.0], [0.0, 0.0]]), + truncated=torch.zeros(2, 2), + ) + + torch.testing.assert_close(normalizer.g_r, torch.tensor([5.0, 3.5])) + torch.testing.assert_close(normalizer.g_r_max, torch.tensor(5.0)) + + +def test_flashsac_critic_update_does_not_advance_reward_stats_from_sampled_batch() -> None: + learner = FlashSACLearner(obs_dim=98, action_dim=29, privileged_dim=3, device="cpu") + learner.update_reward_stats( + rewards=torch.tensor([[1.0, 2.0]]), + terminated=torch.zeros(1, 2), + truncated=torch.zeros(1, 2), + ) + assert learner.reward_normalizer is not None + before = learner.reward_normalizer.g_r.clone() + + learner.update_critic(_make_batch()) + + torch.testing.assert_close(learner.reward_normalizer.g_r, before) diff --git a/tests/algos/test_offpolicy_worker.py b/tests/algos/test_offpolicy_worker.py new file mode 100644 index 000000000..c5d1d8867 --- /dev/null +++ b/tests/algos/test_offpolicy_worker.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +import pytest +import torch + +from unilab.algos.torch.offpolicy.worker import sample_offpolicy_actions + + +class _DummyActor: + def __init__(self) -> None: + self.calls: list[tuple[torch.Tensor, torch.Tensor, bool]] = [] + + def explore( + self, + obs: torch.Tensor, + dones: torch.Tensor | None = None, + deterministic: bool = False, + ) -> torch.Tensor: + assert dones is not None + self.calls.append((obs.clone(), dones.clone(), deterministic)) + return torch.ones(obs.shape[0], 3, dtype=obs.dtype) + + +@pytest.mark.parametrize("algo_type", ["sac", "td3", "flashsac"]) +def test_sample_offpolicy_actions_uses_actor_explore(algo_type: str) -> None: + actor = _DummyActor() + obs = torch.zeros(4, 5) + dones = torch.zeros(4) + + actions = sample_offpolicy_actions( + actor=actor, + algo_type=algo_type, + obs_torch=obs, + prev_dones_torch=dones, + ) + + assert len(actor.calls) == 1 + assert actor.calls[0][2] is False + assert actions.shape == (4, 3) + + +def test_sample_offpolicy_actions_rejects_unknown_algo() -> None: + actor = _DummyActor() + + with pytest.raises(ValueError, match="Unsupported off-policy algo_type"): + sample_offpolicy_actions( + actor=actor, + algo_type="unknown", + obs_torch=torch.zeros(2, 4), + prev_dones_torch=torch.zeros(2), + ) diff --git a/tests/config/test_locomotion_params.py b/tests/config/test_locomotion_params.py index 9d61a386b..e82031eac 100644 --- a/tests/config/test_locomotion_params.py +++ b/tests/config/test_locomotion_params.py @@ -41,10 +41,14 @@ def test_flashsac_config_defaults(): cfg = FlashSACConfig() assert cfg.algo == "flashsac" + assert cfg.num_envs == 1024 assert cfg.batch_size == 2048 + assert cfg.warmup_steps == 100000 + assert cfg.gamma == pytest.approx(0.97) assert cfg.obs_normalization is False assert isinstance(cfg.algo_params, FlashSACAlgoParams) assert cfg.algo_params.normalize_reward is True + assert cfg.algo_params.use_compile is False def test_ppo_config_defaults(): @@ -152,6 +156,33 @@ def test_offpolicy_flashsac_g1_task_overrides(): assert cfg.algo.algo_params.normalize_reward is True +def test_offpolicy_flashsac_g1_joystick_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_joystick/mujoco"], + ) + assert cfg.algo.algo == "flashsac" + assert cfg.training.task_name == "G1JoystickFlatTerrain" + assert cfg.training.sim_backend == "mujoco" + assert cfg.training.use_amp is True + assert cfg.algo.num_envs == 1024 + assert cfg.algo.updates_per_step == 2 + assert cfg.algo.num_envs == 1024 + assert cfg.algo.replay_buffer_n == 9766 + assert cfg.algo.warmup_steps == 100000 + assert cfg.algo.updates_per_step == 2 + assert cfg.algo.gamma == pytest.approx(0.97) + assert cfg.env.control_config.action_scale == pytest.approx(0.5) + assert cfg.env.commands.vel_limit[0] == [-1.0, -0.5, -1.0] + assert cfg.reward.scales.tracking_ang_vel == pytest.approx(0.75) + assert cfg.reward.scales.base_height == pytest.approx(0.0) + + # --------------------------------------------------------------------------- # Hydra YAML loading — appo # --------------------------------------------------------------------------- diff --git a/tests/scripts/test_train_scripts.py b/tests/scripts/test_train_scripts.py index f48607502..14f73ce2a 100644 --- a/tests/scripts/test_train_scripts.py +++ b/tests/scripts/test_train_scripts.py @@ -84,7 +84,11 @@ def _load_script(name: str) -> Any: def _offpolicy_cfg(overrides=None): GlobalHydra.instance().clear() with initialize_config_dir(config_dir=str(_CONF_DIR / "offpolicy"), version_base="1.3"): - return compose("config", overrides=_normalize_overrides(overrides, offpolicy=True)) + return compose( + "config", + overrides=_normalize_overrides(overrides, offpolicy=True), + return_hydra_config=True, + ) def _ppo_cfg(overrides=None): @@ -903,7 +907,10 @@ def value(self, obs): values = mod.get_time_limit_bootstrap_values(state, model, mod.mx.float32) - assert values is not None + if values is None: + raise AssertionError("expected bootstrap values") + if model.last_obs is None: + raise AssertionError("expected model to receive observations") 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()), @@ -1170,16 +1177,22 @@ def test_offpolicy_td3_hydra_default_algo_log_name(): def test_offpolicy_flashsac_hydra_algo_log_name(): - cfg = _offpolicy_cfg(["algo=flashsac", "task=flashsac/g1_sac/mujoco"]) + cfg = _offpolicy_cfg(["algo=flashsac", "task=flashsac/g1_joystick/mujoco"]) assert cfg.algo.algo_log_name == "flash_sac" assert cfg.algo.load_run == "-1" +def test_offpolicy_flashsac_g1_joystick_task_composes() -> None: + cfg = _offpolicy_cfg(["algo=flashsac", "task=flashsac/g1_joystick/mujoco"]) + assert cfg.training.task_name == "G1JoystickFlatTerrain" + assert cfg.training.sim_backend == "mujoco" + + def test_offpolicy_flashsac_rejects_multi_gpu(): cfg = _offpolicy_cfg( [ "algo=flashsac", - "task=flashsac/g1_sac/mujoco", + "task=flashsac/g1_joystick/mujoco", "training.num_gpus=2", ] ) @@ -1188,6 +1201,20 @@ def test_offpolicy_flashsac_rejects_multi_gpu(): _offpolicy().build_runner("flashsac", cfg) +@pytest.mark.parametrize( + ("algo", "task"), + [ + ("flashsac", "sac/g1_sac/mujoco"), + ("sac", "flashsac/g1_sac/mujoco"), + ], +) +def test_offpolicy_rejects_algo_task_owner_mismatch(algo: str, task: str): + cfg = _offpolicy_cfg([f"algo={algo}", f"task={task}"]) + + with pytest.raises(ValueError, match="Off-policy algo/task mismatch"): + _offpolicy().build_runner(algo, 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)