Dev sonic training performance

[PoplarCrystal]

Summary

This PR systematically optimizes the GEAR-SONIC training pipeline by reducing unnecessary CPU/GPU transfers, device synchronizations, repeated computations, and memory operations during rollouts, environment interaction, and PPO updates.

The main changes are:

• Make Universal Token diagnostic and latent caches configurable.
• Share deterministic observation terms between policy and critic groups.
• Enable fused AdamW.
• Defer episode-statistics transfers until the end of each rollout.
• Check gradient finiteness from the global gradient norm.
• Make the critic evaluation chunk size configurable.
• Make environment action logging and adaptive-sampling diagnostics optional.
• Avoid unnecessary reset synchronization.
• Optimize Isaac Lab circular-buffer updates.
• Batch command-reset metric transfers.
• Make the number of PhysX GPU partitions configurable.

All major optimizations have configuration switches that can restore the previous behavior when required.

Motivation

The previous training path contained several high-frequency, small CPU/GPU synchronizations and repeated computations. Although each operation was relatively inexpensive in isolation, their cumulative cost became significant with many parallel environments and long training runs.

The main sources of overhead were:

1. The Universal Token encoder copied token and latent snapshots to the CPU during every forward pass, even when no consumer used them.
2. Policy and critic observation groups independently computed identical deterministic observation terms.
3. Environment actions could be copied to the CPU at every rollout step.
4. Adaptive-sampling minimum, maximum, and mean statistics were collected at every step.
5. Episode rewards and lengths were transferred from the GPU to the CPU whenever an episode ended.
6. Gradient validation scanned every parameter separately for NaN and Inf values.
7. Command-reset metrics used separate scalar transfers, causing repeated device synchronization.
8. Isaac Lab's circular-buffer path performed avoidable initialization and data-copy operations.
9. The reset path could synchronize the device even when the relevant action-transform buffers did not exist.
10. The default PhysX GPU partition count introduced unnecessary overhead for this training workload.

This PR combines, defers, caches, or disables these operations where appropriate while preserving the training semantics.

Changes

1. Configurable Universal Token caches

The Universal Token model now supports:

cache_encoded_outputs: false
cache_full_latent: false

• cache_encoded_outputs controls CPU debug snapshots of encoded tokens and latents.
• cache_full_latent controls the device-side full-latent cache used by optional smoothness bookkeeping.
• The optimized configurations disable caches that are not required by the normal training path.
• Either cache can be re-enabled when an external callback or loss depends on it.

2. Shared deterministic observation cache

The following option is enabled:

sonic_shared_observation_cache: true

Within one observation-manager compute() call, policy and critic groups share the results of deterministic, parameter-free terms:

actions
base_ang_vel
joint_pos
joint_vel

The cache is cleared after each compute() call and therefore never reuses observations across simulation steps.

3. Fused AdamW

The optimizer is changed to:

optim: adamw_torch_fused

PyTorch's fused AdamW implementation reduces the number of kernel launches during optimizer updates. The previous implementation can be restored with optim: adamw_torch.

4. Deferred episode-statistics transfers

The following option is enabled:

defer_episode_buffer_updates: true

Instead of immediately transferring completed episode rewards and lengths to the CPU at every rollout step, the new path collects them on the GPU and performs a consolidated transfer after the rollout.

5. Fast gradient finite check

The following option is enabled:

fast_gradient_finite_check: true

The previous implementation inspected every parameter gradient separately for NaN and Inf values. The optimized path reuses the global gradient norm returned by gradient clipping and performs a single finite check. If the norm is not finite, gradients are still cleared and the optimizer update is skipped.

6. Configurable critic evaluation chunk size

The critic evaluation chunk size is now configurable:

value_evaluate_chunk_size: 4096

Increasing the value from 1024 to 4096 reduces chunking and dispatch overhead. It can still be adjusted for the available GPU memory.

7. Optional environment action logging

The following option is added:

store_env_actions: false

Environment actions are copied to the CPU and stored in extras["env_actions"] only when explicitly requested. Workflows using callbacks such as MultiLatentSaveCallback can restore the previous behavior by setting this option to true.

8. Optional adaptive-sampling diagnostics

The following option is added:

collect_adaptive_sampling_diagnostics: false

