[BugFix] Stabilize LIBERO VLA GRPO training

Add reference rollout comparison and parity diagnostics for OpenVLA on LIBERO.

Align preprocessing, CPU action decoding, inference microbatching, collector lifecycle, replay grouping, evaluator metrics, and policy synchronization with the validated runtime.

Add paper-shaped configuration controls and focused regression coverage for the end-to-end VLA GRPO path.


ghstack-source-id: a88d47e
Pull-Request: #3940

Author: vmoens

sota-implementations/vla_grpo/README.md

@@ -1,19 +1,16 @@
 # SimpleVLA-RL (arXiv:2509.09674) on LIBERO: OpenVLA-OFT token variant (7B),
-# the paper's hyper-parameters. Parallel MuJoCo workers (one process each)
-# feed a single training device; see the README for the multi-GPU notes and
+# the paper's hyper-parameters. H100 fast mode uses shared policy-server
+# MultiCollector rollout workers; see the README for the multi-GPU notes and
 # the LoRA fallback. Requires LIBERO (and its deps) plus transformers/timm.
 #
-# Per-iteration accounting: groups_per_iter x group_size = 512 trajectories,
+# Per-target-wave accounting: groups_per_iter x group_size = 320 trajectories,
 # each up to max_outer_steps = 64 chunk decisions (512 env steps / chunk 8).
 
 env:
   backend: libero
   task_suite: libero_spatial # libero_spatial | libero_object | libero_goal | libero_10
   task_ids: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] # one worker per task (num_envs must cover them)
-  # Parallel MuJoCo processes (training); >= len(task_ids). For GRPO, each
-  # worker owns serial group ids, so keep num_envs <= groups_per_iter (ideally
-  # a divisor, often equal) or no worker may finish a full group per iteration.
-  num_envs: 10
+  num_envs: 80
   # false: each worker replays group_size rollouts serially for one group.
   # true: group_size workers share the same task/init state/group id and
   # collect a GRPO group in parallel. This improves group-completion throughput
@@ -26,24 +23,28 @@ env:
   # batches can overcollect under one policy, but without a group barrier they
   # may also drift; the LIBERO wrapper groups parallel repeats by cycled
   # init-state id to reduce this waste.
-  parallel_group_repeats: false
+  parallel_group_repeats: true
   eval_num_envs: 10 # parallel MuJoCo processes (evaluation); >= len(task_ids)
   camera_height: 256
   camera_width: 256
   render_backend: egl # egl (GPU/headless) | osmesa (CPU fallback)
-  render_gpu_ids: [0] # EGL-visible ids; override to [0,1,...] to spread workers
-  eval_render_gpu_ids: null # null = reuse render_gpu_ids for eval workers
+  render_gpu_ids: [2, 3, 4, 5] # H100 fast mode render GPUs for rollout workers
+  eval_render_gpu_ids: [7] # reserve GPU 7 for eval/video rendering
+  render_gpu_device_zero_fallback: true
+  env_kwargs: null
   chunk_size: 8 # NUM_ACTIONS_CHUNK of the checkpoint
   max_env_steps: 512 # base env steps per episode (the paper's cap)
   max_outer_steps: 64 # = max_env_steps / chunk_size, in chunk decisions
-  train_init_state_mode: random # random | cycle | fixed; eval always cycles
+  train_init_state_mode: cycle # random | cycle | fixed; eval always cycles
+  train_init_state_id: 0 # selected LIBERO init state when train_init_state_mode=fixed
   seed: 0
 
 tokenizer:
   vocab_size: 256 # unused for the openvla backend (codec comes from the checkpoint)
 
 policy:
   backend: openvla
+  mode: tokens # tokens: SimpleVLA-RL GRPO path; l1: official continuous OFT reference/eval path
   # the checkpoint MUST match task_suite (it fixes the SFT policy AND the
   # action statistics unnorm_key resolves into)
   checkpoint: Haozhan72/Openvla-oft-SFT-libero-spatial-traj1
@@ -56,58 +57,54 @@ policy:
   temperature: 1.6 # rollout sampling temperature (greedy eval)
   top_k: null # null = full categorical; small values keep token exploration local
   use_wrist_image: false
+  # L1 reference path: set mode=l1, checkpoint=moojink/openvla-7b-oft-finetuned-libero-spatial,
+  # use_wrist_image=true, use_proprio=true, num_images_in_input=2.
+  action_head_file: action_head--150000_checkpoint.pt
+  proprio_projector_file: proprio_projector--150000_checkpoint.pt
+  use_proprio: false
+  num_images_in_input: 1
   center_crop: true # the SimpleVLA-RL checkpoints are trained with augmentations
-  image_backend: torchvision # OpenVLA image preprocessing backend: torchvision | pil
+  image_backend: torchvision # OpenVLA preprocessing: torchvision | pil | tensorflow (reference)
   # the model emits a continuous gripper; LIBERO/robosuite needs a firm +/-1
   # open/close (demo gripper is binary, -1 open / +1 close). SimpleVLA-RL
