Repository navigation
Feat/multi gpu fastsac - #43
Merged
Merged
Conversation
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 <noreply@anthropic.com>
…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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
…ices duplication - 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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
…ate__ 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 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.