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..9de1f4b79 --- /dev/null +++ b/src/unilab/algos/torch/offpolicy/multi_gpu_runner.py @@ -0,0 +1,422 @@ +"""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 datetime import timedelta +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, timeout=timedelta(seconds=120) + ) + + 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 (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/algos/torch/offpolicy/worker.py b/src/unilab/algos/torch/offpolicy/worker.py index 5a309c285..b1a4497e8 100644 --- a/src/unilab/algos/torch/offpolicy/worker.py +++ b/src/unilab/algos/torch/offpolicy/worker.py @@ -4,6 +4,7 @@ actor policy. Runs in a subprocess; writes to ReplayBuffer. """ +import queue import sys import numpy as np @@ -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.py b/src/unilab/envs/locomotion/g1/joystick.py index 7303d80e8..8c2cde596 100644 --- a/src/unilab/envs/locomotion/g1/joystick.py +++ b/src/unilab/envs/locomotion/g1/joystick.py @@ -14,22 +14,24 @@ 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 +113,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 +152,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/g1/joystick_sac.py b/src/unilab/envs/locomotion/g1/joystick_sac.py index 3a5cb2104..e332eaf56 100644 --- a/src/unilab/envs/locomotion/g1/joystick_sac.py +++ b/src/unilab/envs/locomotion/g1/joystick_sac.py @@ -13,6 +13,29 @@ 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.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,6 +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: ObsConfig = field(default_factory=_g1_sac_obs_config) @registry.env("G1WalkTaskMjSAC", sim_backend="mujoco") @@ -152,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") @@ -169,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") @@ -181,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()) diff --git a/src/unilab/envs/locomotion/go1/joystick.py b/src/unilab/envs/locomotion/go1/joystick.py index dc6a38223..db431e414 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,20 @@ 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 +94,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 +129,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..b6413de19 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,19 @@ 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 +88,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 +119,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 diff --git a/src/unilab/ipc/replay_buffer.py b/src/unilab/ipc/replay_buffer.py index 1446e8b3c..b19da2900 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,47 @@ def __init__(self, capacity: int, obs_dim: int, action_dim: int, device: str): 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. + + 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 +216,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], 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..471862c21 --- /dev/null +++ b/tests/envs/test_env_configs.py @@ -0,0 +1,113 @@ +"""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. + + 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 == 101, ( + f"obs_config total dim is {total}, expected 101. " + "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()