-  # applies sign(2*g - 1) then inverts, i.e. a 0.5 binarization threshold
-  # before the open/close sign flip.
+  # applies sign(2 * g - 1) then inverts, so threshold 0.0 here is equivalent
+  # to threshold 0.5 on the raw decoded gripper scalar.
   gripper_binarize: true
-  gripper_binarize_threshold: 0.5
+  gripper_binarize_threshold: 0.0
   gripper_invert: true
   lora_rank: 32 # de-risk fallback; set to 0/null for full fine-tuning (needs FSDP)
   lora_target_modules: [q_proj, k_proj, v_proj, o_proj]
-  device: null # null = auto (cuda if available)
+  device: cuda:0 # training device
 
 collector:
   group_size: 8 # n: rollouts per initial state (GRPO group)
   candidate_group_size: null # null = group_size; e.g. 16 samples 16 and selects 8
-  # Initial states per iteration (512 trajectories). Must be >= env.num_envs;
-  # for best throughput, use a multiple of env.num_envs so no partial groups
-  # are discarded at the synchronous update boundary. If
-  # env.parallel_group_repeats=true, use env.num_envs / group_size instead.
-  groups_per_iter: 64
-  # Safety cap in target group waves. The collector writes complete
-  # trajectories directly to the replay buffer every one-outer-step internal
-  # poll; one wave is approximately max_outer_steps *
-  # ceil(groups_per_iter * group_size / num_envs) polls. Keeping MCAdvantage
-  # queues across these same-policy polls lets partial GRPO groups finish
-  # without crossing a policy update.
-  max_collect_batches_per_iter: 1
-  # If a capped wave produces too few useful replay decisions, keep the same
-  # rollout policy and queued partial groups for another capped wave instead
-  # of updating on a tiny replay batch and clearing those queues immediately.
-  max_same_policy_collect_attempts: 2
-  min_replay_decisions: null # null/0 = collect one target group wave
+  # Initial states per iteration (320 trajectories). With shared fast mode and
+  # env.parallel_group_repeats=true, use
+  # num_collectors * envs_per_collector / group_size.
+  groups_per_iter: 40
+  # null/0 = one full selected rollout set:
+  # groups_per_iter * group_size * max_outer_steps decisions.
+  min_replay_decisions: null
   total_iters: 100 # the paper's total_epochs
-  policy_device: null # null = policy.device; set e.g. cuda:1 for rollout inference
-  # Execution-mode switches for throughput ablations:
-  # - false/false: regular synchronous TorchRL Collector.
-  # - true/false: async env slots with one request per policy forward.
-  # - true/true: async env slots plus auto-batched policy inference.
-  # - false/true: sync env stepping through the policy server path.
-  async_env: false
-  async_policy: false
-  env_backend: threading # AsyncBatchedCollector env backend: threading | multiprocessing
-  policy_backend: threading # inference transport: threading | multiprocessing | ray | monarch
-  server_backend: thread # process server needs a policy_factory and is not used here
-  server_max_batch_size: null # null = env.num_envs when async_policy=true
+  policy_device: cuda:1 # H100 default: rollout/eval inference server device
+  # null = one model forward for the complete request batch. Set to 1 for
+  # single-environment numerical parity while keeping batched env collection.
+  policy_micro_batch_size: null
+  replay_wait_s: 0.5
+  replay_log_s: 30.0
+  num_collectors: 4
+  envs_per_collector: 80
+  num_threads: 1
+  server_max_batch_size: 1 # one batched ParallelEnv request; avoids mixed eval/rollout modes
   server_min_batch_size: 1
-  server_timeout: 0.01
+  server_timeout: 0.001
   server_collect_stats: true
   server_stats_window_size: 1024
   max_inflight_per_env: 1
-  storing_device: null
+  env_sub_threads: 1
+  storing_device: cpu
+  use_buffers: null
 
 advantage:
   trajectory_return: sum # binary success return per trajectory
@@ -120,6 +117,9 @@ advantage:
 
 buffer:
   device: cpu # decisions hold raw uint8 images; keep the buffer off-GPU
+  shared_init: true
+  capacity_group_waves: 7 # replay capacity as multiples of one target group wave
+  consume_after_n_samples: 1 # sampled decisions leave the buffer after one update pass
 
 loss:
   clip_epsilon: [0.2, 0.28] # asymmetric DAPO Clip-Higher (per-token, see ratio_level)
@@ -142,12 +142,11 @@ logger:
   mode: online
   eval_iter: 4 # evaluate every N iterations (the paper's cadence)
   eval_episodes: 50 # greedy episodes per evaluation (cycled init states)
