From b2b216740961a96d490830ada6ebe4633c7f68b7 Mon Sep 17 00:00:00 2001 From: tatp Date: Mon, 16 Mar 2026 20:45:22 +0000 Subject: [PATCH 1/7] feat: add multi-GPU FastSAC training via NCCL all-reduce MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Enables N-GPU training on a single node by spawning one Learner worker per GPU (torch.multiprocessing.spawn). Each worker samples an independent mini-batch from the shared ReplayBuffer and averages gradients with NCCL all_reduce — equivalent to an N× larger effective batch size per wall-clock second. The single Collector process and single-GPU code path are unchanged. Key changes: - ReplayBuffer: add defer_gpu flag (skips CUDA tensor/stream creation so the buffer is picklable for mp.spawn); add init_local_gpu_cache(device) for per-process GPU cache initialisation; fix sample() to use the actual GPU device rather than hardcoded "cuda" - FastSACLearner: add world_size param, _reduce_gradients() helper, and all-reduce calls in update_critic / update_actor (no-op when world_size=1) - FastSACRunner: forward world_size to FastSACLearner - offpolicy/multi_gpu_runner.py (new): MultiGPUOffPolicyRunner + worker fn - train_offpolicy.py: add --num_gpus flag; routes to MultiGPUOffPolicyRunner when num_gpus > 1, otherwise falls back to existing single-GPU runner Co-Authored-By: Claude Sonnet 4.6 --- scripts/train_offpolicy.py | 66 +++ .../algos/torch/common/shared_buffer.py | 6 +- src/unilab/algos/torch/fast_sac/learner.py | 35 +- src/unilab/algos/torch/fast_sac/runner.py | 2 + src/unilab/algos/torch/offpolicy/__init__.py | 3 +- .../algos/torch/offpolicy/multi_gpu_runner.py | 420 ++++++++++++++++++ src/unilab/ipc/replay_buffer.py | 47 +- 7 files changed, 564 insertions(+), 15 deletions(-) create mode 100644 src/unilab/algos/torch/offpolicy/multi_gpu_runner.py diff --git a/scripts/train_offpolicy.py b/scripts/train_offpolicy.py index 9aad6684c..ff7cf6db8 100644 --- a/scripts/train_offpolicy.py +++ b/scripts/train_offpolicy.py @@ -133,13 +133,79 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument( "--use_amp", action="store_true", help="Enable mixed precision training (FP16)" ) + parser.add_argument( + "--num_gpus", + type=int, + default=1, + help="Number of GPUs for multi-GPU training via NCCL all-reduce (default: 1 = single GPU)", + ) return parser def build_runner(algo_name: str, args, cfg): """Build algorithm runner from unified config.""" if algo_name == "sac": + from unilab.algos.torch.fast_sac.learner import FastSACLearner from unilab.algos.torch.fast_sac.runner import FastSACRunner + from unilab.utils.device_utils import get_default_device, get_env_dims + + # Multi-GPU path + if getattr(args, "num_gpus", 1) > 1: + from unilab.algos.torch.offpolicy.multi_gpu_runner import MultiGPUOffPolicyRunner + from unilab.base import registry + from unilab.utils.algo_utils import ensure_registries + + ensure_registries() + device = args.device or get_default_device() + env = registry.make(args.task, num_envs=1, sim_backend=args.sim_backend) + assert env.observation_space.shape and env.action_space.shape + obs_dim = env.observation_space.shape[0] + action_dim = env.action_space.shape[0] + env.close() + + # Learner kwargs needed to reconstruct FastSACLearner in each worker + learner_kwargs = { + "obs_dim": obs_dim, + "action_dim": action_dim, + "gamma": cfg.gamma, + "tau": cfg.tau, + "actor_lr": cfg.actor_lr, + "critic_lr": cfg.critic_lr, + "alpha_lr": cfg.algo_params.alpha_lr, + "alpha_init": cfg.algo_params.alpha_init, + "target_entropy_ratio": cfg.algo_params.target_entropy_ratio, + "actor_hidden_dim": cfg.actor_hidden_dim, + "critic_hidden_dim": cfg.critic_hidden_dim, + "num_atoms": cfg.num_atoms, + "use_layer_norm": cfg.use_layer_norm, + "max_grad_norm": cfg.algo_params.max_grad_norm, + "use_amp": args.use_amp, + # symmetry not supported in multi-GPU mode (mujoco_model not picklable) + "use_symmetry": False, + } + # Main-process learner (used for WeightSync initialisation and single-GPU fallback) + main_learner = FastSACLearner(device=device, **learner_kwargs) + + return MultiGPUOffPolicyRunner( + learner=main_learner, + env_name=args.task, + algo_type="sac", + learner_kwargs=learner_kwargs, + num_gpus=args.num_gpus, + num_envs=cfg.num_envs, + replay_buffer_n=cfg.replay_buffer_n, + batch_size=cfg.batch_size, + warmup_steps=cfg.warmup_steps, + updates_per_step=cfg.updates_per_step, + policy_frequency=cfg.policy_frequency, + sync_collection=not args.no_sync_collection, + env_steps_per_sync=cfg.env_steps_per_sync, + device=device, + actor_hidden_dim=cfg.actor_hidden_dim, + use_layer_norm=cfg.use_layer_norm, + obs_normalization=False, + sim_backend=args.sim_backend, + ) return FastSACRunner( env_name=args.task, diff --git a/src/unilab/algos/torch/common/shared_buffer.py b/src/unilab/algos/torch/common/shared_buffer.py index 1ecccc5b0..e7e15aa0b 100644 --- a/src/unilab/algos/torch/common/shared_buffer.py +++ b/src/unilab/algos/torch/common/shared_buffer.py @@ -6,14 +6,16 @@ class SharedBufferBase: """Device-adaptive shared memory buffer base class.""" - def __init__(self, capacity: int, device: str): + def __init__(self, capacity: int, device: str, defer_gpu: bool = False): self.capacity = capacity self.device = device self.ptr = torch.zeros(1, dtype=torch.int64).share_memory_() if device == "cuda": self._gpu_synced_ptr = 0 - self._cuda_stream = torch.cuda.Stream() + # defer_gpu=True: skip CUDA stream creation so the buffer is picklable + # (used in multi-GPU mode where workers call init_local_gpu_cache later) + self._cuda_stream = None if defer_gpu else torch.cuda.Stream() # type: ignore[assignment] else: self._gpu_synced_ptr = None # type: ignore[assignment] self._cuda_stream = None # type: ignore[assignment] diff --git a/src/unilab/algos/torch/fast_sac/learner.py b/src/unilab/algos/torch/fast_sac/learner.py index b3b1e652d..99cf6e8f1 100644 --- a/src/unilab/algos/torch/fast_sac/learner.py +++ b/src/unilab/algos/torch/fast_sac/learner.py @@ -14,6 +14,7 @@ from typing import Any, Dict, Tuple import torch +import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F import torch.optim as optim @@ -359,6 +360,7 @@ def __init__( use_amp: bool = False, mujoco_model=None, obs_structure: dict | None = None, + world_size: int = 1, ): self.device = device self.gamma = gamma @@ -366,6 +368,7 @@ def __init__( self.max_grad_norm = max_grad_norm self.use_autotune = use_autotune self.use_amp = use_amp and device == "cuda" + self.world_size = world_size # Build actor self.actor = SACActor( @@ -455,6 +458,27 @@ def __init__( assert obs_structure is not None self.symmetry = G1SymmetryAugmentation(mujoco_model, obs_structure, device=device) + def _reduce_gradients(self, model: nn.Module) -> None: + """All-reduce gradients across all workers and divide by world_size. + + Must be called after ``backward()`` and, when using AMP, after + ``scaler.unscale_(optimizer)`` so that gradients are in full precision. + """ + if self.world_size <= 1: + return + grads = [p.grad.view(-1) for p in model.parameters() if p.grad is not None] + if not grads: + return + flat = torch.cat(grads) + dist.all_reduce(flat, op=dist.ReduceOp.SUM) + flat /= self.world_size + offset = 0 + for p in model.parameters(): + if p.grad is not None: + n = p.grad.numel() + p.grad.copy_(flat[offset : offset + n].view_as(p.grad)) + offset += n + def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: """One critic update step.""" obs = batch["obs"] @@ -504,8 +528,9 @@ def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: self.q_optimizer.zero_grad(set_to_none=True) if self.scaler: self.scaler.scale(qf_loss).backward() + self.scaler.unscale_(self.q_optimizer) + self._reduce_gradients(self.qnet) if self.max_grad_norm > 0: - self.scaler.unscale_(self.q_optimizer) critic_grad_norm = torch.nn.utils.clip_grad_norm_( self.qnet.parameters(), max_norm=self.max_grad_norm ) @@ -515,6 +540,7 @@ def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: self.scaler.update() else: qf_loss.backward() + self._reduce_gradients(self.qnet) if self.max_grad_norm > 0: critic_grad_norm = torch.nn.utils.clip_grad_norm_( self.qnet.parameters(), max_norm=self.max_grad_norm @@ -536,6 +562,9 @@ def update_critic(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: ).mean() if torch.isfinite(alpha_loss): alpha_loss.backward() + if self.world_size > 1 and self.log_alpha.grad is not None: + dist.all_reduce(self.log_alpha.grad, op=dist.ReduceOp.SUM) + self.log_alpha.grad /= self.world_size self.alpha_optimizer.step() return { @@ -574,8 +603,9 @@ def update_actor(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: self.actor_optimizer.zero_grad(set_to_none=True) if self.scaler: self.scaler.scale(actor_loss).backward() + self.scaler.unscale_(self.actor_optimizer) + self._reduce_gradients(self.actor) if self.max_grad_norm > 0: - self.scaler.unscale_(self.actor_optimizer) actor_grad_norm = torch.nn.utils.clip_grad_norm_( self.actor.parameters(), max_norm=self.max_grad_norm ) @@ -585,6 +615,7 @@ def update_actor(self, batch: Dict[str, torch.Tensor]) -> Dict[str, float]: self.scaler.update() else: actor_loss.backward() + self._reduce_gradients(self.actor) if self.max_grad_norm > 0: actor_grad_norm = torch.nn.utils.clip_grad_norm_( self.actor.parameters(), max_norm=self.max_grad_norm diff --git a/src/unilab/algos/torch/fast_sac/runner.py b/src/unilab/algos/torch/fast_sac/runner.py index 6fe0d05b2..e4c214ccd 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_amp: bool = False, sim_backend: str = "mujoco", use_symmetry: bool = False, + world_size: int = 1, ): from unilab.base import registry from unilab.utils.algo_utils import ensure_registries @@ -77,6 +78,7 @@ def __init__( use_symmetry=use_symmetry, mujoco_model=mujoco_model, obs_structure=obs_structure, + world_size=getattr(self, "world_size", world_size), ) # Auto-adjust batch_size when symmetry is enabled diff --git a/src/unilab/algos/torch/offpolicy/__init__.py b/src/unilab/algos/torch/offpolicy/__init__.py index 62289a187..4169e5262 100644 --- a/src/unilab/algos/torch/offpolicy/__init__.py +++ b/src/unilab/algos/torch/offpolicy/__init__.py @@ -1,6 +1,7 @@ """Off-policy RL unified infrastructure.""" +from unilab.algos.torch.offpolicy.multi_gpu_runner import MultiGPUOffPolicyRunner from unilab.algos.torch.offpolicy.runner import OffPolicyRunner from unilab.algos.torch.offpolicy.worker import off_policy_collector_fn -__all__ = ["OffPolicyRunner", "off_policy_collector_fn"] +__all__ = ["OffPolicyRunner", "MultiGPUOffPolicyRunner", "off_policy_collector_fn"] diff --git a/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py new file mode 100644 index 000000000..7712e4b43 --- /dev/null +++ b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py @@ -0,0 +1,420 @@ +"""Multi-GPU off-policy runner using NCCL all-reduce for FastSAC. + +Architecture: + Main process → creates ReplayBuffer (host-only), WeightSync, queues + → spawns Collector subprocess (CPU, env simulation) + → spawns N Learner workers via mp.spawn (one per GPU) + Learner rank i → creates its own GPU cache, communicates via NCCL all_reduce + Collector → talks only to rank 0 via collection_ready_queue / trainer_done_queue +""" + +from __future__ import annotations + +import os +import queue +import socket +import sys +import time +from collections import defaultdict, deque +from typing import Any, Dict, Optional + +import torch +import torch.distributed as dist +import torch.multiprocessing as tmp # torch.multiprocessing for spawn + +from unilab.algos.torch.fast_sac.learner import FastSACLearner +from unilab.algos.torch.offpolicy.runner import OffPolicyRunner +from unilab.algos.torch.offpolicy.worker import off_policy_collector_fn +from unilab.ipc import SharedWeightSync +from unilab.ipc.async_runner import _SPAWN_CTX +from unilab.ipc.replay_buffer import ReplayBuffer +from unilab.utils.offpolicy_logger import OffPolicyLogger + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +def _broadcast_initial_params(learner: FastSACLearner, rank: int) -> None: + """Broadcast rank-0 initial parameters to all workers for consistent starting point.""" + for model in [learner.actor, learner.qnet]: + for p in model.parameters(): + dist.broadcast(p.data, src=0) + dist.broadcast(learner.log_alpha.data, src=0) + + +def _drain_metrics( + metrics_queue: Any, + reward_history: deque, + reward_components: dict, + logger: Optional[OffPolicyLogger], +) -> None: + while not metrics_queue.empty(): + try: + m = metrics_queue.get_nowait() + if "error" in m: + if logger: + logger.log_status(f"[red]Collector ERROR: {m['error']}[/]") + return + + if "mean_ep_reward" in m: + reward_history.append(m["mean_ep_reward"]) + if "reward_components" in m: + reward_components.clear() + reward_components.update(m["reward_components"]) + if "mean_ep_length" in m and logger: + logger.update_ep_length(m["mean_ep_length"]) + if "collector_timing_ms" in m and logger: + logger.update_collector_timing(m["collector_timing_ms"]) + if ("timeout_rate" in m or "terminated_rate" in m) and logger: + logger.update_done_rates( + timeout_rate=float(m.get("timeout_rate", 0.0)), + terminated_rate=float(m.get("terminated_rate", 0.0)), + ) + if "total_steps" in m and "buffer_size" in m and logger: + logger.log_collector( + m["total_steps"], + m["buffer_size"], + m.get("mean_ep_reward", 0.0), + ) + except Exception as e: + print(f"[MultiGPU] metrics drain error: {e}", file=sys.stderr) + break + + +def _learner_worker( + rank: int, + world_size: int, + learner_kwargs: Dict[str, Any], + runner_kwargs: Dict[str, Any], + replay_buffer: ReplayBuffer, + weight_sync_name: str, + weight_sync_lock: Any, + weight_param_shapes: Dict[str, Any], + stop_event: Any, + collection_ready_queue: Any, + trainer_done_queue: Any, + metrics_queue: Any, + master_port: int, +) -> None: + """Worker function executed on each GPU (called via torch.multiprocessing.spawn).""" + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(master_port) + device = f"cuda:{rank}" + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=world_size) + + try: + # 1. Initialise per-process GPU cache (host tensors already shared) + replay_buffer.init_local_gpu_cache(device) + + # 2. Create learner on this device + learner = FastSACLearner(device=device, world_size=world_size, **learner_kwargs) + + # 3. Broadcast rank-0 params so all workers start identically + _broadcast_initial_params(learner, rank) + + # 4. Reconnect to the shared weight-sync buffer + weight_sync = SharedWeightSync( + weight_param_shapes, create=False, shm_name=weight_sync_name, lock=weight_sync_lock + ) + + # 5. Unpack runner config + max_iterations: int = runner_kwargs["max_iterations"] + save_interval: int = runner_kwargs["save_interval"] + log_dir: str = runner_kwargs["log_dir"] + batch_size: int = runner_kwargs["batch_size"] + updates_per_step: int = runner_kwargs["updates_per_step"] + policy_frequency: int = runner_kwargs["policy_frequency"] + sync_collection: bool = runner_kwargs["sync_collection"] + env_steps_per_sync: int = runner_kwargs.get("env_steps_per_sync", 1) + env_name: str = runner_kwargs["env_name"] + num_envs: int = runner_kwargs["num_envs"] + obs_dim: int = runner_kwargs["obs_dim"] + action_dim: int = runner_kwargs["action_dim"] + logger_type: str = runner_kwargs.get("logger_type", "tensorboard") + + # 6. Logger (rank 0 only) + logger: Optional[OffPolicyLogger] = None + if rank == 0: + os.makedirs(log_dir, exist_ok=True) + logger = OffPolicyLogger( + algo_name=f"FastSAC_x{world_size}GPU", + max_iterations=max_iterations, + num_envs=num_envs, + env_name=env_name, + obs_dim=obs_dim, + action_dim=action_dim, + log_dir=log_dir, + log_backend=logger_type, + ) + logger.set_collection_sync(sync_collection, env_steps_per_sync) + logger.start() + + reward_history: deque = deque(maxlen=100) + latest_reward_components: dict = {} + write_read_ema = 0.0 + + # 7. Training loop + for it in range(1, max_iterations + 1): + iter_start = time.time() + + # --- Wait for data (rank 0 only, then barrier syncs everyone) --- + wait_start = time.time() + if rank == 0: + if sync_collection and collection_ready_queue is not None: + while True: + try: + collection_ready_queue.get(timeout=1.0) + break + except queue.Empty: + if stop_event.is_set(): + return + else: + while int(replay_buffer.size[0]) < batch_size: + if stop_event.is_set(): + return + time.sleep(0.1) + _drain_metrics(metrics_queue, reward_history, latest_reward_components, logger) + + dist.barrier() + wait_time = time.time() - wait_start if rank == 0 else 0.0 + collect_time = time.time() - iter_start if rank == 0 else 0.0 + + # --- Training: each rank independently samples a different mini-batch --- + train_start = time.time() + iter_metrics: dict = defaultdict(list) + ptr_before = int(replay_buffer.ptr[0]) if rank == 0 else 0 + + large_batch = replay_buffer.sample(batch_size * updates_per_step) + + for update_idx in range(updates_per_step): + s = update_idx * batch_size + e = s + batch_size + batch = {k: v[s:e] for k, v in large_batch.items()} + + critic_metrics = learner.update_critic(batch) + for k, v in critic_metrics.items(): + iter_metrics[k].append(v) + + if update_idx % policy_frequency == 1: + actor_metrics = learner.update_actor(batch) + for k, v in actor_metrics.items(): + iter_metrics[k].append(v) + + learner.soft_update_target() + + # Barrier: all ranks must finish this iteration before rank 0 proceeds + dist.barrier() + train_time = time.time() - train_start if rank == 0 else 0.0 + + # --- Post-iteration work: rank 0 only --- + if rank == 0: + learner.update_count += 1 + weight_sync.write_weights(learner.actor.state_dict()) + + if sync_collection and trainer_done_queue is not None: + trainer_done_queue.put(1) + + write_delta = int(replay_buffer.ptr[0]) - ptr_before + consume = batch_size * updates_per_step + write_read_ema = 0.9 * write_read_ema + 0.1 * (write_delta / max(consume, 1)) + + import statistics as _stats + + avg_metrics = {k: _stats.mean(v) for k, v in iter_metrics.items() if v} + mean_reward = _stats.mean(reward_history) if reward_history else 0.0 + + if logger: + logger.update_buffer_utilization(write_read_ema) + logger.log_step( + iteration=it, + metrics=avg_metrics, + reward=mean_reward, + reward_components=latest_reward_components, + collect_time=collect_time, + train_time=train_time, + wait_time=wait_time, + ) + + if save_interval > 0 and it % save_interval == 0: + ckpt_path = os.path.join(log_dir, f"model_{it}.pt") + torch.save(learner.get_state_dict(), ckpt_path) + if logger: + logger.log_save(ckpt_path) + + # Final checkpoint (rank 0) + if rank == 0: + ckpt_path = os.path.join(log_dir, f"model_{max_iterations}.pt") + torch.save(learner.get_state_dict(), ckpt_path) + if logger: + logger.log_save(ckpt_path) + logger.finish() + + weight_sync.close() + + finally: + dist.destroy_process_group() + + +class MultiGPUOffPolicyRunner(OffPolicyRunner): + """Multi-GPU off-policy runner. + + Keeps a single Collector on CPU and spawns *num_gpus* Learner workers via + ``torch.multiprocessing.spawn``. Each worker processes independent + mini-batches from the same shared ReplayBuffer; gradients are averaged + with NCCL all_reduce — equivalent to training on a *num_gpus× larger* + effective batch size per wall-clock second. + + Falls back transparently to single-GPU when ``num_gpus <= 1``. + """ + + def __init__( + self, + learner: Any, + env_name: str, + algo_type: str, + learner_kwargs: Dict[str, Any], + num_gpus: int = 1, + **kwargs: Any, + ) -> None: + super().__init__(learner=learner, env_name=env_name, algo_type=algo_type, **kwargs) + self.num_gpus = num_gpus + self.world_size = num_gpus + self._learner_kwargs = learner_kwargs + + def learn( + self, + max_iterations: int = 1500, + save_interval: int = 50, + log_dir: str = "logs", + logger_type: str = "tensorboard", + ) -> None: + if self.num_gpus <= 1: + super().learn( + max_iterations=max_iterations, + save_interval=save_interval, + log_dir=log_dir, + logger_type=logger_type, + ) + return + self._learn_multi_gpu( + max_iterations=max_iterations, + save_interval=save_interval, + log_dir=log_dir, + logger_type=logger_type, + ) + + def _learn_multi_gpu( + self, + max_iterations: int, + save_interval: int, + log_dir: str, + logger_type: str, + ) -> None: + os.makedirs(log_dir, exist_ok=True) + + # --- Shared objects (main process owns, workers share via IPC) --- + buffer_capacity = self.replay_buffer_n * self.num_envs + replay_buffer = ReplayBuffer( + capacity=buffer_capacity, + obs_dim=self.obs_dim, + action_dim=self.action_dim, + device=self.device, + defer_gpu=True, # No GPU tensors in main process; workers call init_local_gpu_cache + ) + self._shared_resources.append(replay_buffer) + + weight_sync = SharedWeightSync.from_state_dict(self.learner.actor.state_dict(), create=True) + self._shared_resources.append(weight_sync) + + collection_ready_queue = None + trainer_done_queue = None + if self.sync_collection: + collection_ready_queue = _SPAWN_CTX.Queue(maxsize=1) + trainer_done_queue = _SPAWN_CTX.Queue(maxsize=1) + trainer_done_queue.put(1) + print( + f"[MultiGPURunner] Collection sync enabled: " + f"env_steps_per_sync={self.env_steps_per_sync}" + ) + + metrics_queue = _SPAWN_CTX.Queue(maxsize=100) + + # --- Start Collector (CPU, single process, unchanged) --- + weight_param_shapes = {k: v.shape for k, v in self.learner.actor.state_dict().items()} + sim_backend: str = self.extra_kwargs.get("sim_backend", "mujoco") + collector_kwargs = { + "env_name": self.env_name, + "num_envs": self.num_envs, + "replay_buffer": replay_buffer, + "weight_sync_name": weight_sync.name, + "weight_sync_lock": weight_sync._lock, + "weight_param_shapes": weight_param_shapes, + "algo_type": self.algo_type, + "actor_hidden_dim": self.actor_hidden_dim, + "use_layer_norm": self.use_layer_norm, + "warmup_steps": self.warmup_steps, + "metrics_queue": metrics_queue, + "sync_collection": self.sync_collection, + "collection_ready_queue": collection_ready_queue, + "trainer_done_queue": trainer_done_queue, + "env_steps_per_sync": self.env_steps_per_sync, + "obs_normalization": False, + "shared_obs_normalizer_stats": None, + "sim_backend": sim_backend, + } + self._start_collector( + target_fn=off_policy_collector_fn, + kwargs={"stop_event": self._stop_event, **collector_kwargs}, + ) + time.sleep(0.5) + if self._collector_process: + print(f"[MultiGPURunner] Collector process alive: {self._collector_process.is_alive()}") + + master_port = _find_free_port() + print( + f"[MultiGPURunner] Spawning {self.num_gpus} Learner workers " + f"(NCCL port {master_port})" + ) + + runner_kwargs: Dict[str, Any] = { + "max_iterations": max_iterations, + "save_interval": save_interval, + "log_dir": log_dir, + "batch_size": self.batch_size, + "updates_per_step": self.updates_per_step, + "policy_frequency": self.policy_frequency, + "sync_collection": self.sync_collection, + "env_steps_per_sync": self.env_steps_per_sync, + "env_name": self.env_name, + "num_envs": self.num_envs, + "obs_dim": self.obs_dim, + "action_dim": self.action_dim, + "logger_type": logger_type, + } + + try: + tmp.spawn( + _learner_worker, + args=( + self.num_gpus, + self._learner_kwargs, + runner_kwargs, + replay_buffer, + weight_sync.name, + weight_sync._lock, + weight_param_shapes, + self._stop_event, + collection_ready_queue, + trainer_done_queue, + metrics_queue, + master_port, + ), + nprocs=self.num_gpus, + join=True, + ) + finally: + self._stop_event.set() diff --git a/src/unilab/ipc/replay_buffer.py b/src/unilab/ipc/replay_buffer.py index 1446e8b3c..5d3390b97 100644 --- a/src/unilab/ipc/replay_buffer.py +++ b/src/unilab/ipc/replay_buffer.py @@ -16,8 +16,15 @@ class ReplayBuffer(SharedBufferBase): - MPS/CPU: single packed host tensor; 1 H2D Metal command per sample() """ - def __init__(self, capacity: int, obs_dim: int, action_dim: int, device: str): - super().__init__(capacity, device) + def __init__( + self, + capacity: int, + obs_dim: int, + action_dim: int, + device: str, + defer_gpu: bool = False, + ): + super().__init__(capacity, device, defer_gpu=defer_gpu) self._obs_dim = obs_dim self._action_dim = action_dim @@ -32,13 +39,14 @@ def __init__(self, capacity: int, obs_dim: int, action_dim: int, device: str): self.dones = torch.zeros(capacity).share_memory_() self.truncated = torch.zeros(capacity).share_memory_() - # GPU cache for zero-copy sampling - self.obs_gpu = torch.empty(capacity, obs_dim, device="cuda") - self.next_obs_gpu = torch.empty(capacity, obs_dim, device="cuda") - self.actions_gpu = torch.empty(capacity, action_dim, device="cuda") - self.rewards_gpu = torch.empty(capacity, device="cuda") - self.dones_gpu = torch.empty(capacity, device="cuda") - self.truncated_gpu = torch.empty(capacity, device="cuda") + if not defer_gpu: + # GPU cache for zero-copy sampling + self.obs_gpu = torch.empty(capacity, obs_dim, device="cuda") + self.next_obs_gpu = torch.empty(capacity, obs_dim, device="cuda") + self.actions_gpu = torch.empty(capacity, action_dim, device="cuda") + self.rewards_gpu = torch.empty(capacity, device="cuda") + self.dones_gpu = torch.empty(capacity, device="cuda") + self.truncated_gpu = torch.empty(capacity, device="cuda") else: # Single packed host tensor; layout: [obs | next_obs | actions | rew | done | trunc] total_dim = 2 * obs_dim + action_dim + 3 @@ -57,6 +65,25 @@ def __init__(self, capacity: int, obs_dim: int, action_dim: int, device: str): c += 1 self._trunc_col = c + def init_local_gpu_cache(self, device: str) -> None: + """Per-process GPU cache initialisation for multi-GPU training. + + Each Learner process calls this once after being spawned, binding + its own GPU cache to ``device``. The shared host tensors (obs, + actions, …) already exist; only the GPU-side copies are created here. + """ + assert self.device == "cuda", "init_local_gpu_cache is only for CUDA buffers" + obs_dim = self.obs.shape[1] + act_dim = self.actions.shape[1] + self.obs_gpu = torch.zeros(self.capacity, obs_dim, device=device) + self.next_obs_gpu = torch.zeros(self.capacity, obs_dim, device=device) + self.actions_gpu = torch.zeros(self.capacity, act_dim, device=device) + self.rewards_gpu = torch.zeros(self.capacity, device=device) + self.dones_gpu = torch.zeros(self.capacity, device=device) + self.truncated_gpu = torch.zeros(self.capacity, device=device) + self._gpu_synced_ptr = 0 + self._cuda_stream = torch.cuda.Stream(device=device) + def add(self, obs, actions, rewards, next_obs, dones, truncated): """Add batch (called by collector).""" n = obs.shape[0] @@ -167,7 +194,7 @@ def sample(self, batch_size: int) -> Dict[str, torch.Tensor]: self._gpu_synced_ptr = ptr - indices = indices.to("cuda") + indices = indices.to(self.obs_gpu.device) return { "obs": self.obs_gpu[indices], "actions": self.actions_gpu[indices], From 2df71a57203f999a395af88e75427fc41306a2f0 Mon Sep 17 00:00:00 2001 From: tatp Date: Mon, 16 Mar 2026 20:58:22 +0000 Subject: [PATCH 2/7] fix: add obs_config to G1JoystickSACCfg; add env config completeness tests G1WalkTaskMjSAC inherits _init_obs_space() from G1JoystickPPO, which reads self._cfg.obs_config. G1JoystickSACCfg extended G1BaseCfg directly (not G1JoystickPPOCfg) so the field was missing, causing: AttributeError: 'G1JoystickSACCfg' object has no attribute 'obs_config' Fix: import obs_cfg from joystick.py and add obs_config: obs_cfg = field(default_factory=obs_cfg) to G1JoystickSACCfg. Tests added (tests/envs/test_env_configs.py): - test_g1_joystick_sac_cfg_has_obs_config: regression test for the exact AttributeError; non-slow, no env.step() required - test_g1_joystick_sac_obs_config_total_dim: verifies obs dim (50) matches what _compute_obs actually concatenates - test_locomotion_env_instantiates (slow, parametrized): actually calls registry.make() for Go1/Go2/G1/G1SAC to catch future init regressions on any registered locomotion env Co-Authored-By: Claude Sonnet 4.6 --- src/unilab/envs/locomotion/g1/joystick_sac.py | 3 +- tests/envs/__init__.py | 0 tests/envs/test_env_configs.py | 111 ++++++++++++++++++ 3 files changed, 113 insertions(+), 1 deletion(-) create mode 100644 tests/envs/__init__.py create mode 100644 tests/envs/test_env_configs.py diff --git a/src/unilab/envs/locomotion/g1/joystick_sac.py b/src/unilab/envs/locomotion/g1/joystick_sac.py index 3a5cb2104..49c2a2e2f 100644 --- a/src/unilab/envs/locomotion/g1/joystick_sac.py +++ b/src/unilab/envs/locomotion/g1/joystick_sac.py @@ -12,7 +12,7 @@ from unilab.base.curriculum import EpisodeLengthTracker, PenaltyCurriculum from unilab.base.dtype_config import get_global_dtype from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv -from unilab.envs.locomotion.g1.joystick import G1JoystickPPO, InitState +from unilab.envs.locomotion.g1.joystick import G1JoystickPPO, InitState, obs_cfg @dataclass @@ -105,6 +105,7 @@ class G1JoystickSACCfg(G1BaseCfg): commands: Commands = field(default_factory=Commands) reward_config: RewardConfigSAC = field(default_factory=RewardConfigSAC) control_config: ControlConfigSAC = field(default_factory=ControlConfigSAC) # type: ignore[assignment] + obs_config: obs_cfg = field(default_factory=obs_cfg) @registry.env("G1WalkTaskMjSAC", sim_backend="mujoco") diff --git a/tests/envs/__init__.py b/tests/envs/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py new file mode 100644 index 000000000..1ed79f0c7 --- /dev/null +++ b/tests/envs/test_env_configs.py @@ -0,0 +1,111 @@ +"""Tests for env config completeness and env instantiation. + +Config-attribute tests (non-slow) verify that config dataclasses expose every +attribute accessed by their paired env class, WITHOUT running a simulation. +They still require MuJoCo to be importable because the config and env classes +live in the same module file. + +Slow tests actually call registry.make() and run a simulation step. +""" + +from __future__ import annotations + +import pytest + +# The G1 env modules import create_backend → mujoco.batch_forward at the top +# level, so all tests in this file need a working MuJoCo installation. +pytest.importorskip("mujoco", reason="mujoco not installed") + +# Some environments also use mujoco.batch_forward (G1 backend). Guard against +# partial MuJoCo installations where the base package installs but platform +# extensions fail (e.g. wrong libstdc++ version). +try: + from mujoco import batch_forward as _ # noqa: F401 +except Exception: + pytest.skip("mujoco.batch_forward not available (platform/libstdc++ issue)", allow_module_level=True) + +from unilab.utils.algo_utils import ensure_registries # noqa: E402 + +# --------------------------------------------------------------------------- +# Non-slow: config attribute completeness (no env.step(), no MuJoCo sim) +# --------------------------------------------------------------------------- + + +def test_g1_joystick_ppo_cfg_has_obs_config(): + """G1JoystickPPOCfg must have obs_config (baseline sanity check).""" + from unilab.envs.locomotion.g1.joystick import G1JoystickPPOCfg + + cfg = G1JoystickPPOCfg() + assert hasattr(cfg, "obs_config"), "G1JoystickPPOCfg missing obs_config" + assert hasattr(cfg.obs_config, "obs_dict") + assert hasattr(cfg.obs_config, "actor_obs") + + +def test_g1_joystick_sac_cfg_has_obs_config(): + """G1JoystickSACCfg must have obs_config because G1WalkTaskMjSAC inherits + _init_obs_space() from G1JoystickPPO which reads self._cfg.obs_config. + + Regression test for: + AttributeError: 'G1JoystickSACCfg' object has no attribute 'obs_config' + """ + from unilab.envs.locomotion.g1.joystick_sac import G1JoystickSACCfg + + cfg = G1JoystickSACCfg() + assert hasattr(cfg, "obs_config"), ( + "G1JoystickSACCfg must declare obs_config so that the inherited " + "G1JoystickPPO._init_obs_space() can resolve self._cfg.obs_config" + ) + assert hasattr(cfg.obs_config, "obs_dict") + assert hasattr(cfg.obs_config, "actor_obs") + + +def test_g1_joystick_sac_obs_config_total_dim(): + """obs_config total dim must match what _compute_obs actually concatenates. + + _compute_obs outputs: + linvel(3) + gyro(3) + gravity(3) + diff(12) + dof_vel(12) + + last_actions(12) + command(3) + gait_phase(2) = 50 + """ + from unilab.envs.locomotion.g1.joystick_sac import G1JoystickSACCfg + + cfg = G1JoystickSACCfg() + total = sum(cfg.obs_config.obs_dict.values()) + assert total == 50, ( + f"obs_config total dim is {total}, expected 50. " + "obs_config.obs_dict does not match _compute_obs output." + ) + + +# --------------------------------------------------------------------------- +# Slow: actual env instantiation (runs MuJoCo physics) +# --------------------------------------------------------------------------- + + +@pytest.mark.slow +@pytest.mark.parametrize( + "env_name", + [ + "Go1JoystickFlatTerrain", + "Go2JoystickFlatTerrain", + "G1JoystickFlatTerrain", + "G1WalkTaskMjSAC", + ], +) +def test_locomotion_env_instantiates(env_name: str): + """Every registered locomotion env must be constructible with num_envs=1 + and expose valid observation/action spaces.""" + ensure_registries() + from unilab.base import registry + + env = registry.make(env_name, num_envs=1, sim_backend="mujoco") + try: + obs_space = env.observation_space + act_space = env.action_space + assert obs_space is not None + assert act_space is not None + assert obs_space.shape is not None and len(obs_space.shape) == 1 + assert act_space.shape is not None and len(act_space.shape) == 1 + assert obs_space.shape[0] > 0 + assert act_space.shape[0] > 0 + finally: + env.close() From 1918ef1332fcb79da197b0fddd7025333bee7d31 Mon Sep 17 00:00:00 2001 From: tatp Date: Mon, 16 Mar 2026 21:18:18 +0000 Subject: [PATCH 3/7] fix: resolve G1WalkTaskMjSAC crashes and multi-GPU hang on trainer crash MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four bugs fixed: 1. joystick_sac.py: Replace broken G1SACObsCfg (class-level dict, no total_dim/actor_indices) and wrong _g1_obs_config (12-DoF PPO) with _g1_sac_obs_config() returning ObsConfig with correct G1 dims (diff:29, dof_vel:29, action:29 → total=101). 2. worker.py: trainer_done_queue.get() had no timeout, so collector blocked forever when the trainer crashed. Replace with a 1-second timeout loop that checks stop_event so the collector exits cleanly. 3. multi_gpu_runner.py: dist.init_process_group used default NCCL timeout (30 min). Add timeout=timedelta(seconds=120) so a worker crash causes a 2-minute hang instead of a 30-minute one. 4. test_env_configs.py: Fix assertion from total==50 (Go1/Go2 PPO dims) to total==101 (G1 29-DoF SAC dims) and update the docstring to match. Co-Authored-By: Claude Sonnet 4.6 --- .../algos/torch/offpolicy/multi_gpu_runner.py | 3 ++- src/unilab/algos/torch/offpolicy/worker.py | 8 +++++- src/unilab/envs/locomotion/g1/joystick_sac.py | 27 +++++++++++++++++-- tests/envs/test_env_configs.py | 10 +++---- 4 files changed, 39 insertions(+), 9 deletions(-) diff --git a/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py index 7712e4b43..9548cf85f 100644 --- a/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py +++ b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py @@ -16,6 +16,7 @@ import sys import time from collections import defaultdict, deque +from datetime import timedelta from typing import Any, Dict, Optional import torch @@ -104,7 +105,7 @@ def _learner_worker( os.environ["MASTER_PORT"] = str(master_port) device = f"cuda:{rank}" torch.cuda.set_device(rank) - dist.init_process_group("nccl", rank=rank, world_size=world_size) + dist.init_process_group("nccl", rank=rank, world_size=world_size, timeout=timedelta(seconds=120)) try: # 1. Initialise per-process GPU cache (host tensors already shared) diff --git a/src/unilab/algos/torch/offpolicy/worker.py b/src/unilab/algos/torch/offpolicy/worker.py index 5a309c285..eda47f17e 100644 --- a/src/unilab/algos/torch/offpolicy/worker.py +++ b/src/unilab/algos/torch/offpolicy/worker.py @@ -5,6 +5,7 @@ """ import sys +import queue import numpy as np import torch @@ -276,7 +277,12 @@ def _run_collector( ): if env_steps_since_sync >= env_steps_per_sync: collection_ready_queue.put(1) - trainer_done_queue.get() # Wait for trainer + while not stop_event.is_set(): + try: + trainer_done_queue.get(timeout=1.0) + break + except queue.Empty: + continue env_steps_since_sync = 0 # Progress log every 2 seconds diff --git a/src/unilab/envs/locomotion/g1/joystick_sac.py b/src/unilab/envs/locomotion/g1/joystick_sac.py index 49c2a2e2f..f87e9e235 100644 --- a/src/unilab/envs/locomotion/g1/joystick_sac.py +++ b/src/unilab/envs/locomotion/g1/joystick_sac.py @@ -12,7 +12,30 @@ from unilab.base.curriculum import EpisodeLengthTracker, PenaltyCurriculum from unilab.base.dtype_config import get_global_dtype from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv -from unilab.envs.locomotion.g1.joystick import G1JoystickPPO, InitState, obs_cfg +from unilab.envs.locomotion.g1.joystick import G1JoystickPPO, InitState +from unilab.envs.locomotion.obs_config import ObsConfig + + +def _g1_sac_obs_config() -> ObsConfig: + """ObsConfig for G1 (29 DoF) SAC. + + Must stay consistent with G1JoystickPPO._compute_obs() which outputs: + vel(3) + gyro(3) + gravity(3) + diff(29) + dof_vel(29) + + action(29) + cmd(3) + phase(2) = 101 total + """ + return ObsConfig( + obs_dict={ + "vel": 3, + "gyro": 3, + "gravity": 3, + "diff": 29, + "dof_vel": 29, + "action": 29, + "cmd": 3, + "phase": 2, + }, + actor_obs=["gyro", "gravity", "diff", "dof_vel", "action", "cmd", "phase"], + ) @dataclass @@ -105,7 +128,7 @@ class G1JoystickSACCfg(G1BaseCfg): commands: Commands = field(default_factory=Commands) reward_config: RewardConfigSAC = field(default_factory=RewardConfigSAC) control_config: ControlConfigSAC = field(default_factory=ControlConfigSAC) # type: ignore[assignment] - obs_config: obs_cfg = field(default_factory=obs_cfg) + obs_config: ObsConfig = field(default_factory=_g1_sac_obs_config) @registry.env("G1WalkTaskMjSAC", sim_backend="mujoco") diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 1ed79f0c7..4329c2f6a 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -62,16 +62,16 @@ def test_g1_joystick_sac_cfg_has_obs_config(): def test_g1_joystick_sac_obs_config_total_dim(): """obs_config total dim must match what _compute_obs actually concatenates. - _compute_obs outputs: - linvel(3) + gyro(3) + gravity(3) + diff(12) + dof_vel(12) - + last_actions(12) + command(3) + gait_phase(2) = 50 + G1JoystickPPO._compute_obs outputs (G1 has 29 DoF): + linvel(3) + gyro(3) + gravity(3) + diff(29) + dof_vel(29) + + last_actions(29) + command(3) + gait_phase(2) = 101 """ from unilab.envs.locomotion.g1.joystick_sac import G1JoystickSACCfg cfg = G1JoystickSACCfg() total = sum(cfg.obs_config.obs_dict.values()) - assert total == 50, ( - f"obs_config total dim is {total}, expected 50. " + assert total == 101, ( + f"obs_config total dim is {total}, expected 101. " "obs_config.obs_dict does not match _compute_obs output." ) From ee7e82339bcf990eeb09635cd260d0fc16c54f45 Mon Sep 17 00:00:00 2001 From: tatp Date: Mon, 16 Mar 2026 21:26:34 +0000 Subject: [PATCH 4/7] refactor: replace obs_cfg with ObsConfig and eliminate _get_actor_indices duplication MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add locomotion/obs_config.py: ObsConfig dataclass with typed instance fields (obs_dict, actor_obs), total_dim property, and actor_indices property using a single-pass offset table (O(n) vs O(n²) before) - Remove three identical _get_actor_indices methods from go1/go2/g1 joystick envs; logic now lives on ObsConfig itself - Replace bare class-level mutable attributes in obs_cfg (not processed by @dataclass, shared across all instances) with proper factory functions (_go1/go2/g1_obs_config) returning per-instance ObsConfig - Remove print(k) / print(indices) debug output from all three envs - Simplify _init_obs_space to obs_cfg.total_dim / obs_cfg.actor_indices Co-Authored-By: Claude Sonnet 4.6 --- src/unilab/envs/locomotion/g1/joystick.py | 45 +++++----------------- src/unilab/envs/locomotion/go1/joystick.py | 45 +++++----------------- src/unilab/envs/locomotion/go2/joystick.py | 44 +++++---------------- src/unilab/envs/locomotion/obs_config.py | 37 ++++++++++++++++++ 4 files changed, 67 insertions(+), 104 deletions(-) create mode 100644 src/unilab/envs/locomotion/obs_config.py diff --git a/src/unilab/envs/locomotion/g1/joystick.py b/src/unilab/envs/locomotion/g1/joystick.py index 7303d80e8..d99ccb7a1 100644 --- a/src/unilab/envs/locomotion/g1/joystick.py +++ b/src/unilab/envs/locomotion/g1/joystick.py @@ -14,22 +14,15 @@ from unilab.base.dtype_config import get_global_dtype from unilab.base.np_env import NpEnvState from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv +from unilab.envs.locomotion.obs_config import ObsConfig from unilab.utils.math_utils import np_quat_mul, np_yaw_to_quat -@dataclass -class obs_cfg: - obs_dict = { - "vel": 3, - "gyro": 3, - "gravity": 3, - "diff": 12, - "dof_vel": 12, - "action": 12, - "cmd": 3, - "phase": 2, - } # 'obs_name': dim - actor_obs = ["gyro", "gravity", "diff", "dof_vel", "action", "cmd", "phase"] +def _g1_obs_config() -> ObsConfig: + return ObsConfig( + obs_dict={"vel": 3, "gyro": 3, "gravity": 3, "diff": 12, "dof_vel": 12, "action": 12, "cmd": 3, "phase": 2}, + actor_obs=["gyro", "gravity", "diff", "dof_vel", "action", "cmd", "phase"], + ) @dataclass @@ -111,7 +104,7 @@ class G1JoystickPPOCfg(G1BaseCfg): init_state: InitState = field(default_factory=InitState) commands: Commands = field(default_factory=Commands) reward_config: RewardConfigPPO = field(default_factory=RewardConfigPPO) - obs_config: obs_cfg = field(default_factory=obs_cfg) + obs_config: ObsConfig = field(default_factory=_g1_obs_config) @registry.env("G1JoystickFlatTerrain", sim_backend="mujoco") @@ -150,29 +143,11 @@ def _init_reward_functions(self): } def _init_obs_space(self): - num_obs = 0 - for value in self._cfg.obs_config.obs_dict.values(): - num_obs += value + obs_cfg = self._cfg.obs_config self._observation_space = gym.spaces.Box( - low=-float("inf"), high=float("inf"), shape=(num_obs,), dtype=float + low=-float("inf"), high=float("inf"), shape=(obs_cfg.total_dim,), dtype=float ) - self.actor_indices = self._get_actor_indices() - - def _get_actor_indices(self): - s = 0 - indices = [] - for i in self._cfg.obs_config.actor_obs: - s = 0 - for k, v in self._cfg.obs_config.obs_dict.items(): - if k == i: - print(k) - for q in range(s, s + v): - indices.append(q) - s += v - print(indices) - break - s += v - return indices + self.actor_indices = obs_cfg.actor_indices @property def observation_space(self) -> gym.spaces.Box: diff --git a/src/unilab/envs/locomotion/go1/joystick.py b/src/unilab/envs/locomotion/go1/joystick.py index dc6a38223..2b6fb1c93 100644 --- a/src/unilab/envs/locomotion/go1/joystick.py +++ b/src/unilab/envs/locomotion/go1/joystick.py @@ -11,6 +11,7 @@ from unilab.base.dtype_config import get_global_dtype from unilab.base.np_env import NpEnvState from unilab.envs.locomotion.go1.base import Go1BaseCfg, Go1BaseEnv +from unilab.envs.locomotion.obs_config import ObsConfig from unilab.utils.math_utils import np_quat_mul, np_yaw_to_quat @@ -67,19 +68,11 @@ class Domain_Rand: max_force = [1, 1, 0.5] -@dataclass -class obs_cfg: - obs_dict = { - "vel": 3, - "gyro": 3, - "gravity": 3, - "diff": 12, - "dof_vel": 12, - "action": 12, - "cmd": 3, - "phase": 4, - } # 'obs_name': dim - actor_obs = ["gyro", "gravity", "diff", "dof_vel", "action", "cmd", "phase"] +def _go1_obs_config() -> ObsConfig: + return ObsConfig( + obs_dict={"vel": 3, "gyro": 3, "gravity": 3, "diff": 12, "dof_vel": 12, "action": 12, "cmd": 3, "phase": 4}, + actor_obs=["gyro", "gravity", "diff", "dof_vel", "action", "cmd", "phase"], + ) @registry.envcfg("Go1JoystickFlatTerrain") @@ -92,7 +85,7 @@ class Go1JoystickCfg(Go1BaseCfg): reward_config: RewardConfig = field(default_factory=RewardConfig) sensor: JoystickSensor = field(default_factory=JoystickSensor) # type: ignore[assignment] domain_rand: Domain_Rand = field(default_factory=Domain_Rand) - obs_config: obs_cfg = field(default_factory=obs_cfg) + obs_config: ObsConfig = field(default_factory=_go1_obs_config) @registry.env("Go1JoystickFlatTerrain", sim_backend="mujoco") @@ -127,29 +120,11 @@ def _init_reward_functions(self): } def _init_obs_space(self): - num_obs = 0 - for value in self._cfg.obs_config.obs_dict.values(): - num_obs += value + obs_cfg = self._cfg.obs_config self._observation_space = gym.spaces.Box( - low=-float("inf"), high=float("inf"), shape=(num_obs,), dtype=float + low=-float("inf"), high=float("inf"), shape=(obs_cfg.total_dim,), dtype=float ) - self.actor_indices = self._get_actor_indices() - - def _get_actor_indices(self): - s = 0 - indices = [] - for i in self._cfg.obs_config.actor_obs: - s = 0 - for k, v in self._cfg.obs_config.obs_dict.items(): - if k == i: - print(k) - for q in range(s, s + v): - indices.append(q) - s += v - print(indices) - break - s += v - return indices + self.actor_indices = obs_cfg.actor_indices @property def observation_space(self) -> gym.spaces.Box: diff --git a/src/unilab/envs/locomotion/go2/joystick.py b/src/unilab/envs/locomotion/go2/joystick.py index 7dcc5f0dc..6cfe5c6b4 100644 --- a/src/unilab/envs/locomotion/go2/joystick.py +++ b/src/unilab/envs/locomotion/go2/joystick.py @@ -10,6 +10,7 @@ from unilab.base.backend import create_backend from unilab.base.np_env import NpEnvState from unilab.envs.locomotion.go2.base import Go2BaseCfg, Go2BaseEnv +from unilab.envs.locomotion.obs_config import ObsConfig from unilab.utils.math_utils import np_quat_mul, np_yaw_to_quat @@ -18,18 +19,11 @@ class InitState: pos = [0.0, 0.0, 0.42] -@dataclass -class obs_cfg: - obs_dict = { - "vel": 3, - "gyro": 3, - "gravity": 3, - "diff": 12, - "dof_vel": 12, - "action": 12, - "cmd": 3, - } # 'obs_name': dim - actor_obs = ["gyro", "gravity", "diff", "dof_vel", "action", "cmd"] +def _go2_obs_config() -> ObsConfig: + return ObsConfig( + obs_dict={"vel": 3, "gyro": 3, "gravity": 3, "diff": 12, "dof_vel": 12, "action": 12, "cmd": 3}, + actor_obs=["gyro", "gravity", "diff", "dof_vel", "action", "cmd"], + ) @dataclass @@ -86,7 +80,7 @@ class Go2JoystickCfg(Go2BaseCfg): commands: Commands = field(default_factory=Commands) reward_config: RewardConfig = field(default_factory=RewardConfig) domain_rand: Domain_Rand = field(default_factory=Domain_Rand) - obs_config: obs_cfg = field(default_factory=obs_cfg) + obs_config: ObsConfig = field(default_factory=_go2_obs_config) @registry.env("Go2JoystickFlatTerrain", sim_backend="mujoco") @@ -117,29 +111,11 @@ def _init_reward_functions(self): } def _init_obs_space(self): - num_obs = 0 - for value in self._cfg.obs_config.obs_dict.values(): - num_obs += value + obs_cfg = self._cfg.obs_config self._observation_space = gym.spaces.Box( - low=-float("inf"), high=float("inf"), shape=(num_obs,), dtype=float + low=-float("inf"), high=float("inf"), shape=(obs_cfg.total_dim,), dtype=float ) - self.actor_indices = self._get_actor_indices() - - def _get_actor_indices(self): - s = 0 - indices = [] - for i in self._cfg.obs_config.actor_obs: - s = 0 - for k, v in self._cfg.obs_config.obs_dict.items(): - if k == i: - print(k) - for q in range(s, s + v): - indices.append(q) - s += v - print(indices) - break - s += v - return indices + self.actor_indices = obs_cfg.actor_indices @property def observation_space(self) -> gym.spaces.Box: diff --git a/src/unilab/envs/locomotion/obs_config.py b/src/unilab/envs/locomotion/obs_config.py new file mode 100644 index 000000000..37d9478bb --- /dev/null +++ b/src/unilab/envs/locomotion/obs_config.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass +class ObsConfig: + """Observation layout for locomotion environments. + + ``obs_dict`` maps observation name → dimension in the order the + fields are concatenated into the full observation vector. + ``actor_obs`` lists the subset of keys exposed to the policy. + ``actor_indices`` derives the corresponding flat index slice + automatically, so callers never have to compute offsets by hand. + """ + + obs_dict: dict[str, int] = field(default_factory=dict) + actor_obs: list[str] = field(default_factory=list) + + @property + def total_dim(self) -> int: + return sum(self.obs_dict.values()) + + @property + def actor_indices(self) -> list[int]: + """Flat indices into the full obs vector for the ``actor_obs`` keys.""" + offsets: dict[str, int] = {} + s = 0 + for k, v in self.obs_dict.items(): + offsets[k] = s + s += v + indices: list[int] = [] + for key in self.actor_obs: + if key in offsets: + start = offsets[key] + indices.extend(range(start, start + self.obs_dict[key])) + return indices From d1d0c827eb1e5b65c2494e08c8bbf0d63b327b6d Mon Sep 17 00:00:00 2001 From: tatp Date: Mon, 16 Mar 2026 21:26:36 +0000 Subject: [PATCH 5/7] fix: remove spurious qpos param from all SAC reward functions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The parent _compute_reward calls reward fns with 6 args: (info, linvel, gyro, gravity, dof_pos, dof_vel) Every override in G1WalkTaskMjSAC had a 7th `qpos` parameter, causing TypeError at runtime. Also fix _reward_base_height which was reading qpos[:, 2] (never provided) — use self._backend.get_base_pos()[:, 2] instead, matching the parent class implementation. Co-Authored-By: Claude Sonnet 4.6 --- src/unilab/envs/locomotion/g1/joystick_sac.py | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/src/unilab/envs/locomotion/g1/joystick_sac.py b/src/unilab/envs/locomotion/g1/joystick_sac.py index f87e9e235..e332eaf56 100644 --- a/src/unilab/envs/locomotion/g1/joystick_sac.py +++ b/src/unilab/envs/locomotion/g1/joystick_sac.py @@ -176,16 +176,16 @@ def _init_reward_functions(self): "alive": self._reward_alive, } - def _reward_orientation(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_orientation(self, info, linvel, gyro, gravity, dof_pos, dof_vel): """惩罚姿态偏差（roll/pitch）""" return np.square(gravity[:, 0]) + np.square(gravity[:, 1]) - def _reward_pose(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_pose(self, info, linvel, gyro, gravity, dof_pos, dof_vel): """加权惩罚偏离默认姿态""" diff = dof_pos - self.default_angles return np.sum(self._pose_weights * np.square(diff), axis=1) - def _reward_close_feet_xy(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_close_feet_xy(self, info, linvel, gyro, gravity, dof_pos, dof_vel): """惩罚双脚过近""" left_foot = self._backend.get_sensor_data("left_foot_pos") right_foot = self._backend.get_sensor_data("right_foot_pos") @@ -193,7 +193,7 @@ def _reward_close_feet_xy(self, info, linvel, gyro, gravity, dof_pos, dof_vel, q threshold = self._cfg.reward_config.close_feet_threshold return np.where(feet_dist < threshold, np.square(feet_dist - threshold), 0.0) - def _reward_feet_ori(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_feet_ori(self, info, linvel, gyro, gravity, dof_pos, dof_vel): """惩罚脚部姿态偏差""" left_foot_quat = self._backend.get_sensor_data("left_foot_quat") right_foot_quat = self._backend.get_sensor_data("right_foot_quat") @@ -205,45 +205,45 @@ def _reward_feet_ori(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + np.square(right_foot_quat[:, 2]) ) - def _reward_alive(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_alive(self, info, linvel, gyro, gravity, dof_pos, dof_vel): return np.ones((self._num_envs,), dtype=get_global_dtype()) - def _reward_lin_vel_z(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_lin_vel_z(self, info, linvel, gyro, gravity, dof_pos, dof_vel): """惩罚 z 方向线速度""" return np.square(linvel[:, 2]) - def _reward_base_height(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_base_height(self, info, linvel, gyro, gravity, dof_pos, dof_vel): """惩罚基座高度偏差""" - base_height = qpos[:, 2] + base_height = self._backend.get_base_pos()[:, 2] return np.square(base_height - self._cfg.reward_config.base_height_target) - def _reward_torques(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_torques(self, info, linvel, gyro, gravity, dof_pos, dof_vel): """惩罚力矩""" torques = info.get( "torques", np.zeros((self._num_envs, self._num_action), dtype=get_global_dtype()) ) return np.sum(np.abs(torques), axis=1) - def _reward_energy(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_energy(self, info, linvel, gyro, gravity, dof_pos, dof_vel): """惩罚能量消耗""" torques = info.get( "torques", np.zeros((self._num_envs, self._num_action), dtype=get_global_dtype()) ) return np.sum(np.abs(dof_vel) * np.abs(torques), axis=1) - def _reward_dof_acc(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_dof_acc(self, info, linvel, gyro, gravity, dof_pos, dof_vel): """惩罚关节加速度""" qacc = info.get( "qacc", np.zeros((self._num_envs, self._num_action), dtype=get_global_dtype()) ) return np.sum(np.square(qacc), axis=1) - def _reward_upright(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_upright(self, info, linvel, gyro, gravity, dof_pos, dof_vel): """奖励直立姿态（mjlab flat_orientation）""" xy_squared = np.sum(np.square(gravity[:, :2]), axis=1) return np.exp(-xy_squared / 0.25) - def _reward_feet_air_time(self, info, linvel, gyro, gravity, dof_pos, dof_vel, qpos): + def _reward_feet_air_time(self, info, linvel, gyro, gravity, dof_pos, dof_vel): """奖励脚离地时间""" air_time = info.get( "feet_air_time", np.zeros((self._num_envs, 2), dtype=get_global_dtype()) From ec89e9d519a859c1b821b79013ae96b6176e9efd Mon Sep 17 00:00:00 2001 From: tatp Date: Mon, 16 Mar 2026 21:38:18 +0000 Subject: [PATCH 6/7] fix: make ReplayBuffer picklable by stripping CUDA objects in __getstate__ MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit torch.cuda.Stream and GPU tensors (obs_gpu, etc.) cannot be pickled, but the ReplayBuffer is passed to the collector subprocess via mp.Process which requires pickling. This caused: TypeError: cannot pickle 'Stream' object The collector only calls add() which writes to CPU shared-memory tensors; it never calls sample() and never needs the GPU cache or async stream. __getstate__ strips those fields from the pickled copy only — the original ReplayBuffer in the learner process keeps all GPU objects intact. The multi-GPU path already avoided this by using defer_gpu=True; this fix covers the single-GPU path without needing to change the runner. Co-Authored-By: Claude Sonnet 4.6 --- src/unilab/ipc/replay_buffer.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/src/unilab/ipc/replay_buffer.py b/src/unilab/ipc/replay_buffer.py index 5d3390b97..0aa3bd79a 100644 --- a/src/unilab/ipc/replay_buffer.py +++ b/src/unilab/ipc/replay_buffer.py @@ -65,6 +65,21 @@ def __init__( c += 1 self._trunc_col = c + def __getstate__(self) -> dict: + """Custom pickle support. + + The collector subprocess only calls add(), which writes to the CPU + shared-memory tensors (self.obs, self.actions, …). It never calls + sample(), so it doesn't need the GPU cache or the CUDA stream. + Neither torch.cuda.Stream nor CUDA tensors are picklable, so we strip + them here. The original object in the learner process is unaffected. + """ + state = self.__dict__.copy() + state["_cuda_stream"] = None + for key in ("obs_gpu", "next_obs_gpu", "actions_gpu", "rewards_gpu", "dones_gpu", "truncated_gpu"): + state.pop(key, None) + return state + def init_local_gpu_cache(self, device: str) -> None: """Per-process GPU cache initialisation for multi-GPU training. From d88a34a4b44af23501611e5f44194bf4338858ff Mon Sep 17 00:00:00 2001 From: tatp Date: Mon, 16 Mar 2026 21:51:37 +0000 Subject: [PATCH 7/7] style: apply ruff line-length formatting Co-Authored-By: Claude Sonnet 4.6 --- src/unilab/algos/torch/offpolicy/multi_gpu_runner.py | 7 ++++--- src/unilab/algos/torch/offpolicy/worker.py | 2 +- src/unilab/envs/locomotion/g1/joystick.py | 11 ++++++++++- src/unilab/envs/locomotion/go1/joystick.py | 11 ++++++++++- src/unilab/envs/locomotion/go2/joystick.py | 10 +++++++++- src/unilab/ipc/replay_buffer.py | 9 ++++++++- tests/envs/test_env_configs.py | 4 +++- 7 files changed, 45 insertions(+), 9 deletions(-) diff --git a/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py index 9548cf85f..9de1f4b79 100644 --- a/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py +++ b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py @@ -105,7 +105,9 @@ def _learner_worker( os.environ["MASTER_PORT"] = str(master_port) device = f"cuda:{rank}" torch.cuda.set_device(rank) - dist.init_process_group("nccl", rank=rank, world_size=world_size, timeout=timedelta(seconds=120)) + dist.init_process_group( + "nccl", rank=rank, world_size=world_size, timeout=timedelta(seconds=120) + ) try: # 1. Initialise per-process GPU cache (host tensors already shared) @@ -377,8 +379,7 @@ def _learn_multi_gpu( master_port = _find_free_port() print( - f"[MultiGPURunner] Spawning {self.num_gpus} Learner workers " - f"(NCCL port {master_port})" + f"[MultiGPURunner] Spawning {self.num_gpus} Learner workers (NCCL port {master_port})" ) runner_kwargs: Dict[str, Any] = { diff --git a/src/unilab/algos/torch/offpolicy/worker.py b/src/unilab/algos/torch/offpolicy/worker.py index eda47f17e..b1a4497e8 100644 --- a/src/unilab/algos/torch/offpolicy/worker.py +++ b/src/unilab/algos/torch/offpolicy/worker.py @@ -4,8 +4,8 @@ actor policy. Runs in a subprocess; writes to ReplayBuffer. """ -import sys import queue +import sys import numpy as np import torch diff --git a/src/unilab/envs/locomotion/g1/joystick.py b/src/unilab/envs/locomotion/g1/joystick.py index d99ccb7a1..8c2cde596 100644 --- a/src/unilab/envs/locomotion/g1/joystick.py +++ b/src/unilab/envs/locomotion/g1/joystick.py @@ -20,7 +20,16 @@ def _g1_obs_config() -> ObsConfig: return ObsConfig( - obs_dict={"vel": 3, "gyro": 3, "gravity": 3, "diff": 12, "dof_vel": 12, "action": 12, "cmd": 3, "phase": 2}, + obs_dict={ + "vel": 3, + "gyro": 3, + "gravity": 3, + "diff": 12, + "dof_vel": 12, + "action": 12, + "cmd": 3, + "phase": 2, + }, actor_obs=["gyro", "gravity", "diff", "dof_vel", "action", "cmd", "phase"], ) diff --git a/src/unilab/envs/locomotion/go1/joystick.py b/src/unilab/envs/locomotion/go1/joystick.py index 2b6fb1c93..db431e414 100644 --- a/src/unilab/envs/locomotion/go1/joystick.py +++ b/src/unilab/envs/locomotion/go1/joystick.py @@ -70,7 +70,16 @@ class Domain_Rand: def _go1_obs_config() -> ObsConfig: return ObsConfig( - obs_dict={"vel": 3, "gyro": 3, "gravity": 3, "diff": 12, "dof_vel": 12, "action": 12, "cmd": 3, "phase": 4}, + obs_dict={ + "vel": 3, + "gyro": 3, + "gravity": 3, + "diff": 12, + "dof_vel": 12, + "action": 12, + "cmd": 3, + "phase": 4, + }, actor_obs=["gyro", "gravity", "diff", "dof_vel", "action", "cmd", "phase"], ) diff --git a/src/unilab/envs/locomotion/go2/joystick.py b/src/unilab/envs/locomotion/go2/joystick.py index 6cfe5c6b4..b6413de19 100644 --- a/src/unilab/envs/locomotion/go2/joystick.py +++ b/src/unilab/envs/locomotion/go2/joystick.py @@ -21,7 +21,15 @@ class InitState: def _go2_obs_config() -> ObsConfig: return ObsConfig( - obs_dict={"vel": 3, "gyro": 3, "gravity": 3, "diff": 12, "dof_vel": 12, "action": 12, "cmd": 3}, + obs_dict={ + "vel": 3, + "gyro": 3, + "gravity": 3, + "diff": 12, + "dof_vel": 12, + "action": 12, + "cmd": 3, + }, actor_obs=["gyro", "gravity", "diff", "dof_vel", "action", "cmd"], ) diff --git a/src/unilab/ipc/replay_buffer.py b/src/unilab/ipc/replay_buffer.py index 0aa3bd79a..b19da2900 100644 --- a/src/unilab/ipc/replay_buffer.py +++ b/src/unilab/ipc/replay_buffer.py @@ -76,7 +76,14 @@ def __getstate__(self) -> dict: """ state = self.__dict__.copy() state["_cuda_stream"] = None - for key in ("obs_gpu", "next_obs_gpu", "actions_gpu", "rewards_gpu", "dones_gpu", "truncated_gpu"): + for key in ( + "obs_gpu", + "next_obs_gpu", + "actions_gpu", + "rewards_gpu", + "dones_gpu", + "truncated_gpu", + ): state.pop(key, None) return state diff --git a/tests/envs/test_env_configs.py b/tests/envs/test_env_configs.py index 4329c2f6a..471862c21 100644 --- a/tests/envs/test_env_configs.py +++ b/tests/envs/test_env_configs.py @@ -22,7 +22,9 @@ try: from mujoco import batch_forward as _ # noqa: F401 except Exception: - pytest.skip("mujoco.batch_forward not available (platform/libstdc++ issue)", allow_module_level=True) + pytest.skip( + "mujoco.batch_forward not available (platform/libstdc++ issue)", allow_module_level=True + ) from unilab.utils.algo_utils import ensure_registries # noqa: E402