Repository navigation
Feature/refactor appo - #37
Merged
Merged
Conversation
Implements async PPO using adapter pattern to reuse rsl-rl PPO: - OnPolicyReplayBuffer: device-adaptive storage (CUDA/MPS/CPU) - AsyncPPOLearner: adapts buffer data to RolloutStorage - async_ppo_collector_fn: collects rollouts using PPO.act() - AsyncPPORunner: orchestrates async training Total: 420 lines (5 files) 100% rsl-rl code reuse via thin adapter layer Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Add async_ppo_config() to locomotion_params.py - Create scripts/train_async_ppo.py training entry - Enhance AsyncPPORunner with full logging and checkpointing - Support TensorBoard integration via OffPolicyLogger Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Use pinned memory for faster H2D transfer (CUDA) - Add dedicated CUDA stream for async data transfer - Reduce data copies in worker (zero-copy when possible) - Improve buffer sync efficiency with non_blocking=True Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Create HardwareMonitor utility for CPU/GPU/memory tracking - Integrate monitoring into AsyncPPORunner - Log hardware metrics to TensorBoard every 10 iterations - Support NVML for detailed GPU utilization Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Fix distribution_cfg with correct class name and parameters - Add multi_gpu config field for rsl_rl v4 - Add config conversion in worker process - Fix SharedWeightSync attribute names (name, not shm_name) - Add num_actions attribute to env for PPO construction Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Support None lock in subprocess for lock-free read/write - Fix worker to properly initialize combined state dict Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Add actor/critic prefixes to weight keys to prevent overwriting - Manually initialize distribution_params in learner by calling actor forward - Clone distribution params to avoid view-of-view gradient issues - Remove stochastic_output parameter from PPO.act() call - Add lock-free weight sync support in worker subprocess
…r pipeline - Add initialization wait for collector before starting training - Reset buffer count after learner consumes rollout to avoid reuse - Add debug logging for collector and buffer operations - Increase first iteration timeout to 300s for collector initialization - Fix distribution_params initialization by manually calling actor forward - Successfully tested: collector collects data, learner consumes and trains
- Calculate env steps from rollout dimensions in learner - Accumulate total env steps in runner - Update logger with env steps via log_collector method - Successfully displays total env steps in training UI
- Add weight_sync_lock parameter to collector function - Pass lock from runner to worker via _start_collector - Use lock when creating SharedWeightSync in subprocess - Fixes potential deadlock in weight synchronization
- Create metrics_queue and pass to collector - Add _drain_metrics method to process collector metrics - Separate wait_time from collect_time in logging - Send periodic metrics from collector worker Note: Collector timing metrics remain 0 due to async architecture, but key metrics (env steps, train time, buffer) work correctly
- Add detailed timing collection in worker (mlp_infer_ms, env_step_total_ms, step_core_ms, update_state_ms, reset_done_ms) - Fix collect_time calculation in runner to match off-policy standard - Optimize learner by passing distribution params from collector, avoiding redundant actor forward pass - Update buffer to store and transfer action_mean and action_sigma - Reduces learner computation and improves GPU utilization Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
The transition.action_mean/action_sigma attributes are not set in collector context. Use actor.output_mean and actor.output_std properties instead. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…port) rsl-rl 5.0.1's RolloutStorage doesn't have mu/sigma fields. PPO.update() internally recomputes distribution params, causing slow training. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Major performance fix: - Add MPS device support in _get_default_device() - Set distribution_params in learner by running actor forward pass - Upgrade local rsl-rl to 5.0.1 (matching PyPI version) Performance improvement: - Before: 10+ seconds (CUDA/CPU fallback) - After: 2.3 seconds (MPS) - Sync PPO baseline: 0.8-1.0 seconds Remaining 2-3x overhead is due to: - H2D data transfer from collector to learner - Recomputing distribution_params (collector on CPU) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Add logger_type parameter support (tensorboard/wandb/none/no_print) - Collect and display reward components from state.info["log"] - Pass reward_components to logger.log_step() - Update _drain_metrics to handle reward_components Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Add device_utils.py with get_default_device() and get_env_dims() - Add SharedBufferBase for common buffer initialization - Add BaseCollector for shared collector patterns - Update all runners to use device_utils - Update buffers to inherit from SharedBufferBase - Fix async_ppo worker to send metrics every rollout Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…tion - runner.py/worker.py: use is_rsl_rl_v5()+convert_config_v5 for rsl-rl 5.x instead of always using convert_config_v3_to_v4; add multi_gpu=None fallback required by PPO.construct_algorithm - buffer.py: add dist_params_dim param; extend packed storage (MPS/CPU) and named tensors (CUDA) to carry distribution_params per rollout - worker.py: save ppo.transition.distribution_params each step into dist_params_buf; wrap all init+loop in try/except so crashes are forwarded to main process via metrics_queue instead of silently hanging - learner.py: _fill_storage reads dist_params from buffer directly (split into mean/std), removing the extra full-batch actor forward pass that was dispatched asynchronously and inflating ppo_update_ms timing
The collector process holds a stale copy of the critic weights. Using stale st.values from the buffer mixed with a fresh last_values from compute_returns produced inconsistent GAE estimates, inflating advantages and causing the surrogate loss to diverge. The policy never converged (reward stayed negative, entropy stuck at init). Fix: after filling storage from the buffer rollout, recompute st.values for all T*N transitions using the current learner critic before calling compute_returns. This ensures a consistent value baseline throughout GAE. Result: surrogate stabilizes near 0 (was diverging), entropy decreases (15.9 vs 17 stuck), mean reward improves (+0.319 vs -0.38) over 151 iters.
… PPO worker Previously dones = state.terminated only. When an episode ends by timeout (truncated=True, terminated=False), the env auto-resets internally but current_episode_lengths/rewards were never reset, causing episode lengths to accumulate unboundedly across episode boundaries (reaching 2500+). Also fixed: timeout_rate was never counted because the timeout check was gated on state.terminated > 0. Fix: - dones_buf (for GAE) stays as state.terminated — truncated episodes correctly bootstrap V(s) at the boundary - episode boundary detection uses state.done (terminated | truncated) - timeout/terminate classification uses state.truncated[i] directly
…PO worker
rsl-rl adds γ·V(s_T) to rewards at truncated episode boundaries so GAE
sees an infinite-horizon return. async PPO was missing this: the env
auto-resets on truncation, so obs_{t+1} in the buffer is the reset obs.
With done=0 (correct for GAE), compute_returns bootstraps V(reset obs)
instead of V(s_T), systematically underestimating truncated returns.
Fix: after computing rewards, add γ·V(s_T) for any truncated env.
V(s_T) is already available as ppo.transition.values from the act() call
just before env.step(). γ is read from ppo.gamma.
Also fix episode reward tracking to use raw env reward (state.reward)
rather than the bootstrapped reward, so logged mean_reward reflects
true undiscounted episode return.
- Remove async_ppo module and train_async_ppo.py - Restore appo module from origin/refactor/appo (V-trace, target network, IMPACT IS correction, clean worker-only-actor architecture) - Add appo_config() to locomotion_params.py - Fix rsl-rl version check in runner.py, worker.py, train_appo.py: is_rsl_rl_v5() now short-circuits before convert_config_v3_to_v4 so appo_config (already v5-native actor/critic format) is not corrupted
…tribution
- Delete APPOActorWrapper: MLPModel (stochastic=True) has identical API
(get_output_log_prob, output_mean, output_std, output_entropy)
- Update appo_config actor to use distribution_cfg v5 format instead of
legacy init_noise_std/noise_std_type/stochastic (fixes TypeError on v5)
- runner._build_learner / worker: deepcopy actor_cfg so MLPModel.__init__'s
distribution_cfg.pop("class_name") doesn't mutate shared rl_cfg
- learner.update: fix stochastic_output=False → True for target_actor KL prep
(MLPModel only populates output_mean/std on the stochastic forward path)
- Export SharedOnPolicyStorage from unilab.ipc.__init__
- Add untracked shared_onpolicy_storage.py to git
Benefits: weight format unified with rsl-rl PPO checkpoints, JIT/ONNX
export compatible, no more hand-rolled std param or torch.distributions.Normal.
…g, drain_metrics) worker.py: - Track ep_timeouts / ep_terminates, report timeout_rate + terminated_rate per window - EMA timing for MLP inference (mlp_infer_ms) and env.step (env_step_total_ms) - Send collector_timing_ms dict in every metrics message runner.py: - Add _drain_metrics() static method (mirrors OffPolicyRunner) — drains full queue, calls update_ep_length / update_done_rates / update_collector_timing / log_collector - Split wait_time (blocking wait for shared storage) from collect_time (read latency) - Pass wait_time to log_step() so terminal shows learner wait colored red if > 1ms - Maintain reward_history deque (200), proper rolling mean for terminal display - Remove dead-code deque()/time.time() assignments learner.py: - Rename return keys: surrogate→surrogate_loss, value_function→value_loss so OffPolicyLogger terminal sorts them under "Losses" section correctly TensorBoard coverage now matches FastSAC: train/surrogate_loss, train/value_loss, train/entropy, train/kl episode/timeout_rate, episode/terminated_rate, episode/length timing/learner_collect_ms, timing/learner_train_ms timing/collector_mlp_infer_ms, timing/collector_env_step_ms perf/steps_per_sec, perf/iter_ms, perf/collect_train_ratio
- APPORunner: add sync_collection (default True) with collection_ready_queue /
trainer_done_queue handshake mirroring FastSAC; collector waits for
trainer_done before each rollout, signals collection_ready after write
- OffPolicyLogger: replace 5 hardcoded _collector_*_ms fields with a single
_collector_timing dict; update_collector_timing does a plain dict.update();
TB/W&B emit timing/collector_{key} for whatever the collector reports;
terminal table renders rows dynamically in pairs — no more zero pollution
from algo-mismatched fields; add timing/learner_wait_ms to TB and W&B
- train_appo.py: CLI numeric args default to None and fall back to appo_config
values, so --max_iterations etc. only override when explicitly specified
- learner: simplify surrogate to standard PPO ratio (drop IMPACT IS term); V-trace handles off-policy correction in advantages, not the surrogate - learner: vectorize vtrace_advantages (numpy backward scan, cat next_values/next_vs) - learner: cache target_actor mu/sigma in process_batch, skip duplicate forward in update - learner: remove unnecessary rewards.clone() (combined is already a fresh cat tensor) - runner: drain all available ring-buffer slots per iteration to keep GPU busy - runner: epoch adjustment increase-only (5% threshold); removing reduce branch - runner: default epoch_adjust_interval 10→50, EMA α 0.05→0.02, warmup 2×interval - worker: pre-alloc obs TensorDict; update in-place, remove redundant astype(float32) - train_appo: add --load_run argument for play mode
Replace standard PPO ratio with min(π_current/π_target, π_current/π_b) = clamp(π_b/π_target, max=1) * π_current/π_b. Fresh data (π_b≈π_target): reduces to standard PPO. Stale data: clipped_rho < 1 downweights the update without amplifying it, matching the V-trace policy gradient (ρ_t * ∇log π * A_vtrace).
added 3 commits
March 16, 2026 04:43
epochs 固定由 config 控制，不再运行时自动调整。 同步删除 min_epochs/max_epochs/epoch_adjust_interval 参数及 --max_epochs CLI flag。 新增 available_on_arrive 到 metrics 方便观察 ring buffer 积压情况。
… queue train_rsl_rl.py play_only fixes: - suppress tensorboard logger in play mode (runner["logger"] = "none") - pass log_dir=None to OnPolicyRunner (avoids persistent play_temp dir) - call env.close() after training before play starts APPO: - runner: call logger.update_replay_queue() each iteration - offpolicy_logger: add replay queue occupancy row to timing panel Env configs: - go1/go2/g1 joystick: expand vel_limit from [0.5,0,0] to [-0.6/-0.4/-0.8, 1.0/0.4/0.8] - locomotion_params: reduce default APPO max_iterations 500→150, save_interval 100→50 - locomotion_params: consolidate G1/G1WalkTask into shared config block
…xist - get_latest_run: skip dirs whose name doesn't start with a digit (filters out 'appo-...', 'play_temp', etc. accidentally in rsl_rl_train/) - play_rsl_rl: validate checkpoint has actor_state_dict before runner.load(); print clear error instead of cryptic KeyError if format mismatch
This was referenced Sep 5, 2026
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.