-  eval_process: false # true = run eval in a dedicated process
-  eval_async: false # with eval_process, submit eval and log result later
-  eval_device: null # null = policy.device/auto; set e.g. cuda:7 with eval_process
-  eval_state_dir: null # null = tempfile dir for trainable-weight snapshots
-  record_video: true # render eval rollouts via VideoRecorder (needs a logger)
-  video_episodes: 2 # greedy episodes per recorded video (one camera stream)
+  eval_backend: process # TorchRL Evaluator backend: thread | process | ray
+  eval_async: false # submit eval and log result later
+  eval_busy_policy: skip # skip | error | queue
+  eval_timeout_s: 1800.0
+  eval_device: null # null = rollout device; set e.g. cuda:7 for eval
   video_fps: 8 # frames per second of the logged video (one frame per decision)
 
 checkpoint:

sota-implementations/vla_grpo/compare_simplevla_rollouts.py

@@ -11,7 +11,15 @@ env:
   success_tol: 0.35 # sized so a random policy succeeds sometimes (cold-start signal)
   max_outer_steps: 6 # episode truncation, in chunk decisions
   render_size: 64 # side length of the from_pixels eval-video frame
-  num_envs: 1 # async-env workers; ToyVLAEnv grouped rollouts are per worker
+  num_envs: 1
+  eval_num_envs: 1
+  parallel_group_repeats: false
+  train_init_state_mode: random
+  render_backend: null
+  render_gpu_ids: null
+  eval_render_gpu_ids: null
+  render_gpu_device_zero_fallback: true
+  env_kwargs: null
   seed: 0
 
 tokenizer:
@@ -26,33 +34,24 @@ collector:
   group_size: 4 # n: rollouts per initial state (GRPO group)
   candidate_group_size: null # null = group_size; larger values oversample then select group_size
   groups_per_iter: 8 # initial states per iteration
-  # Safety cap in target group waves. The collector writes complete
-  # trajectories directly to the replay buffer every one-outer-step internal
-  # poll; one wave is approximately max_outer_steps *
-  # ceil(groups_per_iter * group_size / num_envs) polls.
-  max_collect_batches_per_iter: 1
-  # Retry a capped wave once under the same policy if too few useful replay
-  # decisions have reached the replay buffer.
-  max_same_policy_collect_attempts: 2
   min_replay_decisions: null # null/0 = collect one target group wave
   total_iters: 200
-  # Execution-mode switches for throughput ablations:
-  # - false/false: regular synchronous TorchRL Collector.
-  # - true/false: async env slots with one request per policy forward.
-  # - true/true: async env slots plus auto-batched policy inference.
-  # - false/true: sync env stepping through the policy server path.
-  async_env: false
-  async_policy: false
-  env_backend: threading # AsyncBatchedCollector env backend: threading | multiprocessing
-  policy_backend: threading # inference transport: threading | multiprocessing | ray | monarch
-  server_backend: thread # process server needs a policy_factory and is not used here
-  server_max_batch_size: null # null = number of async envs when async_policy=true
+  replay_wait_s: 0.05
+  replay_log_s: 30.0
+  policy_device: null
+  policy_micro_batch_size: null
+  num_collectors: 1
+  envs_per_collector: 1
+  num_threads: 1
+  env_sub_threads: 1
+  server_max_batch_size: 1
   server_min_batch_size: 1
   server_timeout: 0.01
   server_collect_stats: true
   server_stats_window_size: 1024
   max_inflight_per_env: 1
-  storing_device: null
+  storing_device: cpu
+  use_buffers: null
 
 advantage:
   trajectory_return: sum # binary success return per trajectory
@@ -63,6 +62,9 @@ advantage:
 
 buffer:
   device: null # null = training device
+  shared_init: false
+  capacity_group_waves: 1 # replay capacity as multiples of one target group wave
+  consume_after_n_samples: 1 # sampled decisions leave the buffer after one update pass
 
 loss:
   clip_epsilon: [0.2, 0.28] # asymmetric DAPO clip; a float gives symmetric clipping
@@ -85,14 +87,13 @@ logger:
   mode: online
   eval_iter: 10 # evaluate every N iterations
   eval_episodes: 50 # greedy episodes per evaluation
-  eval_process: false # true = run eval in a dedicated process
-  eval_async: false # with eval_process, submit eval and log result later
-  eval_device: null # null = policy.device/auto; set e.g. cuda:7 with eval_process
-  eval_state_dir: null # null = tempfile dir for trainable-weight snapshots
-  record_video: true # render eval rollouts via VideoRecorder (needs a logger)
-  video_episodes: 2 # greedy episodes per recorded video
+  eval_backend: thread # TorchRL Evaluator backend: thread | process | ray
+  eval_async: false # submit eval and log result later
+  eval_busy_policy: skip # skip | error | queue
+  eval_timeout_s: 300.0
+  eval_device: null # null = rollout device; set e.g. cuda:7 for eval
   video_fps: 4 # frames per second of the logged video (one frame per decision)
 
 checkpoint:
   save_iter: 25 # save every N iterations (0 disables)
-  resume: null # path to a checkpoint_latest.pt to resume from
+  resume: null # path to a checkpoint_latest TensorDict directory to resume from