Adaptive sampling continues to operate normally, but per-step minimum, maximum, and mean diagnostic statistics are disabled by default. They can be re-enabled when debugging adaptive sampling.

9. Avoid unnecessary reset synchronization

The reset condition now checks whether the relevant action-transform buffers exist before evaluating reset_mask.any(). This avoids an unnecessary GPU synchronization when the feature is not in use.

10. Faster circular-buffer updates

The following option is enabled:

sonic_fast_circular_buffer: true

The optimized Isaac Lab CircularBuffer path:

• avoids redundant copies during the first append;
• consolidates buffer initialization;
• preserves per-batch reset behavior;
• uses torch.roll to construct the chronological buffer view.

The original Isaac Lab implementation can be restored by disabling this option.

11. Batched command-reset metric transfers

The following option is enabled:

sonic_fast_command_reset: true

Command metrics are stacked before being transferred to the CPU. This replaces multiple scalar transfers with one batched values.cpu().tolist() operation and reduces device synchronization.

12. Configurable PhysX GPU partitions

The environment now supports:

gpu_max_num_partitions: 1

The configured value is propagated to self.sim.physx.gpu_max_num_partitions. This workload uses one partition to reduce partition-management overhead. Setting it to 8 restores the original Isaac Lab default behavior.

Configuration

The relevant optimized settings are:

# gear_sonic/config/base.yaml
sonic_shared_observation_cache: true
sonic_fast_circular_buffer: true
sonic_fast_command_reset: true

# gear_sonic/config/algo/ppo_im_phc.yaml
defer_episode_buffer_updates: true
fast_gradient_finite_check: true
value_evaluate_chunk_size: 4096

# gear_sonic/config/algo/trl/ppo.yaml
optim: adamw_torch_fused

# gear_sonic/config/manager_env/base_env.yaml
store_env_actions: false
collect_adaptive_sampling_diagnostics: false
gpu_max_num_partitions: 1

# Universal Token model configuration
cache_encoded_outputs: false
cache_full_latent: false

Validation

Static validation

• All 12 optimization commits were applied and reviewed independently.
• All newly introduced comments and configuration descriptions are in English.
• git diff --check passes.
• All five modified Python files pass Python AST syntax parsing.
• No unrelated tracked files are included.

Performance validation

The optimized and baseline runs were compared at approximately the same wall-clock time:

Elapsed time: approximately 29.94 hours

Optimized run:
28,556,918,784 total timesteps
approximately 264.93 thousand timesteps/second

Baseline run:
24,647,565,312 total timesteps
approximately 228.66 thousand timesteps/second

The performance improvement is calculated as:

Throughput ratio
= 28,556,918,784 / 24,647,565,312
= 1.1586x

Performance improvement
= (1.1586 - 1) * 100%
= 15.86%

At the same elapsed time, the optimized run completed an additional:

3,909,353,472 timesteps

For this benchmark configuration, training throughput increased from approximately 228.66K timesteps/s to 264.93K timesteps/s, an improvement of approximately 15.9% (1.16x).

This is an observed result for the tested configuration. The exact improvement may vary with the GPU model, number of parallel environments, model configuration, and enabled logging options.

Based on the iteration counts measured at approximately the same elapsed time (48,560 optimized versus 41,937 baseline), completing 100,000 iterations is estimated to decrease from approximately 4.0 days to 3.4 days, saving about 0.54 days (13.0 hours) and reducing the total training time by approximately 13.6%.

This is an observed result for the tested configuration. The exact improvement may vary with the GPU model, number of parallel environments, model configuration, and enabled logging options.

Scope

This PR only optimizes GEAR-SONIC training performance. It does not change:

• policy or critic network architecture;
• checkpoint format;
• observation or action dimensions and ordering;
• reward definitions;
• the PPO objective;
• environment dynamics;
• deployment or ONNX inference interfaces.

Some diagnostic data and debug caches are now disabled by default. Workflows that depend on encoded token/latent CPU snapshots, the full-latent cache, extras["env_actions"], or per-step adaptive-sampling diagnostics must explicitly re-enable the corresponding options.

Co-authored-by: songzhan songzhan@baidu.com
Co-authored-by: bizaorong bizaorong@baidu.com