diff --git a/sota-implementations/vla_grpo/README.md b/sota-implementations/vla_grpo/README.md index bda1af31c4e..cdddd67f13c 100644 --- a/sota-implementations/vla_grpo/README.md +++ b/sota-implementations/vla_grpo/README.md @@ -166,8 +166,8 @@ vendored modeling code under `openvla_oft/` comes from [SimpleVLA-RL](https://github.com/PRIME-RL/SimpleVLA-RL) (MIT). Important compatibility note: the official continuous-head OpenVLA-OFT -checkpoints are not interchangeable with this token-head variant. Use the -SimpleVLA-RL SFT checkpoints, for example `Haozhan72/*`: +checkpoints are not interchangeable with this token-head variant for GRPO +training. Use the SimpleVLA-RL SFT checkpoints, for example `Haozhan72/*`: ```python from openvla import OpenVLAOFTWrapper @@ -196,29 +196,47 @@ pytest sota-implementations/vla_grpo/test_openvla.py ## What gets logged With a logger configured (`logger.backend=wandb`, the default), each iteration -logs reward curves (`train/reward_mean`, `train/reward_max`), success rate, and -throughput split into collection and optimization: +logs the training success rate (`train/success_rate`), trajectory-return +aggregates (`collector/trajectory_return_sum`, +`collector/trajectory_return_max`), and throughput split into collection and +optimization: - `throughput/inference_env_steps_per_s` - `throughput/inference_decisions_per_s` - `throughput/train_decisions_per_s` - `throughput/optim_steps_per_s` -Eval rollouts can also be rendered to video (`logger.record_video=true`, on by -default). A dedicated single-environment recorder is built with -`from_pixels=True`: `ToyVLAEnv` renders the tracking scene, while `LiberoEnv` -exposes its camera. `torchrl.record.VideoRecorder` writes -`logger.video_episodes` greedy episodes to `eval/video` on every eval. wandb -video encoding needs `moviepy` from the `dev` dependency group. Disable videos -with `logger.record_video=false`. +The collector path is fixed: a TorchRL `MultiCollector` launches rollout +workers, each worker owns a sync `ParallelEnv(envs_per_collector)`, and all +workers plus the evaluator share one process policy server. + +```bash +python sota-implementations/vla_grpo/vla-grpo.py --config-name vla_grpo_libero +``` + +Rollout clients request random sampling; eval and video clients request +deterministic decoding from the same server, so rollout and eval are synced by +one explicit TensorDict weight update after each optimizer step. The replay +buffer is passed to the collector and receives complete trajectories as they +finish; training waits until the consuming replay buffer has enough sampleable +decisions. -Checkpointing is shared by the toy and LIBERO configs. `checkpoint_latest.pt` -is written to the hydra run directory every `checkpoint.save_iter` iterations; +With the thread evaluator backend, eval rollouts are rendered to video whenever +a logger is configured. A dedicated single-environment evaluator is built with +`from_pixels=True`: `ToyVLAEnv` renders the tracking scene, while `LiberoEnv` +exposes its camera. `torchrl.record.VideoRecorder` writes the evaluator rollout +to `eval/video` on every eval. wandb video encoding needs `moviepy` from the +`dev` dependency group. Process-backend evaluator video dumping still needs a +TorchRL-side remote `VideoRecorder.dump` path. + +Checkpointing is shared by the toy and LIBERO configs. `checkpoint_latest` +is written as a TensorDict directory in the hydra run directory every +`checkpoint.save_iter` iterations; resume with: ```bash python sota-implementations/vla_grpo/vla-grpo.py \ - checkpoint.resume=/path/to/checkpoint_latest.pt + checkpoint.resume=/path/to/checkpoint_latest ``` ## LIBERO configuration details @@ -226,8 +244,10 @@ python sota-implementations/vla_grpo/vla-grpo.py \ The full LIBERO config follows the SimpleVLA-RL hyper-parameter shape: - groups of `n=8` rollouts per initial state; -- 64 initial states per iteration, for 512 trajectories before dynamic - filtering; +- 40 initial states per iteration (`collector.groups_per_iter`), for 320 + trajectories before dynamic filtering -- one aligned group wave across the + 320 rollout envs; the paper uses 64 initial states (512 trajectories) per + iteration, which is a known deviation of the shipped config; - 512 base environment steps, or 64 chunk decisions, per episode; - rollout temperature 1.6 and greedy evaluation; - dynamic sampling bounds `(0.1, 0.9)` to drop groups that are all failure or @@ -241,47 +261,53 @@ A sequence-level ratio remains available as a config switch for ablations, but for a 56-token action chunk it saturates the clip range much more easily than per-token ratios. -LIBERO simulation runs in parallel worker processes (`env.num_envs`, one MuJoCo -instance each), and policy inference batches across workers. Group accounting is -the main thing to keep in mind: GRPO needs repeated attempts from the same -initial state under the same policy. Each worker owns a disjoint `group_id` -block so advantages never mix across unrelated groups. - -Because groups are repeated serially within each worker, `env.num_envs` should -not exceed `collector.groups_per_iter`; otherwise many same-policy collection -polls are needed before each worker can finish all `group_size` rollouts for a -group and the replay buffer receives advantaged decisions. For best throughput, -set `env.num_envs` to a divisor of `collector.groups_per_iter`, often the same -value. - -When `env.parallel_group_repeats=true`, `env.num_envs / collector.group_size` -logical workers each run one repeated-initial-state group in parallel. In this -mode, prefer setting `collector.groups_per_iter` to that logical worker count -so one target group wave is aligned. If `collector.candidate_group_size` is -larger than `collector.group_size`, each worker in a logical group repeats the -same initial state serially enough times to produce up to the requested -candidate count. For example, 8 parallel workers x 2 serial repeats gives at -most 16 candidates. Groups can be written earlier if the candidates already -contain a useful selected subset. - -The replay-buffer writer polls the collector at one outer step per worker, so -complete trajectories are handed to the replay buffer shortly after they finish -instead of waiting for a full max-length rollout from every worker. +LIBERO simulation runs through `collector.num_collectors` MultiCollector +workers. Each worker hosts a synchronous +`ParallelEnv(collector.envs_per_collector)`. Policy inference runs on the +shared process server and each worker owns a disjoint `group_id` block so +advantages never mix across unrelated groups. + +Without `env.parallel_group_repeats`, groups are repeated serially within each +worker, so the total rollout worker count should not exceed +`collector.groups_per_iter`. For that serial mode, set +`collector.num_collectors * collector.envs_per_collector` to a divisor of +`collector.groups_per_iter`, often the same value. + +When `env.parallel_group_repeats=true`, the shared replay buffer centralizes +`MCAdvantage` write state, so same-initial-state groups may straddle +subcollectors. The logical worker count is the total rollout worker count +divided by `collector.group_size`. In this mode, prefer setting +`collector.groups_per_iter` to that logical worker count so one target group +wave is aligned. If `collector.candidate_group_size` is larger than +`collector.group_size`, each worker in a logical group repeats the same initial +state serially enough times to produce up to the requested candidate count. For +example, 8 parallel workers x 2 serial repeats gives at most 16 candidates. +Groups can be written earlier if the candidates already contain a useful +selected subset. + +The training script starts the collector once, waits until the consuming replay +buffer has enough sampleable decisions, pauses collection, runs the PPO update, +clears incomplete same-policy advantage queues and partial trajectories, pushes +the TensorDict policy weights to the shared policy server, and then lets the +collector resume. `MCAdvantage` runs as the replay-buffer transform and keeps incomplete groups -queued across same-policy polls until all siblings arrive. -`max_collect_batches_per_iter` sets the safety cap in target group waves, and -`collector.min_replay_decisions` can require a minimum number of useful replay -decisions before the PPO update. +queued only within a single policy window. `collector.min_replay_decisions` can +require a minimum number of useful replay decisions before the PPO update. +Set `TORCHRL_MC_ADVANTAGE_LOCAL_QUEUES=1` to keep grouping state in each replay +writer instead of a multiprocessing manager. At every policy boundary the +trainer reads those worker-local counters while collection is paused, clears +their queues, and resets in-flight collector trajectories before the policy +version advances. Candidate selection is delegated to `MCAdvantageSelector` (`first`, `uniform`, or `balanced`), so the replay-buffer transform owns the sample-selection policy -while the collector only supplies same-policy completed trajectories. At the -policy-update boundary the replay buffer, incomplete advantage queues, and -in-flight collector trajectories are cleared before the next policy is rolled -out. LIBERO workers stamp parallel-repeat group ids from the cycled -initial-state id so fast and slow sibling workers can still complete a -same-initial-state GRPO group under the same policy even when their episode -lengths differ. +while the collector only supplies same-policy completed trajectories. The +consuming replay buffer removes sampled decisions after +`buffer.consume_after_n_samples` samples, and the policy-boundary pause keeps +rollout and optimization phases explicit. LIBERO workers stamp +parallel-repeat group ids from the cycled initial-state id so fast and slow +sibling workers can still complete a same-initial-state GRPO group under the +same policy even when their episode lengths differ. Run the LIBERO recipe with: @@ -295,31 +321,82 @@ Requirements beyond the toy scale: LIBERO (see the `torchrl.envs.LiberoEnv` docs for install notes), `transformers`, `timm`, `Pillow`, and `peft` when `policy.lora_rank` is set. +For reference-parity rollouts, set `policy.image_backend=tensorflow`. This uses +the SimpleVLA JPEG, Lanczos resize, and center-crop order. Normalized +vocabulary-tail action tokens are detokenized through the NumPy float64 CPU +path before the gripper transform is applied once in the environment. Use +`env.train_init_state_mode=fixed env.train_init_state_id=` for a fixed +LIBERO initial state. `collector.policy_micro_batch_size` only slices actual +model calls inside the inference-server policy; it does not change PPO +minibatching. + ## Hardware notes -- The default configuration trains a LoRA adapter (`policy.lora_rank: 32`) on a - single GPU while the simulation workers occupy CPU cores. Rollout wall-clock - dominates, so scale `env.num_envs` with the available cores first, while - keeping it within the GRPO grouping constraint above. -- Set `collector.policy_device` to a different CUDA device to keep rollout - inference on a separate policy replica. The training loop copies only the - trainable state dict after optimizer updates, so this split is intended for - LoRA/adapters rather than full-parameter fine-tuning. +- The default H100 configuration trains a LoRA adapter + (`policy.lora_rank: 32`) on `policy.device: cuda:0` and serves rollout plus + evaluator inference from `collector.policy_device: cuda:1`. Four + collectors each run `ParallelEnv(80)` for 320 rollout envs total. On + single-GPU runs, override `policy.device=null`, + `collector.policy_device=null`, `collector.num_collectors=1`, and + `collector.envs_per_collector` to the number of local envs. +- Rollout wall-clock dominates, so scale `collector.num_collectors` and + `collector.envs_per_collector` with the available CPU cores while keeping + the GRPO grouping constraint above. The H100 default uses + `collector.num_collectors=4`, `collector.envs_per_collector=80`, + `collector.groups_per_iter=40`, and parallel group repeats enabled. The + training loop pushes TensorDict policy weights to the shared policy server + after optimizer updates. - Headless LIBERO rendering uses MuJoCo/robosuite EGL by default (`env.render_backend: egl`). `env.render_gpu_ids` controls the EGL-visible - render device ids assigned to workers, round-robin. The default `[0]` works - on a single-GPU allocation; on a multi-GPU node, override it, for example - `env.render_gpu_ids=[0,1,2,3]`, to spread render workers across GPUs. These - ids are the devices visible to EGL inside the process/container and may not - match global CUDA ordinals. -- Set `logger.eval_process=true` to move greedy eval into a dedicated process. - Use `logger.eval_device` for its policy device and `env.eval_render_gpu_ids` - for its EGL render workers; when the latter is left null, eval reuses - `env.render_gpu_ids`. + render device ids assigned to rollout workers, round-robin. The H100 default + spreads rollout rendering over `[2,3,4,5]` and reserves + `env.eval_render_gpu_ids=[7]` for eval/video rendering. These ids are the + devices visible to EGL inside the process/container and may not match global + CUDA ordinals. +- Use `logger.eval_backend` for the TorchRL evaluator backend. The evaluator + shares the same policy server as rollout and uses `env.eval_render_gpu_ids` + for EGL rendering; when the latter is left null, eval reuses + `env.render_gpu_ids`. The LIBERO default uses `process` to isolate simulator + work; use `thread` only when local VideoRecorder dumping is required. Note + the caveat: with a logger, the thread backend swaps the eval env for a + single-env video recorder bound to the first task, so on a multi-task suite + `eval/success_rate` covers one task instead of the whole suite (and + `env.eval_num_envs` is ignored). This requires the explicit opt-in + `logger.record_video_single_task=true`. - Minimal CUDA containers often lack the NVIDIA EGL/GLVND userspace stack. Before debugging TorchRL, verify that `libEGL_nvidia`, `libnvidia-eglcore`, `libGLX_nvidia`, and `/usr/share/glvnd/egl_vendor.d/10_nvidia.json` are visible in the runtime. +- On an H200 container with the 595 driver, the userspace libraries can be + extracted without installing Debian packages into the image: + + ```bash + mkdir -p /opt/nvidia-595-deb/download /opt/nvidia-595-deb/extract + cd /opt/nvidia-595-deb/download + apt-get download \ + libnvidia-gl-595 libegl1 libglvnd0 libopengl0 libgl1 libgles2 libglx0 + for deb in ./*.deb; do + dpkg-deb -x "$deb" /opt/nvidia-595-deb/extract + done + + export LIBDIR=/opt/nvidia-595-deb/extract/usr/lib/x86_64-linux-gnu + export LD_LIBRARY_PATH="$LIBDIR:${LD_LIBRARY_PATH:-}" + export LD_PRELOAD="$LIBDIR/libOpenGL.so.0${LD_PRELOAD:+:$LD_PRELOAD}" + export __EGL_VENDOR_LIBRARY_FILENAMES=/opt/nvidia-595-deb/extract/usr/share/glvnd/egl_vendor.d/10_nvidia.json + export MUJOCO_GL=egl PYOPENGL_PLATFORM=egl ROBOT_PLATFORM=LIBERO + + python - <<'PY' + from OpenGL import EGL, GL + import mujoco + from libero.libero import benchmark + + assert EGL is not None and GL.glGetError is not None + assert mujoco is not None and benchmark is not None + PY + ``` + + The validated parity runtime pins `mujoco==3.2.3`, `robosuite==1.4.1`, + `transformers==4.40.1`, and `peft==0.11.1`. - Full-parameter fine-tuning of the 7B model requires sharded training (FSDP) and a multi-GPU inference/training split with explicit weight synchronization. That topology should be sized on the target hardware: diff --git a/sota-implementations/vla_grpo/config/vla_grpo_libero.yaml b/sota-implementations/vla_grpo/config/vla_grpo_libero.yaml index d9c90e8881e..1dae8961e97 100644 --- a/sota-implementations/vla_grpo/config/vla_grpo_libero.yaml +++ b/sota-implementations/vla_grpo/config/vla_grpo_libero.yaml @@ -1,42 +1,47 @@ # 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 + task_ids: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] # one rollout worker per task; the collector env total must cover them + # Only consumed by the standalone make_env tooling path (e.g. probe scripts); + # the training env count is collector.num_collectors * collector.envs_per_collector + # and evaluation uses env.eval_num_envs. + 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 - # but requires num_envs to be a multiple of collector.group_size and, to cover - # all tasks in one wave, num_envs >= len(task_ids) * collector.group_size. + # but requires the collector env total (collector.num_collectors * + # collector.envs_per_collector) to be a multiple of collector.group_size and, + # to cover all tasks in one wave, that total to be at least + # len(task_ids) * collector.group_size. # For lowest boundary waste, set collector.groups_per_iter to the logical - # worker count (num_envs / collector.group_size). Asking each logical worker - # to advance through multiple group ids in one collector batch creates - # partial groups when episodes have variable lengths. Additional collector - # 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 + # worker count (collector env total / collector.group_size). Asking each + # logical worker to advance through multiple group ids in one collector batch + # creates partial groups when episodes have variable lengths. Additional + # collector 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: 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: @@ -44,6 +49,7 @@ tokenizer: 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,41 +62,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 + 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.001 + server_collect_stats: true + server_stats_window_size: 1024 + max_inflight_per_env: 1 + env_sub_threads: 1 + storing_device: cpu + use_buffers: null advantage: trajectory_return: sum # binary success return per trajectory @@ -103,11 +122,15 @@ 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) ratio_level: token # the paper's semantics: one clipped ratio per action token - ppo_epochs: 1 + # epochs are controlled by buffer.consume_after_n_samples: each decision + # leaves the buffer after that many sampling passes mini_batch_size: 16 # decisions per micro-batch (one forward) accumulate_batches: 8 # optimizer step every N micro-batches (grad accumulation) @@ -117,6 +140,13 @@ optim: warmup_updates: 10 # constant-with-warmup, in optimizer steps max_grad_norm: 1.0 +train: + # true: push only requires_grad=True parameters (the LoRA adapters) to the + # rollout policy server on each sync; the server applies the partial + # TensorDict in place, so frozen base weights stay untouched. false ships + # the full parameter/buffer set (~14 GB in bf16 for the 7B policy). + weight_sync_trainable_only: true + logger: backend: wandb project_name: torchrl-sota-vla-grpo @@ -125,13 +155,16 @@ 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) + # eval_backend=thread with a logger records video from a single env bound to + # task_ids[0]; on a multi-task suite this silently narrows eval/success_rate + # to one task, so it requires this explicit opt-in. + record_video_single_task: false checkpoint: save_iter: 25 diff --git a/sota-implementations/vla_grpo/config/vla_grpo_toy.yaml b/sota-implementations/vla_grpo/config/vla_grpo_toy.yaml index e9a85d1d920..c42d66c6b3e 100644 --- a/sota-implementations/vla_grpo/config/vla_grpo_toy.yaml +++ b/sota-implementations/vla_grpo/config/vla_grpo_toy.yaml @@ -11,6 +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 + 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: @@ -25,16 +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 + 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: cpu + use_buffers: null advantage: trajectory_return: sum # binary success return per trajectory @@ -44,12 +61,19 @@ advantage: candidate_selection_max_combinations: 100000 buffer: - device: null # null = training device + # The collector's worker-init path shares the buffer as a CPU memmap, so + # a GPU buffer device cannot be honored; keep it off-GPU (the update loop + # moves sampled batches to the training device). + device: cpu + 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 ratio_level: sequence # sequence: one ratio per decision; token: one per action token - ppo_epochs: 1 + # epochs are controlled by buffer.consume_after_n_samples: each decision + # leaves the buffer after that many sampling passes mini_batch_size: 64 # decisions per micro-batch accumulate_batches: 1 # optimizer step every N micro-batches @@ -59,6 +83,13 @@ optim: warmup_updates: 1 # linear LR warmup, in optimizer steps (1 = no warmup) max_grad_norm: 1.0 +train: + # true: push only requires_grad=True parameters to the rollout policy server + # on each sync (partial in-place update on the receiving side); false ships + # the full parameter/buffer set. TinyVLA trains all parameters, so this only + # skips buffers at toy scale; it matters for the LoRA LIBERO config. + weight_sync_trainable_only: true + logger: backend: wandb project_name: torchrl-sota-vla-grpo @@ -67,14 +98,17 @@ 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) + # eval_backend=thread with a logger records video from a single env bound to + # the first task; only relevant for multi-task suites (the toy env has one + # task, so recording never narrows evaluation here). + record_video_single_task: false 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 diff --git a/sota-implementations/vla_grpo/openvla.py b/sota-implementations/vla_grpo/openvla.py index 01747b450e8..3b5b456aa49 100644 --- a/sota-implementations/vla_grpo/openvla.py +++ b/sota-implementations/vla_grpo/openvla.py @@ -2,17 +2,16 @@ # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -"""OpenVLA-OFT (token variant) policy wrapper for the VLA GRPO recipe. +"""OpenVLA-OFT policy wrappers for the VLA GRPO recipe. -Wraps the vendored SimpleVLA-RL token-OFT model (see ``openvla_oft/``) as a -:class:`~torchrl.modules.vla.VLAWrapperBase`: one forward emits the whole -action chunk (parallel decoding -- ``chunk_size * action_dim`` tokens in a -single language-model pass, no autoregressive loop), sampled from the 256-way -categorical over the action-token window at the tail of the LLaMA-2 -vocabulary. +The token wrapper covers the vendored SimpleVLA-RL token-OFT model (see +``openvla_oft/``): one forward emits the whole action chunk (parallel decoding +-- ``chunk_size * action_dim`` tokens in a single language-model pass, no +autoregressive loop), sampled from the 256-way categorical over the +action-token window at the tail of the LLaMA-2 vocabulary. -The wrapper owns ALL model-side preprocessing (prompt construction, image -transforms) and the temperature scaling, and both rollout sampling and +The wrappers own ALL model-side preprocessing (prompt construction, image +transforms) and, for token heads, the temperature scaling. Both rollout and loss-time recomputation go through the same :meth:`get_dist`, so with identical weights the PPO importance ratio is exactly 1 -- the temperature contract a T != 1 rollout requires. @@ -21,15 +20,24 @@ import importlib.util import io +import json +import os from typing import Literal import numpy as np import torch from tensordict import TensorDictBase -from tensordict.nn import InteractionType +from tensordict.nn import InteractionType, TensorDictSequential +from tensordict.utils import NestedKey, unravel_key from torch.nn.utils.rnn import pad_sequence -from torchrl.data.vla import OpenVLAImagePreprocessor, VocabTailActionTokenizer +from torchrl.data.vla import ( + ACTION_CHUNK_KEY, + OpenVLAImagePreprocessor, + VocabTailActionTokenizer, +) +from torchrl.envs.transforms import ActionTokenizerTransform +from torchrl.envs.transforms._base import Transform from torchrl.modules.vla import VLAWrapperBase from torchrl.modules.vla.common import LogProbsMode @@ -37,6 +45,8 @@ _has_transformers = importlib.util.find_spec("transformers") is not None _has_timm = importlib.util.find_spec("timm") is not None _has_pil = importlib.util.find_spec("PIL") is not None +_has_huggingface_hub = importlib.util.find_spec("huggingface_hub") is not None +_hf_hub_download = None # the special empty token LLaMA-2 emits after "Out:" at training time _EMPTY_TOKEN_ID = 29871 @@ -59,6 +69,20 @@ # pixels, costing closed-loop precision. _OPENVLA_IMAGE_SIZE = 224 _JPEG_QUALITY = 95 +_OPENVLA_INPUT_IDS_KEY = ("openvla", "input_ids") +_OPENVLA_ATTENTION_MASK_KEY = ("openvla", "attention_mask") +_OPENVLA_PIXEL_VALUES_KEY = ("openvla", "pixel_values") + + +def _hf_hub_download_fn(): + global _hf_hub_download + if not _has_huggingface_hub: + raise ImportError("Hugging Face Hub is required to download OpenVLA weights.") + if _hf_hub_download is None: + from huggingface_hub import hf_hub_download + + _hf_hub_download = hf_hub_download + return _hf_hub_download def _load_dataset_statistics(spec: str) -> dict: @@ -67,19 +91,532 @@ def _load_dataset_statistics(spec: str) -> dict: ``spec`` is either a filesystem path to the json or a HuggingFace repo id that ships a ``dataset_statistics.json`` at its root. """ - import json - import os - if os.path.isfile(spec): path = spec else: - from huggingface_hub import hf_hub_download - - path = hf_hub_download(spec, "dataset_statistics.json") + path = _hf_hub_download_fn()(spec, "dataset_statistics.json") with open(path) as f: return json.load(f) +class OpenVLAInputTransform(Transform): + """Build OpenVLA prompt tokens and image tensors from canonical VLA inputs. + + The transform reads TorchRL's canonical VLA observation keys and writes the + preprocessed prompt/image entries consumed by :class:`OpenVLAModelTransform`. + It is intentionally a :class:`~torchrl.envs.Transform`: the same object can + be used in an environment, a replay-buffer transform, or a + :class:`~tensordict.nn.TensorDictSequential` policy stack. + """ + + def __init__( + self, + processor, + *, + use_wrist_image: bool = False, + center_crop: bool = False, + image_backend: Literal["torchvision", "pil", "tensorflow"] = "torchvision", + image_key: NestedKey = ("observation", "image"), + wrist_image_key: NestedKey | None = ("observation", "wrist_image"), + instruction_key: NestedKey = "language_instruction", + input_ids_key: NestedKey = _OPENVLA_INPUT_IDS_KEY, + attention_mask_key: NestedKey = _OPENVLA_ATTENTION_MASK_KEY, + pixel_values_key: NestedKey = _OPENVLA_PIXEL_VALUES_KEY, + device: torch.device | str | None = None, + dtype: torch.dtype | None = None, + ) -> None: + image_key = unravel_key(image_key) + wrist_image_key = ( + None if wrist_image_key is None else unravel_key(wrist_image_key) + ) + instruction_key = unravel_key(instruction_key) + input_ids_key = unravel_key(input_ids_key) + attention_mask_key = unravel_key(attention_mask_key) + pixel_values_key = unravel_key(pixel_values_key) + in_keys = [image_key, instruction_key] + if use_wrist_image and wrist_image_key is not None: + in_keys.append(wrist_image_key) + super().__init__( + in_keys=in_keys, + out_keys=[ + input_ids_key, + attention_mask_key, + pixel_values_key, + ], + ) + self.processor = processor + self.use_wrist_image = bool(use_wrist_image) + self.center_crop = bool(center_crop) + self.image_backend = image_backend + self.image_key = image_key + self.wrist_image_key = wrist_image_key + self.instruction_key = instruction_key + self.input_ids_key = input_ids_key + self.attention_mask_key = attention_mask_key + self.pixel_values_key = pixel_values_key + self.device = None if device is None else torch.device(device) + self.dtype = dtype + self._prompt_cache: dict[str, torch.Tensor] = {} + self._image_preprocessor = self._make_image_preprocessor( + processor, image_backend + ) + if self._image_preprocessor is not None: + self.image_backend = self._image_preprocessor.backend + + def _to_device( + self, tensor: torch.Tensor, *, dtype: torch.dtype | None = None + ) -> torch.Tensor: + device = self.device + if device is None and dtype is None: + return tensor + return tensor.to(device=device, dtype=dtype) + + def _to_pil(self, image: torch.Tensor): + from PIL import Image + + array = image.permute(1, 2, 0).cpu().numpy().astype(np.uint8) + pil = Image.fromarray(array).convert("RGB") + size = _OPENVLA_IMAGE_SIZE + if pil.size != (size, size): + pil = pil.resize((size, size), Image.LANCZOS) + buffer = io.BytesIO() + pil.save(buffer, format="JPEG", quality=_JPEG_QUALITY) + buffer.seek(0) + pil = Image.open(buffer).convert("RGB") + if self.center_crop: + width, height = pil.size + crop_w = int(round(width * _CROP_LINEAR_SCALE)) + crop_h = int(round(height * _CROP_LINEAR_SCALE)) + left = (width - crop_w) // 2 + top = (height - crop_h) // 2 + pil = pil.crop((left, top, left + crop_w, top + crop_h)).resize( + (width, height), Image.LANCZOS + ) + return pil + + def _instructions(self, tensordict: TensorDictBase, batch: int) -> list[str]: + instruction = tensordict.get(self.instruction_key) + data = getattr(instruction, "tolist", lambda: instruction)() + if isinstance(data, str): + data = [data] * batch + elif len(data) != batch: + flat_data = [] + stack = list(data) + while stack: + item = stack.pop(0) + if isinstance(item, str): + flat_data.append(item) + elif isinstance(item, (list, tuple)): + stack[:0] = list(item) + else: + flat_data.append(item) + if len(flat_data) == 1: + data = flat_data * batch + else: + data = flat_data + return [str(item) for item in data] + + def _prompt_input_ids(self, instruction: str) -> torch.Tensor | None: + tokenizer = getattr(self.processor, "tokenizer", None) + if not callable(tokenizer): + return None + prompt = PROMPT_TEMPLATE.format(instruction=instruction.lower()) + input_ids = self._prompt_cache.get(prompt) + if input_ids is None: + tokenized = tokenizer(prompt, return_tensors="pt") + input_ids = tokenized.input_ids.squeeze(0).to( + dtype=torch.long, device="cpu" + ) + if input_ids.numel() == 0 or int(input_ids[-1]) != _EMPTY_TOKEN_ID: + input_ids = torch.cat( + (input_ids, torch.tensor([_EMPTY_TOKEN_ID], dtype=torch.long)) + ) + self._prompt_cache[prompt] = input_ids + return input_ids + + @staticmethod + def _image_processor_config( + image_processor, + ) -> tuple[int, torch.Tensor, torch.Tensor] | None: + input_sizes = getattr(image_processor, "input_sizes", None) + normalize_params = getattr(image_processor, "tvf_normalize_params", None) + if ( + input_sizes is None + or normalize_params is None + or not input_sizes + or len(input_sizes) != len(normalize_params) + or any( + tuple(input_size[-2:]) != (_OPENVLA_IMAGE_SIZE, _OPENVLA_IMAGE_SIZE) + for input_size in input_sizes + ) + ): + return None + return ( + int(input_sizes[0][-1]), + torch.as_tensor( + [params["mean"] for params in normalize_params], + dtype=torch.float32, + ), + torch.as_tensor( + [params["std"] for params in normalize_params], + dtype=torch.float32, + ), + ) + + def _make_image_preprocessor( + self, + processor, + image_backend: Literal["torchvision", "pil", "tensorflow"], + ) -> OpenVLAImagePreprocessor | None: + image_processor = getattr(processor, "image_processor", None) + if image_processor is None: + return None + config = self._image_processor_config(image_processor) + if config is None: + return None + size, mean, std = config + try: + return OpenVLAImagePreprocessor( + size=size, + jpeg_quality=_JPEG_QUALITY, + center_crop=self.center_crop, + backend=image_backend, + mean=mean, + std=std, + ) + except ImportError: + if image_backend != "torchvision": + raise + return OpenVLAImagePreprocessor( + size=size, + jpeg_quality=_JPEG_QUALITY, + center_crop=self.center_crop, + backend="pil", + mean=mean, + std=std, + ) + + def _pixel_values_from_images(self, images: torch.Tensor) -> torch.Tensor | None: + if self._image_preprocessor is None: + return None + return self._image_preprocessor(images) + + def _preprocess(self, images, wrist_images, instructions): + pad_token_id = self.processor.tokenizer.pad_token_id + input_ids_list = [ + self._prompt_input_ids(instruction) for instruction in instructions + ] + if all(input_ids is not None for input_ids in input_ids_list): + pixel_values = self._pixel_values_from_images(images) + if pixel_values is not None: + if wrist_images is not None: + wrist_pixel_values = self._pixel_values_from_images(wrist_images) + if wrist_pixel_values is None: + return self._preprocess_slow(images, wrist_images, instructions) + pixel_values = torch.cat((pixel_values, wrist_pixel_values), dim=1) + input_ids = pad_sequence( + input_ids_list, batch_first=True, padding_value=pad_token_id + ) + attention_mask = input_ids.ne(pad_token_id).long() + return ( + self._to_device(input_ids), + self._to_device(attention_mask), + self._to_device(pixel_values, dtype=self.dtype), + ) + return self._preprocess_slow(images, wrist_images, instructions) + + def _preprocess_slow(self, images, wrist_images, instructions): + pad_token_id = self.processor.tokenizer.pad_token_id + input_ids_list, pixel_values_list = [], [] + for index, instruction in enumerate(instructions): + prompt = PROMPT_TEMPLATE.format(instruction=instruction.lower()) + image = self._to_pil(images[index]) + features = self.processor(prompt, image) + input_ids = features["input_ids"] + pixel_values = [features["pixel_values"]] + if wrist_images is not None: + wrist = self.processor(prompt, self._to_pil(wrist_images[index])) + pixel_values.append(wrist["pixel_values"]) + if not torch.all(input_ids[:, -1] == _EMPTY_TOKEN_ID): + input_ids = torch.cat( + (input_ids, torch.tensor([[_EMPTY_TOKEN_ID]], dtype=torch.long)), + dim=1, + ) + input_ids_list.append(input_ids.squeeze(0)) + pixel_values_list.append(torch.cat(pixel_values, dim=1)) + input_ids = pad_sequence( + input_ids_list, batch_first=True, padding_value=pad_token_id + ) + attention_mask = input_ids.ne(pad_token_id).long() + pixel_values = torch.cat(pixel_values_list, dim=0) + return ( + self._to_device(input_ids), + self._to_device(attention_mask), + self._to_device(pixel_values, dtype=self.dtype), + ) + + def forward(self, tensordict: TensorDictBase) -> TensorDictBase: + images = tensordict.get(self.image_key) + batch_dims = images.shape[:-3] + if images.ndim == 3: + images = images.unsqueeze(0) + else: + images = images.reshape(-1, *images.shape[-3:]) + wrist_images = None + if self.use_wrist_image: + wrist_images = tensordict.get(self.wrist_image_key) + if wrist_images.ndim == 3: + wrist_images = wrist_images.unsqueeze(0) + else: + wrist_images = wrist_images.reshape(-1, *wrist_images.shape[-3:]) + instructions = self._instructions(tensordict, images.shape[0]) + input_ids, attention_mask, pixel_values = self._preprocess( + images, wrist_images, instructions + ) + if batch_dims: + input_ids = input_ids.reshape(*batch_dims, input_ids.shape[-1]) + attention_mask = attention_mask.reshape( + *batch_dims, attention_mask.shape[-1] + ) + pixel_values = pixel_values.reshape(*batch_dims, *pixel_values.shape[-3:]) + tensordict.set(self.input_ids_key, input_ids) + tensordict.set(self.attention_mask_key, attention_mask) + tensordict.set(self.pixel_values_key, pixel_values) + return tensordict + + def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase: + return self.forward(next_tensordict) + + +class OpenVLAModelTransform(Transform): + """Map preprocessed OpenVLA inputs to action-token logits.""" + + def __init__( + self, + model, + processor, + *, + chunk_size: int, + action_dim: int, + num_bins: int, + temperature: float = 1.0, + top_k: int | None = None, + micro_batch_size: int | None = None, + image_key: NestedKey = ("observation", "image"), + input_ids_key: NestedKey = _OPENVLA_INPUT_IDS_KEY, + attention_mask_key: NestedKey = _OPENVLA_ATTENTION_MASK_KEY, + pixel_values_key: NestedKey = _OPENVLA_PIXEL_VALUES_KEY, + logits_key: NestedKey = ("vla_action", "logits"), + ) -> None: + if temperature <= 0: + raise ValueError(f"temperature must be > 0, got {temperature}.") + if top_k is not None and int(top_k) < 1: + raise ValueError(f"top_k must be >= 1 when provided, got {top_k}.") + if micro_batch_size is not None and int(micro_batch_size) < 1: + raise ValueError( + f"micro_batch_size must be >= 1 when provided, got {micro_batch_size}." + ) + image_key = unravel_key(image_key) + input_ids_key = unravel_key(input_ids_key) + attention_mask_key = unravel_key(attention_mask_key) + pixel_values_key = unravel_key(pixel_values_key) + logits_key = unravel_key(logits_key) + super().__init__( + in_keys=[ + image_key, + input_ids_key, + attention_mask_key, + pixel_values_key, + ], + out_keys=[logits_key], + ) + self.model = model + self.processor = processor + self.chunk_size = int(chunk_size) + self.action_dim = int(action_dim) + self.num_bins = int(num_bins) + self.temperature = float(temperature) + self.top_k = int(top_k) if top_k is not None else None + self.micro_batch_size = ( + int(micro_batch_size) if micro_batch_size is not None else None + ) + self.image_key = image_key + self.input_ids_key = input_ids_key + self.attention_mask_key = attention_mask_key + self.pixel_values_key = pixel_values_key + self.logits_key = logits_key + + def _window_logits(self, input_ids, attention_mask, pixel_values): + from openvla_oft.constants import IGNORE_INDEX + + model = self.model + pad_token_id = self.processor.tokenizer.pad_token_id + labels = input_ids.clone() + labels[:] = IGNORE_INDEX + num_prompt_tokens = input_ids.ne(pad_token_id).sum(dim=1) - 1 + input_ids, attention_mask = model._prepare_input_for_action_prediction( + input_ids, attention_mask + ) + labels = model._prepare_labels_for_action_prediction(labels, input_ids) + padding_mask = (~input_ids.ne(pad_token_id)).int() + sorted_indices = torch.argsort( + padding_mask, dim=1, descending=False, stable=True + ) + input_ids = torch.gather(input_ids, 1, sorted_indices) + attention_mask = torch.gather(attention_mask, 1, sorted_indices) + labels = torch.gather(labels, 1, sorted_indices) + input_embeddings = model.get_input_embeddings()(input_ids) + all_actions_mask = model._process_action_masks(labels) + language_embeddings = input_embeddings[~all_actions_mask].reshape( + input_embeddings.shape[0], -1, input_embeddings.shape[2] + ) + projected_patch_embeddings = model._process_vision_features( + pixel_values, language_embeddings, False + ) + num_patches = ( + model.vision_backbone.get_num_patches() + * model.vision_backbone.get_num_images_in_input() + ) + input_embeddings = input_embeddings * ~all_actions_mask.unsqueeze(-1) + ( + multimodal_embeddings, + multimodal_attention_mask, + ) = model._build_multimodal_attention( + input_embeddings, projected_patch_embeddings, attention_mask + ) + output = model.language_model( + input_ids=None, + attention_mask=multimodal_attention_mask, + inputs_embeds=multimodal_embeddings, + use_cache=False, + return_dict=True, + ) + batch_size = output.logits.shape[0] + device = output.logits.device + start = (num_prompt_tokens.to(device) + num_patches).unsqueeze(1) + offsets = torch.arange( + self.chunk_size * self.action_dim, device=device + ).unsqueeze(0) + positions = start + offsets + logits = output.logits[ + torch.arange(batch_size, device=device).unsqueeze(-1), positions, : + ] + return logits[..., model.vocab_size - self.num_bins : model.vocab_size] + + def forward(self, tensordict: TensorDictBase) -> TensorDictBase: + image = tensordict.get(self.image_key) + batch_dims = image.shape[:-3] + input_ids = tensordict.get(self.input_ids_key) + attention_mask = tensordict.get(self.attention_mask_key) + pixel_values = tensordict.get(self.pixel_values_key) + device = next(self.model.parameters()).device + dtype = next(self.model.parameters()).dtype + input_ids = input_ids.reshape(-1, input_ids.shape[-1]).to(device) + attention_mask = attention_mask.reshape(-1, attention_mask.shape[-1]).to(device) + pixel_values = pixel_values.reshape(-1, *pixel_values.shape[-3:]).to( + device=device, dtype=dtype + ) + micro_batch_size = self.micro_batch_size + if micro_batch_size is None or input_ids.shape[0] <= micro_batch_size: + window = self._window_logits(input_ids, attention_mask, pixel_values) + else: + window = torch.cat( + [ + self._window_logits( + input_ids[start : start + micro_batch_size], + attention_mask[start : start + micro_batch_size], + pixel_values[start : start + micro_batch_size], + ) + for start in range(0, input_ids.shape[0], micro_batch_size) + ], + dim=0, + ) + window = window.float() / self.temperature + if self.top_k is not None and self.top_k < window.shape[-1]: + indices = torch.argsort(window, dim=-1, descending=True, stable=True)[ + ..., : self.top_k + ] + values = torch.gather(window, -1, indices) + masked = torch.full_like(window, -torch.inf) + window = masked.scatter(-1, indices, values) + logits = window.reshape( + *batch_dims, self.chunk_size, self.action_dim, self.num_bins + ) + tensordict.set(self.logits_key, logits) + return tensordict + + def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase: + return self.forward(next_tensordict) + + +class GripperPostProcessTransform(Transform): + """Apply the SimpleVLA/OpenVLA gripper convention to decoded actions. + + SimpleVLA decodes the gripper scalar, maps it through ``2 * g - 1``, + binarizes by sign, then optionally inverts for LIBERO. Keeping this as a + transform makes the action-tokenizer a pure token<->continuous codec and + lets env-side decode and policy-side decode share the same postprocess. + """ + + def __init__( + self, + *, + action_key: NestedKey = ACTION_CHUNK_KEY, + out_key: NestedKey | None = None, + rescale: bool = True, + binarize: bool = True, + threshold: float = 0.0, + invert: bool = False, + ) -> None: + action_key = unravel_key(action_key) + out_key = action_key if out_key is None else unravel_key(out_key) + self.action_key = action_key + self.out_key = out_key + self.rescale = bool(rescale) + self.binarize = bool(binarize) + self.threshold = float(threshold) + self.invert = bool(invert) + super().__init__( + in_keys=[action_key], + out_keys=[out_key], + in_keys_inv=[action_key], + out_keys_inv=[out_key], + ) + + def postprocess(self, actions: torch.Tensor) -> torch.Tensor: + """Return actions with SimpleVLA/OpenVLA gripper post-processing.""" + gripper = actions[..., -1] + if self.rescale: + gripper = 2.0 * gripper - 1.0 + if self.binarize: + gripper = (gripper > self.threshold).to(actions.dtype) * 2.0 - 1.0 + if self.invert: + gripper = -gripper + actions = actions.clone() + actions[..., -1] = gripper + return actions + + def _postprocess(self, actions: torch.Tensor) -> torch.Tensor: + return self.postprocess(actions) + + def _apply_transform(self, actions: torch.Tensor) -> torch.Tensor: + return self.postprocess(actions) + + def _inv_apply_transform(self, actions: torch.Tensor) -> torch.Tensor: + return self.postprocess(actions) + + def forward(self, tensordict: TensorDictBase) -> TensorDictBase: + """Postprocess actions when present and ignore observation-only calls.""" + return self._call(tensordict) + + def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase: + actions = next_tensordict.get(self.action_key, None) + if actions is None: + return next_tensordict + next_tensordict.set(self.out_key, self.postprocess(actions)) + return next_tensordict + + def register_openvla_oft() -> None: """Register the vendored token-OFT classes with the transformers Auto* APIs.""" from openvla_oft.configuration_prismatic import OpenVLAConfig @@ -148,6 +685,10 @@ class OpenVLAOFTWrapper(VLAWrapperBase): recomputation. This keeps exploration local while preserving the rollout/loss distribution contract. Defaults to ``None`` (full categorical). + micro_batch_size (int, optional): maximum number of observations in + each model forward. This can reproduce single-environment + inference while the surrounding collector remains batched. + Defaults to ``None`` (one model forward for the complete batch). default_interaction_type (InteractionType, optional): token readout when no exploration context is active (``RANDOM`` samples, else argmax); the forward otherwise follows the ambient @@ -157,6 +698,9 @@ class OpenVLAOFTWrapper(VLAWrapperBase): log_probs_mode (str, optional): ``"sequence"`` or ``"token"``. See :class:`~torchrl.modules.vla.VLAWrapperBase`. Defaults to ``"sequence"``. + output_mode (str, optional): ``"tokens"`` keeps the default token-head + output, ``"chunk"`` decodes to continuous action chunks, and + ``"both"`` writes both representations. Defaults to ``"tokens"``. use_wrist_image (bool, optional): read ``("observation", "wrist_image")`` as the second camera (the checkpoint must have been trained with two input images). @@ -167,18 +711,18 @@ class OpenVLAOFTWrapper(VLAWrapperBase): image_backend (str, optional): backend passed to :class:`~torchrl.data.vla.OpenVLAImagePreprocessor` for the fast batched image path. ``"torchvision"`` keeps preprocessing in tensor - form when available; ``"pil"`` is the reference path. Defaults to - ``"torchvision"``. + form when available; ``"tensorflow"`` is the exact LIBERO + reference path and ``"pil"`` is a lightweight debugging path. + Defaults to ``"torchvision"``. gripper_binarize (bool, optional): binarize the decoded gripper action to +/-1. The model emits a continuous gripper value but robots (LIBERO/robosuite) need a firm open/close, so without this the gripper half-actuates and never secures a grasp. Defaults to ``False``. gripper_binarize_threshold (float, optional): threshold used when - ``gripper_binarize=True``. SimpleVLA-RL's LIBERO post-processing - first maps the [0, 1] gripper output to [-1, 1] and signs it, which - is equivalent to a 0.5 threshold before the optional inversion. - Defaults to ``0.0``. + ``gripper_binarize=True`` after the SimpleVLA ``2 * g - 1`` + remap. A zero threshold here is therefore equivalent to a ``0.5`` + threshold on the raw decoded gripper scalar. Defaults to ``0.0``. gripper_invert (bool, optional): flip the gripper open/close sign after (optional) binarization, for checkpoints whose convention is opposite the environment's. Defaults to ``False``. @@ -192,14 +736,17 @@ def __init__( unnorm_key: str | None = None, temperature: float = 1.0, top_k: int | None = None, + micro_batch_size: int | None = None, default_interaction_type: InteractionType = InteractionType.DETERMINISTIC, log_probs_mode: LogProbsMode = "sequence", + output_mode: Literal["chunk", "tokens", "both"] | None = None, use_wrist_image: bool = False, center_crop: bool = False, - image_backend: Literal["torchvision", "pil"] = "torchvision", + image_backend: Literal["torchvision", "pil", "tensorflow"] = "torchvision", gripper_binarize: bool = False, gripper_binarize_threshold: float = 0.0, gripper_invert: bool = False, + return_vla_action_container: bool = True, ) -> None: from openvla_oft.constants import ACTION_DIM, NUM_ACTIONS_CHUNK @@ -221,9 +768,6 @@ def __init__( model.norm_stats, unnorm_key, num_bins=num_bins, - gripper_binarize=gripper_binarize, - gripper_binarize_threshold=gripper_binarize_threshold, - gripper_invert=gripper_invert, ) super().__init__( action_dim=ACTION_DIM, @@ -234,11 +778,16 @@ def __init__( use_state=False, default_interaction_type=default_interaction_type, log_probs_mode=log_probs_mode, + output_mode=output_mode, + return_vla_action_container=return_vla_action_container, ) self.model = model self.processor = processor self.temperature = float(temperature) self.top_k = int(top_k) if top_k is not None else None + self.micro_batch_size = ( + int(micro_batch_size) if micro_batch_size is not None else None + ) self.use_wrist_image = bool(use_wrist_image) self.center_crop = bool(center_crop) self.image_backend = image_backend @@ -246,13 +795,57 @@ def __init__( self.gripper_binarize_threshold = float(gripper_binarize_threshold) self.gripper_invert = bool(gripper_invert) self.num_bins = num_bins - self._prompt_cache: dict[str, torch.Tensor] = {} self.unnorm_key = unnorm_key - self._image_preprocessor = self._make_image_preprocessor( - processor, image_backend + device = next(model.parameters()).device + dtype = next(model.parameters()).dtype + self.input_transform = OpenVLAInputTransform( + processor, + use_wrist_image=use_wrist_image, + center_crop=center_crop, + image_backend=image_backend, + image_key=self.tensor_keys.image, + wrist_image_key=self.tensor_keys.wrist_image, + instruction_key=self.tensor_keys.instruction, + device=device, + dtype=dtype, ) - if self._image_preprocessor is not None: - self.image_backend = self._image_preprocessor.backend + self.model_transform = OpenVLAModelTransform( + model, + processor, + chunk_size=NUM_ACTIONS_CHUNK, + action_dim=ACTION_DIM, + num_bins=num_bins, + temperature=temperature, + top_k=top_k, + micro_batch_size=micro_batch_size, + image_key=self.tensor_keys.image, + logits_key=self.tensor_keys.action_logits, + ) + self.policy_stack = TensorDictSequential( + self.input_transform, + self.model_transform, + ) + self.gripper_postprocess = GripperPostProcessTransform( + action_key=self.tensor_keys.action_chunk, + rescale=True, + binarize=gripper_binarize, + threshold=gripper_binarize_threshold, + invert=gripper_invert, + ) + self.decode_stack = None + if self.action_tokenizer is not None: + self.decode_stack = TensorDictSequential( + ActionTokenizerTransform( + self.action_tokenizer, + in_key=self.tensor_keys.action_chunk, + out_key=self.tensor_keys.action_tokens, + mode="decode", + ), + self.gripper_postprocess, + ) + self._prompt_cache = self.input_transform._prompt_cache + self._image_preprocessor = self.input_transform._image_preprocessor + self.image_backend = self.input_transform.image_backend if self.use_wrist_image: self.in_keys = [*self.in_keys, ("observation", "wrist_image")] @@ -328,56 +921,13 @@ def make_action_tokenizer(self) -> VocabTailActionTokenizer: # -- preprocessing (shared by rollout and loss-time recompute) --------- def _to_pil(self, image: torch.Tensor): - from PIL import Image - - array = image.permute(1, 2, 0).cpu().numpy().astype(np.uint8) - pil = Image.fromarray(array).convert("RGB") - # Match OpenVLA-OFT's eval chain: resize to 224 (lanczos), JPEG - # round-trip to mirror the RLDS training compression, THEN the - # area-0.9 center crop -- in that order. - size = _OPENVLA_IMAGE_SIZE - if pil.size != (size, size): - pil = pil.resize((size, size), Image.LANCZOS) - buffer = io.BytesIO() - pil.save(buffer, format="JPEG", quality=_JPEG_QUALITY) - buffer.seek(0) - pil = Image.open(buffer).convert("RGB") - if self.center_crop: - width, height = pil.size - crop_w = int(round(width * _CROP_LINEAR_SCALE)) - crop_h = int(round(height * _CROP_LINEAR_SCALE)) - left = (width - crop_w) // 2 - top = (height - crop_h) // 2 - # LANCZOS to mirror the reference tf.image.resize(method="lanczos3") - pil = pil.crop((left, top, left + crop_w, top + crop_h)).resize( - (width, height), Image.LANCZOS - ) - return pil + return self.input_transform._to_pil(image) def _instructions(self, tensordict: TensorDictBase, batch: int) -> list[str]: - instruction = tensordict.get(self.tensor_keys.instruction) - data = getattr(instruction, "tolist", lambda: instruction)() - if isinstance(data, str): - data = [data] * batch - return [str(item) for item in data] + return self.input_transform._instructions(tensordict, batch) def _prompt_input_ids(self, instruction: str) -> torch.Tensor | None: - tokenizer = getattr(self.processor, "tokenizer", None) - if not callable(tokenizer): - return None - prompt = PROMPT_TEMPLATE.format(instruction=instruction.lower()) - input_ids = self._prompt_cache.get(prompt) - if input_ids is None: - tokenized = tokenizer(prompt, return_tensors="pt") - input_ids = tokenized.input_ids.squeeze(0).to( - dtype=torch.long, device="cpu" - ) - if input_ids.numel() == 0 or int(input_ids[-1]) != _EMPTY_TOKEN_ID: - input_ids = torch.cat( - (input_ids, torch.tensor([_EMPTY_TOKEN_ID], dtype=torch.long)) - ) - self._prompt_cache[prompt] = input_ids - return input_ids + return self.input_transform._prompt_input_ids(instruction) @staticmethod def _image_processor_config( @@ -400,189 +950,76 @@ def _image_processor_config( ) def _make_image_preprocessor( - self, processor, image_backend: Literal["torchvision", "pil"] + self, + processor, + image_backend: Literal["torchvision", "pil", "tensorflow"], ) -> OpenVLAImagePreprocessor | None: - image_processor = getattr(processor, "image_processor", None) - if image_processor is None: - return None - config = self._image_processor_config(image_processor) - if config is None: - return None - size, mean, std = config - try: - return OpenVLAImagePreprocessor( - size=size, - jpeg_quality=_JPEG_QUALITY, - center_crop=self.center_crop, - backend=image_backend, - mean=mean, - std=std, - ) - except ImportError: - if image_backend != "torchvision": - raise - return OpenVLAImagePreprocessor( - size=size, - jpeg_quality=_JPEG_QUALITY, - center_crop=self.center_crop, - backend="pil", - mean=mean, - std=std, - ) + return self.input_transform._make_image_preprocessor(processor, image_backend) def _pixel_values_from_images(self, images: torch.Tensor) -> torch.Tensor | None: - if self._image_preprocessor is None: - return None - return self._image_preprocessor(images) + return self.input_transform._pixel_values_from_images(images) def _preprocess(self, images, wrist_images, instructions): - pad_token_id = self.processor.tokenizer.pad_token_id - input_ids_list = [ - self._prompt_input_ids(instruction) for instruction in instructions - ] - if all(input_ids is not None for input_ids in input_ids_list): - pixel_values = self._pixel_values_from_images(images) - if pixel_values is not None: - if wrist_images is not None: - wrist_pixel_values = self._pixel_values_from_images(wrist_images) - if wrist_pixel_values is None: - return self._preprocess_slow(images, wrist_images, instructions) - pixel_values = torch.cat((pixel_values, wrist_pixel_values), dim=1) - input_ids = pad_sequence( - input_ids_list, batch_first=True, padding_value=pad_token_id - ) - attention_mask = input_ids.ne(pad_token_id).long() - device = next(self.model.parameters()).device - dtype = next(self.model.parameters()).dtype - return ( - input_ids.to(device), - attention_mask.to(device), - pixel_values.to(device=device, dtype=dtype), - ) - return self._preprocess_slow(images, wrist_images, instructions) + return self.input_transform._preprocess(images, wrist_images, instructions) def _preprocess_slow(self, images, wrist_images, instructions): - pad_token_id = self.processor.tokenizer.pad_token_id - input_ids_list, pixel_values_list = [], [] - for index, instruction in enumerate(instructions): - prompt = PROMPT_TEMPLATE.format(instruction=instruction.lower()) - image = self._to_pil(images[index]) - features = self.processor(prompt, image) - input_ids = features["input_ids"] - pixel_values = [features["pixel_values"]] - if wrist_images is not None: - wrist = self.processor(prompt, self._to_pil(wrist_images[index])) - pixel_values.append(wrist["pixel_values"]) - if not torch.all(input_ids[:, -1] == _EMPTY_TOKEN_ID): - input_ids = torch.cat( - (input_ids, torch.tensor([[_EMPTY_TOKEN_ID]], dtype=torch.long)), - dim=1, - ) - input_ids_list.append(input_ids.squeeze(0)) - pixel_values_list.append(torch.cat(pixel_values, dim=1)) - input_ids = pad_sequence( - input_ids_list, batch_first=True, padding_value=pad_token_id - ) - attention_mask = input_ids.ne(pad_token_id).long() - pixel_values = torch.cat(pixel_values_list, dim=0) - device = next(self.model.parameters()).device - dtype = next(self.model.parameters()).dtype - return ( - input_ids.to(device), - attention_mask.to(device), - pixel_values.to(device=device, dtype=dtype), - ) + return self.input_transform._preprocess_slow(images, wrist_images, instructions) # -- the parallel-decoding forward ------------------------------------- def _window_logits(self, input_ids, attention_mask, pixel_values): - from openvla_oft.constants import IGNORE_INDEX - - model = self.model - pad_token_id = self.processor.tokenizer.pad_token_id - labels = input_ids.clone() - labels[:] = IGNORE_INDEX - num_prompt_tokens = input_ids.ne(pad_token_id).sum(dim=1) - 1 - input_ids, attention_mask = model._prepare_input_for_action_prediction( - input_ids, attention_mask - ) - labels = model._prepare_labels_for_action_prediction(labels, input_ids) - # _prepare_input_for_action_prediction appends the action placeholders - # and stop token AFTER the padding of shorter rows: sort the padding - # to the end of each row (input ids, attention mask AND labels with - # the same permutation, as in the vendored generate_action_verl) so - # that the action block sits at num_prompt_tokens for every row. - padding_mask = (~input_ids.ne(pad_token_id)).int() - sorted_indices = torch.argsort( - padding_mask, dim=1, descending=False, stable=True - ) - input_ids = torch.gather(input_ids, 1, sorted_indices) - attention_mask = torch.gather(attention_mask, 1, sorted_indices) - labels = torch.gather(labels, 1, sorted_indices) - input_embeddings = model.get_input_embeddings()(input_ids) - all_actions_mask = model._process_action_masks(labels) - language_embeddings = input_embeddings[~all_actions_mask].reshape( - input_embeddings.shape[0], -1, input_embeddings.shape[2] - ) - projected_patch_embeddings = model._process_vision_features( - pixel_values, language_embeddings, False - ) - num_patches = ( - model.vision_backbone.get_num_patches() - * model.vision_backbone.get_num_images_in_input() + return self.model_transform._window_logits( + input_ids, attention_mask, pixel_values ) - input_embeddings = input_embeddings * ~all_actions_mask.unsqueeze(-1) - ( - multimodal_embeddings, - multimodal_attention_mask, - ) = model._build_multimodal_attention( - input_embeddings, projected_patch_embeddings, attention_mask - ) - output = model.language_model( - input_ids=None, - attention_mask=multimodal_attention_mask, - inputs_embeds=multimodal_embeddings, - use_cache=False, - return_dict=True, - ) - batch_size = output.logits.shape[0] - device = output.logits.device - start = (num_prompt_tokens.to(device) + num_patches).unsqueeze(1) - offsets = torch.arange( - self.chunk_size * self.action_dim, device=device - ).unsqueeze(0) - positions = start + offsets - logits = output.logits[ - torch.arange(batch_size, device=device).unsqueeze(-1), positions, : - ] - # restrict to the action-token window at the tail of the (true, - # unpadded) vocabulary: the policy is a num_bins-way categorical - window = logits[..., model.vocab_size - self.num_bins : model.vocab_size] - return window def _action_logits(self, tensordict: TensorDictBase) -> torch.Tensor: - images = tensordict.get(self.tensor_keys.image) - batch_dims = images.shape[:-3] - images = images.reshape(-1, *images.shape[-3:]) - wrist_images = None - if self.use_wrist_image: - wrist_images = tensordict.get(("observation", "wrist_image")) - wrist_images = wrist_images.reshape(-1, *wrist_images.shape[-3:]) - instructions = self._instructions(tensordict, images.shape[0]) - input_ids, attention_mask, pixel_values = self._preprocess( - images, wrist_images, instructions - ) - window = self._window_logits(input_ids, attention_mask, pixel_values) - window = window.float() / self.temperature - if self.top_k is not None and self.top_k < window.shape[-1]: - # Stable sort gives the same tie-breaking as ``argmax`` for - # ``top_k=1`` (lowest index among equal logits), unlike - # ``torch.topk`` whose tie indices are not stable. - indices = torch.argsort(window, dim=-1, descending=True, stable=True)[ - ..., : self.top_k - ] - values = torch.gather(window, -1, indices) - masked = torch.full_like(window, -torch.inf) - window = masked.scatter(-1, indices, values) - return window.reshape( - *batch_dims, self.chunk_size, self.action_dim, self.num_bins + out = self.policy_stack(tensordict.clone(False)) + return out.get(self.tensor_keys.action_logits) + + def forward( + self, + tensordict: TensorDictBase, + *, + tensordict_out: TensorDictBase | None = None, + logits_only: bool = False, + **kwargs, + ) -> TensorDictBase: + out = super().forward( + tensordict, + tensordict_out=tensordict_out, + logits_only=logits_only, + **kwargs, ) + if self.action_head != "tokens": + return out + chunk = out.get(self.tensor_keys.action_chunk, None) + if chunk is None: + return out + chunk = self.gripper_postprocess.postprocess(chunk) + out.set(self.tensor_keys.action_chunk, chunk) + action = out.get(self.tensor_keys.vla_action, None) + if hasattr(action, "chunk"): + action.chunk = chunk + return out + + +def _resolve_component_checkpoint( + pretrained_model_name_or_path: str, filename_or_path: str, pattern: str +) -> str: + if os.path.isfile(filename_or_path): + return filename_or_path + if os.path.isdir(pretrained_model_name_or_path): + exact = os.path.join(pretrained_model_name_or_path, filename_or_path) + if os.path.isfile(exact): + return exact + matches = [ + os.path.join(pretrained_model_name_or_path, filename) + for filename in os.listdir(pretrained_model_name_or_path) + if pattern in filename and "checkpoint" in filename + ] + if len(matches) != 1: + raise FileNotFoundError( + f"expected one {pattern!r} checkpoint in " + f"{pretrained_model_name_or_path!r}, found {len(matches)}." + ) + return matches[0] + return _hf_hub_download_fn()(pretrained_model_name_or_path, filename_or_path) diff --git a/sota-implementations/vla_grpo/test_openvla.py b/sota-implementations/vla_grpo/test_openvla.py index 1047119a3de..f388f07d950 100644 --- a/sota-implementations/vla_grpo/test_openvla.py +++ b/sota-implementations/vla_grpo/test_openvla.py @@ -15,22 +15,22 @@ Requires ``transformers``, ``timm`` and ``Pillow`` (the vendored modeling module imports them). """ + from __future__ import annotations import argparse import importlib.util import os import sys -import warnings from types import SimpleNamespace import numpy as np import pytest import torch from tensordict import NonTensorStack, TensorDict -from tensordict.nn import InteractionType +from tensordict.nn import InteractionType, TensorDictModule from torch import nn -from torchrl.data.vla import ACTION_TOKENS_KEY +from torchrl.data.vla import ACTION_CHUNK_KEY, ACTION_TOKENS_KEY from torchrl.objectives import ClipPPOLoss sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) @@ -43,12 +43,150 @@ ) if _has_deps: - from openvla import OpenVLAOFTWrapper + from openvla import GripperPostProcessTransform, OpenVLAOFTWrapper from openvla_oft.modeling_prismatic import OpenVLAForActionPrediction CHUNK, ACT_DIM, N_BINS = 8, 7, 256 TRUE_VOCAB, PADDED_VOCAB, DIM = 32000, 32064, 16 LOG_PROBS_KEY = ("vla_action", "log_probs") +_REAL_CHECKPOINT = os.environ.get("TORCHRL_OPENVLA_TEST_CHECKPOINT") +_REAL_OBSERVATIONS = os.environ.get("TORCHRL_OPENVLA_TEST_OBSERVATIONS") + + +def _fake_token_policy(): + return TensorDictModule( + lambda action_tokens: action_tokens, + in_keys=[ACTION_TOKENS_KEY], + out_keys=[ACTION_TOKENS_KEY], + ) + + +class TestAdvantageMetrics: + def test_local_advantage_worker_metrics_and_reset(self, monkeypatch): + monkeypatch.setenv("TORCHRL_MC_ADVANTAGE_LOCAL_QUEUES", "1") + advantage = utils.MCAdvantage( + grpo_size=2, + prompt_key="group_id", + trajectory_return="sum", + ) + + class FakeCollector: + def __init__(self): + self.calls = [] + + def map_fn(self, method_name): + self.calls.append(method_name) + if method_name.endswith("get_stats"): + stats = advantage.get_stats() + stats.update( + completed_groups=2, + written_groups=1, + dropped_groups=1, + completed_trajectories=4, + completed_decisions=12, + successful_trajectories=1, + trajectory_return_sum=1.0, + trajectory_return_max=1.0, + ) + return [stats] + return [None] + + collector = FakeCollector() + metrics = utils.advantage_metrics(advantage, collector) + assert metrics["buffer/complete_groups"] == 2 + assert metrics["buffer/kept_groups"] == 1 + assert metrics["buffer/skipped_groups"] == 1 + assert metrics["collector/completed_trajectories"] == 4 + assert metrics["collector/successful_trajectories"] == 1 + + utils.reset_collection_state(advantage, collector) + assert collector.calls[-3:] == [ + "replay_buffer._transform[0].clear_queues", + "replay_buffer._transform[0].reset_stats", + "reset", + ] + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +@pytest.mark.skipif(not _has_deps, reason="OpenVLA dependencies are missing") +@pytest.mark.skipif( + not (_REAL_CHECKPOINT and _REAL_OBSERVATIONS), + reason="set the real OpenVLA checkpoint and observation fixture paths", +) +class TestRealCheckpoint: + def test_real_checkpoint_microbatch_tokens_and_actions_match(self): + """Exercise microbatch equivalence with a real checkpoint when configured.""" + cfg = SimpleNamespace( + policy=SimpleNamespace( + backend="openvla", + mode="tokens", + checkpoint=_REAL_CHECKPOINT, + unnorm_key="libero_spatial_no_noops", + dataset_statistics=os.environ.get( + "TORCHRL_OPENVLA_TEST_DATASET_STATISTICS", + "moojink/openvla-7b-oft-finetuned-libero-spatial", + ), + dtype="bfloat16", + temperature=0.7, + top_k=None, + use_wrist_image=False, + center_crop=True, + image_backend="tensorflow", + gripper_binarize=True, + gripper_binarize_threshold=0.0, + gripper_invert=True, + lora_rank=32, + lora_target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], + ), + loss=SimpleNamespace(ratio_level="token"), + ) + fixture = torch.load(_REAL_OBSERVATIONS, map_location="cpu", weights_only=False) + rows = TensorDict( + { + "observation": {"image": fixture["images"][:8]}, + "language_instruction": NonTensorStack(*fixture["instructions"][:8]), + }, + batch_size=[8], + ) + policy = utils.make_policy(cfg, torch.device("cuda:0")) + policy.eval() + logits = {} + try: + with torch.inference_mode(): + for microbatch in (1, 2, 4, 8): + policy.model_transform.micro_batch_size = microbatch + logits[microbatch] = policy._action_logits(rows.clone(False)) + reference = logits[1] + generator = torch.Generator(device=reference.device).manual_seed(1234) + uniform = torch.rand( + reference.shape, + generator=generator, + device=reference.device, + ).clamp_(1e-7, 1.0 - 1e-7) + gumbel = -torch.log(-torch.log(uniform)) + reference_tokens = (reference + gumbel).argmax(-1) + reference_log_probs = reference.log_softmax(-1).gather( + -1, reference_tokens.unsqueeze(-1) + ) + reference_actions = policy.action_tokenizer.decode(reference_tokens.cpu()) + reference_actions = policy.gripper_postprocess.postprocess( + reference_actions + ) + for microbatch in (2, 4, 8): + assert torch.equal(reference, logits[microbatch]) + tokens = (logits[microbatch] + gumbel).argmax(-1) + assert torch.equal(reference_tokens, tokens) + log_probs = ( + logits[microbatch].log_softmax(-1).gather(-1, tokens.unsqueeze(-1)) + ) + assert torch.equal(reference_log_probs, log_probs) + actions = policy.action_tokenizer.decode(tokens.cpu()) + actions = policy.gripper_postprocess.postprocess(actions) + assert torch.equal(reference_actions, actions) + finally: + del policy + torch.cuda.empty_cache() class _TinyVision(nn.Module): @@ -240,250 +378,499 @@ def policy(): ) -def test_make_collector_casts_deviceless_env_to_cpu(monkeypatch): - captured = {} +def _complete_collector_cfg(**collector_overrides): + collector = { + "groups_per_iter": 4, + "group_size": 2, + "candidate_group_size": None, + "num_collectors": 2, + "envs_per_collector": 4, + "server_max_batch_size": 1, + "server_min_batch_size": 1, + "server_timeout": 0.01, + "server_collect_stats": True, + "server_stats_window_size": 1024, + "policy_micro_batch_size": None, + "max_inflight_per_env": 1, + "num_threads": 1, + "env_sub_threads": 1, + "storing_device": "cpu", + "use_buffers": None, + } + collector.update(collector_overrides) + return SimpleNamespace(**collector) + + +def _complete_toy_env_cfg(**env_overrides): + env = { + "backend": "toy", + "action_dim": 2, + "state_dim": 4, + "image_shape": (3, 8, 8), + "render_size": 16, + "success_steps": 2, + "success_tol": 0.25, + "max_outer_steps": 3, + "num_envs": 4, + "eval_num_envs": 1, + "seed": 0, + "parallel_group_repeats": True, + "train_init_state_mode": "cycle", + "render_backend": None, + "render_gpu_ids": [2, 3], + "eval_render_gpu_ids": None, + "render_gpu_device_zero_fallback": True, + "env_kwargs": None, + } + env.update(env_overrides) + return SimpleNamespace(**env) + + +def _complete_logger_cfg(**logger_overrides): + logger = { + "eval_episodes": 4, + "eval_backend": "thread", + "eval_busy_policy": "skip", + "record_video_single_task": False, + } + logger.update(logger_overrides) + return SimpleNamespace(**logger) + + +class TestCollectorFactory: + def test_make_collector_uses_multicollector_policy_server(self, monkeypatch): + captured = {} + + class _FakeMultiCollector: + def __init__(self, *args, **kwargs): + captured["multi_args"] = args + captured["multi_kwargs"] = kwargs + + class _FakeServer: + def __init__(self, **kwargs): + captured["server_kwargs"] = kwargs + self.transport = kwargs["transport"] + + def start(self): + captured["server_started"] = True + return self + + class _FakeTransport: + def __init__(self, **kwargs): + captured["transport_kwargs"] = kwargs + self.clients = [] + + def client(self): + client = object() + self.clients.append(client) + return client + + policy = _fake_token_policy() + cfg = SimpleNamespace( + collector=_complete_collector_cfg(policy_micro_batch_size=1), + env=_complete_toy_env_cfg(), + ) + replay_buffer = object() + + monkeypatch.setattr(utils, "MultiCollector", _FakeMultiCollector) + monkeypatch.setattr(utils, "ProcessInferenceServer", _FakeServer) + monkeypatch.setattr(utils, "MPTransport", _FakeTransport) - class _FakeCollector: - def __init__(self, *args, **kwargs): - captured["kwargs"] = kwargs + collector, server, eval_policy = utils.make_collector( + cfg, + policy, + torch.device("cpu"), + tokenizer=utils.UniformActionTokenizer(16, low=-1.0, high=1.0), + replay_buffer=replay_buffer, + ) - class _FakeEnv: - batch_size = torch.Size([2]) - device = None + assert isinstance(collector, _FakeMultiCollector) + assert server is not None + assert isinstance(eval_policy, utils.PolicyClientModule) + assert captured["transport_kwargs"]["use_manager"] + assert captured["server_started"] + assert captured["server_kwargs"]["server_config"].max_batch_size == 1 + assert ( + captured["server_kwargs"]["policy_factory"].keywords[ + "policy_micro_batch_size" + ] + == 1 + ) + env_factories = captured["multi_args"][0] + assert [ + factory.keywords["render_gpu_device_id"] for factory in env_factories + ] == [ + 2, + 3, + ] + assert [factory.keywords["worker_idx_offset"] for factory in env_factories] == [ + 0, + 4, + ] + assert captured["multi_kwargs"]["policy"] is None + assert len(captured["multi_kwargs"]["policy_factory"]) == 2 + assert captured["multi_kwargs"]["replay_buffer"] is replay_buffer + assert captured["multi_kwargs"]["trajs_per_batch"] == 1 + assert captured["multi_kwargs"]["traj_format"] == "cat" + assert captured["multi_kwargs"]["storing_device"] == torch.device("cpu") + + def test_make_collector_rejects_cross_subcollector_parallel_groups_without_shared_rb( + self, + ): + policy = _fake_token_policy() + cfg = SimpleNamespace( + collector=_complete_collector_cfg( + num_collectors=4, + envs_per_collector=20, + group_size=8, + groups_per_iter=10, + ), + env=_complete_toy_env_cfg(parallel_group_repeats=True), + ) - cfg = SimpleNamespace( - collector=SimpleNamespace(groups_per_iter=2, group_size=1), - env=SimpleNamespace(max_outer_steps=3), - ) - monkeypatch.setattr(utils, "Collector", _FakeCollector) + with pytest.raises(ValueError, match="collector.envs_per_collector"): + utils.make_collector( + cfg, + policy, + torch.device("cpu"), + tokenizer=utils.UniformActionTokenizer(16, low=-1.0, high=1.0), + replay_buffer=SimpleNamespace(shared=False), + ) - with warnings.catch_warnings(): - warnings.simplefilter("error") - collector = utils.make_collector( - cfg, _FakeEnv(), object(), torch.device("cuda:0") + def test_make_collector_allows_cross_subcollector_groups_with_shared_rb( + self, monkeypatch + ): + captured = {} + + class _FakeMultiCollector: + def __init__(self, *args, **kwargs): + captured["multi_args"] = args + captured["multi_kwargs"] = kwargs + + class _FakeServer: + def __init__(self, **kwargs): + self.transport = kwargs["transport"] + + def start(self): + return self + + class _FakeTransport: + def __init__(self, **kwargs): + pass + + def client(self): + return object() + + policy = _fake_token_policy() + cfg = SimpleNamespace( + collector=_complete_collector_cfg( + num_collectors=4, + envs_per_collector=20, + group_size=8, + groups_per_iter=10, + ), + env=_complete_toy_env_cfg(parallel_group_repeats=True), ) + replay_buffer = SimpleNamespace(shared=True) - assert isinstance(collector, _FakeCollector) - assert captured["kwargs"]["policy_device"] == torch.device("cuda:0") - assert captured["kwargs"]["env_device"] == torch.device("cpu") + monkeypatch.setattr(utils, "MultiCollector", _FakeMultiCollector) + monkeypatch.setattr(utils, "ProcessInferenceServer", _FakeServer) + monkeypatch.setattr(utils, "MPTransport", _FakeTransport) + collector, _, _ = utils.make_collector( + cfg, + policy, + torch.device("cpu"), + tokenizer=utils.UniformActionTokenizer(16, low=-1.0, high=1.0), + replay_buffer=replay_buffer, + ) -def test_make_collector_parallel_groups_use_logical_worker_count(monkeypatch): - captured = {} + assert isinstance(collector, _FakeMultiCollector) + assert len(captured["multi_args"][0]) == 4 - class _FakeCollector: - def __init__(self, *args, **kwargs): - captured["kwargs"] = kwargs - class _FakeEnv: - batch_size = torch.Size([32]) - device = torch.device("cpu") +class TestEvaluatorFactory: + def test_make_evaluator_process_backend_uses_factories(self, monkeypatch): + captured = {} - cfg = SimpleNamespace( - collector=SimpleNamespace(groups_per_iter=16, group_size=8), - env=SimpleNamespace(max_outer_steps=3, parallel_group_repeats=True), - ) - monkeypatch.setattr(utils, "Collector", _FakeCollector) + class _FakeEvaluator: + def __init__(self, env, policy=None, policy_factory=None, **kwargs): + captured["env"] = env + captured["policy"] = policy + captured["policy_factory"] = policy_factory + captured["kwargs"] = kwargs - with pytest.warns(UserWarning, match="parallel_group_repeats=true"): - collector = utils.make_collector( - cfg, _FakeEnv(), object(), torch.device("cuda:0") - ) + monkeypatch.setattr(utils, "Evaluator", _FakeEvaluator) - assert isinstance(collector, _FakeCollector) - assert captured["kwargs"]["trajs_per_batch"] == 16 * 8 - assert captured["kwargs"]["frames_per_batch"] == 32 * 3 + policy = _fake_token_policy() + cfg = SimpleNamespace( + env=_complete_toy_env_cfg(), + logger=_complete_logger_cfg(eval_backend="process"), + ) + evaluator = utils.make_evaluator( + cfg, + utils.UniformActionTokenizer(16, low=-1.0, high=1.0), + policy, + logger=object(), + device=torch.device("cpu"), + ) + assert isinstance(evaluator, _FakeEvaluator) + assert callable(captured["env"]) + assert captured["policy"] is None + assert captured["policy_factory"]() is policy + assert not captured["kwargs"]["dump_video"] + assert captured["kwargs"]["backend"] == "process" + assert captured["kwargs"]["max_steps"] == 0 + + +class TestReplayBufferFactory: + def test_make_replay_buffer_scales_capacity_with_overcollection(self): + cfg = SimpleNamespace( + collector=SimpleNamespace( + groups_per_iter=2, + group_size=3, + candidate_group_size=None, + ), + env=SimpleNamespace(max_outer_steps=5), + advantage=SimpleNamespace( + trajectory_return="sum", + keep_return_bounds=None, + candidate_selection="balanced", + candidate_selection_min_size=None, + candidate_selection_max_combinations=100000, + ), + loss=SimpleNamespace(mini_batch_size=2), + buffer=SimpleNamespace( + shared_init=True, + capacity_group_waves=4, + consume_after_n_samples=1, + ), + ) -def test_make_collector_parallel_group_wave_has_no_warning(monkeypatch): - captured = {} + replay_buffer, _ = utils.make_replay_buffer(cfg, torch.device("cpu")) + + assert replay_buffer._storage.max_size == 2 * 3 * 5 * 4 + assert replay_buffer._storage.shared_init + assert replay_buffer.shared + + def test_make_replay_buffer_scales_capacity_with_candidate_group_size(self): + cfg = SimpleNamespace( + collector=SimpleNamespace( + groups_per_iter=2, + group_size=3, + candidate_group_size=6, + ), + env=SimpleNamespace(max_outer_steps=5), + advantage=SimpleNamespace( + trajectory_return="sum", + keep_return_bounds=None, + candidate_selection="balanced", + candidate_selection_min_size=4, + candidate_selection_max_combinations=100000, + ), + loss=SimpleNamespace(mini_batch_size=2), + buffer=SimpleNamespace( + shared_init=False, + capacity_group_waves=4, + consume_after_n_samples=1, + ), + ) - class _FakeCollector: - def __init__(self, *args, **kwargs): - captured["kwargs"] = kwargs + replay_buffer, advantage = utils.make_replay_buffer(cfg, torch.device("cpu")) - class _FakeEnv: - batch_size = torch.Size([32]) - device = torch.device("cpu") + assert replay_buffer._storage.max_size == 2 * 6 * 5 * 4 + assert advantage.grpo_size == 3 + assert advantage.candidate_group_size == 6 + assert advantage.candidate_selection_min_size == 4 - cfg = SimpleNamespace( - collector=SimpleNamespace(groups_per_iter=4, group_size=8), - env=SimpleNamespace(max_outer_steps=3, parallel_group_repeats=True), - ) - monkeypatch.setattr(utils, "Collector", _FakeCollector) - collector = utils.make_collector(cfg, _FakeEnv(), object(), torch.device("cuda:0")) +class TestTokenizerAndTransforms: + def test_chunk_transform_without_tokenizer_consumes_vla_chunk(self): + cfg = SimpleNamespace(env=SimpleNamespace(max_outer_steps=1)) - assert isinstance(collector, _FakeCollector) - assert captured["kwargs"]["trajs_per_batch"] == 4 * 8 - assert captured["kwargs"]["frames_per_batch"] == 32 * 3 + transform = utils._chunk_transform(cfg, None) + assert transform[0].out_keys_inv == [ACTION_CHUNK_KEY] -def test_make_collector_can_write_to_replay_buffer(monkeypatch): - captured = {} + def test_chunk_transform_openvla_tokens_decodes_then_postprocesses(self): + cfg = SimpleNamespace( + env=SimpleNamespace(max_outer_steps=1), + policy=SimpleNamespace( + backend="openvla", + mode="tokens", + gripper_binarize=True, + gripper_binarize_threshold=0.0, + gripper_invert=True, + ), + ) - class _FakeCollector: - def __init__(self, *args, **kwargs): - captured["kwargs"] = kwargs + transform = utils._chunk_transform( + cfg, + utils.UniformActionTokenizer(16, low=-1.0, high=1.0), + ) - class _FakeEnv: - batch_size = torch.Size([4]) - device = torch.device("cpu") + assert transform[0].__class__.__name__ == "MultiAction" + assert transform[1].__class__.__name__ == "GripperPostProcessTransform" + assert transform[2].__class__.__name__ == "ActionTokenizerTransform" + assert transform[1].in_keys_inv == ["action"] + assert transform[1].out_keys_inv == ["action"] + + +class TestLiberoWorkers: + def test_libero_worker_assignment_serial_and_parallel_groups(self): + class _EnvCfg(dict): + def __getattr__(self, name): + return self[name] + + cfg = SimpleNamespace( + collector=SimpleNamespace(group_size=8, candidate_group_size=None), + env=_EnvCfg( + { + "task_ids": [10, 11, 12], + "parallel_group_repeats": False, + } + ), + ) - cfg = SimpleNamespace( - collector=SimpleNamespace(groups_per_iter=4, group_size=2), - env=SimpleNamespace(max_outer_steps=3), - ) - replay_buffer = object() + assert utils._libero_worker_assignment(cfg, 4, group_repeats=8) == ( + 11, + 8, + 4 * utils.GROUP_ID_OFFSET, + ) - def hook(_): - return None + cfg.env["parallel_group_repeats"] = True - monkeypatch.setattr(utils, "Collector", _FakeCollector) + assert utils._libero_worker_assignment(cfg, 0, group_repeats=8) == (10, 1, 0) + assert utils._libero_worker_assignment(cfg, 7, group_repeats=8) == (10, 1, 0) + assert utils._libero_worker_assignment(cfg, 8, group_repeats=8) == ( + 11, + 1, + utils.GROUP_ID_OFFSET, + ) + cfg.collector.candidate_group_size = 16 + assert utils._libero_worker_assignment(cfg, 7, group_repeats=8) == (10, 2, 0) - collector = utils.make_collector( - cfg, - _FakeEnv(), - object(), - torch.device("cpu"), - replay_buffer=replay_buffer, - post_collect_hook=hook, - ) + def test_make_libero_worker_parallel_groups_by_init_state(self, monkeypatch): + class _Cfg(dict): + def __getattr__(self, name): + return self[name] - assert isinstance(collector, _FakeCollector) - assert captured["kwargs"]["replay_buffer"] is replay_buffer - assert captured["kwargs"]["post_collect_hook"] is hook - assert captured["kwargs"]["trajs_per_batch"] == 4 * 2 - assert captured["kwargs"]["frames_per_batch"] == 4 - - -def test_make_replay_buffer_scales_capacity_with_overcollection(): - cfg = SimpleNamespace( - collector=SimpleNamespace( - groups_per_iter=2, - group_size=3, - max_collect_batches_per_iter=4, - ), - env=SimpleNamespace(max_outer_steps=5), - advantage=SimpleNamespace( - trajectory_return="sum", - keep_return_bounds=None, - ), - loss=SimpleNamespace(mini_batch_size=2), - ) + captured = {} - replay_buffer, _ = utils.make_replay_buffer(cfg, torch.device("cpu")) + def _fake_libero_env(*args, **kwargs): + captured["args"] = args + captured["kwargs"] = kwargs + return object() + + cfg = SimpleNamespace( + env=_Cfg( + task_suite="libero_spatial", + task_ids=[0, 1], + camera_height=64, + camera_width=64, + render_backend="egl", + render_gpu_ids=[2, 3], + eval_render_gpu_ids=None, + render_gpu_device_zero_fallback=True, + env_kwargs=None, + max_env_steps=512, + train_init_state_mode="cycle", + train_init_state_id=3, + parallel_group_repeats=True, + ), + collector=SimpleNamespace(group_size=8, candidate_group_size=None), + policy=SimpleNamespace(use_wrist_image=False), + ) + monkeypatch.setattr(utils, "LiberoEnv", _fake_libero_env) - assert replay_buffer._storage.max_size == 2 * 3 * 5 * 4 + utils._make_libero_worker(cfg, 7, group_repeats=8) + assert captured["args"] == ("libero_spatial",) + assert captured["kwargs"]["task_id"] == 0 + assert captured["kwargs"]["group_repeats"] == 1 + assert captured["kwargs"]["group_id_offset"] == 0 + assert captured["kwargs"]["group_id_mode"] == "init_state" + assert captured["kwargs"]["init_state_id"] == 3 + assert captured["kwargs"]["env_kwargs"]["render_gpu_device_id"] == 3 -def test_make_replay_buffer_scales_capacity_with_candidate_group_size(): - cfg = SimpleNamespace( - collector=SimpleNamespace( - groups_per_iter=2, - group_size=3, - candidate_group_size=6, - max_collect_batches_per_iter=4, - ), - env=SimpleNamespace(max_outer_steps=5), - advantage=SimpleNamespace( - trajectory_return="sum", - keep_return_bounds=None, - candidate_selection="balanced", - candidate_selection_min_size=4, - ), - loss=SimpleNamespace(mini_batch_size=2), - ) - replay_buffer, advantage = utils.make_replay_buffer(cfg, torch.device("cpu")) +@pytest.mark.skipif(not _has_deps, reason="transformers/timm/PIL not found") +class TestOpenVLAOFTWrapper: + def test_forward_shapes(self, policy): + out = policy(_make_obs()) + assert out[ACTION_TOKENS_KEY].shape == (2, CHUNK, ACT_DIM) + assert out[ACTION_TOKENS_KEY].min() >= 0 + assert out[ACTION_TOKENS_KEY].max() < N_BINS + assert out[LOG_PROBS_KEY].shape == (2,) - assert replay_buffer._storage.max_size == 2 * 6 * 5 * 4 - assert advantage.grpo_size == 3 - assert advantage.candidate_group_size == 6 - assert advantage.candidate_selection_min_size == 4 + def test_forward_nested_batch_shapes(self, policy): + out = policy(_make_obs(batch=8).reshape(1, 8)) + assert out[ACTION_TOKENS_KEY].shape == (1, 8, CHUNK, ACT_DIM) + assert out[LOG_PROBS_KEY].shape == (1, 8) + def test_policy_stack_matches_action_logits(self, policy): + obs = _make_obs() -def test_libero_worker_assignment_serial_and_parallel_groups(): - class _EnvCfg(dict): - def __getattr__(self, name): - return self[name] + with torch.no_grad(): + stacked = policy.policy_stack(obs.clone()) + logits = policy._action_logits(obs.clone()) - cfg = SimpleNamespace( - collector=SimpleNamespace(group_size=8), - env=_EnvCfg( - { - "task_ids": [10, 11, 12], - "parallel_group_repeats": False, - } - ), - ) + torch.testing.assert_close( + stacked.get(policy.tensor_keys.action_logits), + logits, + ) + assert stacked.get(policy.input_transform.input_ids_key).shape[0] == 2 + assert stacked.get(policy.input_transform.pixel_values_key).shape[0] == 2 + + def test_gripper_postprocess_matches_simplevla_order(self): + transform = GripperPostProcessTransform( + action_key="action", + rescale=True, + binarize=True, + threshold=0.0, + invert=True, + ) + actions = torch.zeros(2, CHUNK, ACT_DIM) + actions[0, :, -1] = 0.25 + actions[1, :, -1] = 0.75 - assert utils._libero_worker_assignment(cfg, 4, group_repeats=8) == ( - 11, - 8, - 4 * utils.GROUP_ID_OFFSET, - ) + out = transform(TensorDict({"action": actions}, batch_size=[2]))["action"] - cfg.env["parallel_group_repeats"] = True + torch.testing.assert_close(out[0, :, -1], torch.ones(CHUNK)) + torch.testing.assert_close(out[1, :, -1], -torch.ones(CHUNK)) - assert utils._libero_worker_assignment(cfg, 0, group_repeats=8) == (10, 1, 0) - assert utils._libero_worker_assignment(cfg, 7, group_repeats=8) == (10, 1, 0) - assert utils._libero_worker_assignment(cfg, 8, group_repeats=8) == ( - 11, - 1, - utils.GROUP_ID_OFFSET, - ) - cfg.collector.candidate_group_size = 16 - assert utils._libero_worker_assignment(cfg, 7, group_repeats=8) == (10, 2, 0) - - -def test_make_libero_worker_parallel_groups_by_init_state(monkeypatch): - class _Cfg(dict): - def __getattr__(self, name): - return self[name] - - captured = {} - - def _fake_libero_env(*args, **kwargs): - captured["args"] = args - captured["kwargs"] = kwargs - return object() - - cfg = SimpleNamespace( - env=_Cfg( - task_suite="libero_spatial", - task_ids=[0, 1], - camera_height=64, - camera_width=64, - render_gpu_ids=[2, 3], - eval_render_gpu_ids=None, - max_env_steps=512, - train_init_state_mode="cycle", - parallel_group_repeats=True, - ), - collector=SimpleNamespace(group_size=8), - policy=SimpleNamespace(use_wrist_image=False), - ) - monkeypatch.setattr(utils, "LiberoEnv", _fake_libero_env) + def test_decode_stack_applies_tokenizer_then_gripper_postprocess(self): + policy = OpenVLAOFTWrapper( + _TinyOFT(), + _FakeProcessor(), + output_mode="both", + gripper_binarize=True, + gripper_invert=True, + ) + obs = _make_obs() - utils._make_libero_worker(cfg, 7, group_repeats=8) + with torch.no_grad(): + out = policy(obs) - assert captured["args"] == ("libero_spatial",) - assert captured["kwargs"]["task_id"] == 0 - assert captured["kwargs"]["group_repeats"] == 1 - assert captured["kwargs"]["group_id_offset"] == 0 - assert captured["kwargs"]["group_id_mode"] == "init_state" - assert captured["kwargs"]["env_kwargs"]["render_gpu_device_id"] == 3 + decoded = policy.action_tokenizer.decode(out[ACTION_TOKENS_KEY]) + expected = policy.gripper_postprocess.postprocess(decoded) + torch.testing.assert_close(out[ACTION_CHUNK_KEY], expected) + decoded_td = policy.decode_stack( + TensorDict( + {ACTION_TOKENS_KEY: out[ACTION_TOKENS_KEY]}, + batch_size=out.batch_size, + ) + ) + torch.testing.assert_close(decoded_td[ACTION_CHUNK_KEY], expected) -@pytest.mark.skipif(not _has_deps, reason="transformers/timm/PIL not found") -class TestOpenVLAOFTWrapper: - def test_forward_shapes(self, policy): - out = policy(_make_obs()) - assert out[ACTION_TOKENS_KEY].shape == (2, CHUNK, ACT_DIM) - assert out[ACTION_TOKENS_KEY].min() >= 0 - assert out[ACTION_TOKENS_KEY].max() < N_BINS - assert out[LOG_PROBS_KEY].shape == (2,) + def test_gripper_postprocess_skips_observation_forward_path(self): + transform = GripperPostProcessTransform(action_key="action") + observation = TensorDict({"observation": torch.zeros(3)}, batch_size=[]) + out = transform(observation) + assert "action" not in out + torch.testing.assert_close(out["observation"], torch.zeros(3)) def test_temperature_contract_ratio_one(self, policy): # the go/no-go contract: with identical weights, the log-probs written @@ -548,6 +935,20 @@ def test_mixed_prompt_lengths_consistent(self, policy): msg=f"row {row} differs between batched and single forward", ) + def test_model_microbatch_matches_sequential(self): + policy = OpenVLAOFTWrapper(_TinyOFT(), _FakeProcessor(), micro_batch_size=1) + obs = _make_obs(batch=3) + with torch.no_grad(): + microbatched = policy._action_logits(obs.clone()) + sequential = torch.cat( + [policy._action_logits(obs[row : row + 1].clone()) for row in range(3)], + dim=0, + ) + torch.testing.assert_close(microbatched, sequential) + assert policy.model_transform.micro_batch_size == 1 + with pytest.raises(ValueError, match="micro_batch_size"): + OpenVLAOFTWrapper(_TinyOFT(), _FakeProcessor(), micro_batch_size=0) + def test_token_log_probs_mode(self): torch.manual_seed(0) policy = OpenVLAOFTWrapper(_TinyOFT(), _FakeProcessor(), log_probs_mode="token") diff --git a/sota-implementations/vla_grpo/utils.py b/sota-implementations/vla_grpo/utils.py index bdef29a624b..fd7b2c9d01b 100644 --- a/sota-implementations/vla_grpo/utils.py +++ b/sota-implementations/vla_grpo/utils.py @@ -21,17 +21,21 @@ """ from __future__ import annotations +import importlib.util +import multiprocessing as mp import os +import time import warnings from collections.abc import Callable +from contextlib import contextmanager from functools import partial import torch -from tensordict import TensorDictBase -from torchrl.collectors import Collector +from tensordict import NonTensorData, TensorDict, TensorDictBase +from torchrl._utils import logger as torchrl_logger, timeit +from torchrl.collectors import Evaluator, MultiCollector from torchrl.data import LazyTensorStorage, TensorDictReplayBuffer -from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement from torchrl.data.vla import ( ACTION_TOKENS_KEY, ActionTokenizerBase, @@ -40,7 +44,6 @@ from torchrl.envs import ( ActionTokenizerTransform, Compose, - EnvBase, LiberoEnv, MultiAction, ParallelEnv, @@ -49,30 +52,131 @@ ToyVLAEnv, TransformedEnv, ) -from torchrl.envs.utils import ExplorationType, set_exploration_type +from torchrl.envs.utils import ExplorationType +from torchrl.modules.inference_server import ( + InferenceServerConfig, + MPTransport, + PolicyClientModule, + ProcessInferenceServer, +) from torchrl.modules.vla import TinyVLA, VLAWrapperBase from torchrl.objectives import ClipPPOLoss from torchrl.objectives.llm import MCAdvantage, MCAdvantageSelector from torchrl.record import VideoRecorder +from torchrl.record.loggers import generate_exp_name, get_logger +from torchrl.weight_update import WeightStrategy # group ids must be unique across parallel workers: each worker gets a # disjoint offset block GROUP_ID_OFFSET = 10**6 +_has_robosuite = importlib.util.find_spec("robosuite") is not None +_has_openvla = importlib.util.find_spec("openvla") is not None +_has_peft = importlib.util.find_spec("peft") is not None +_ROBOSUITE_EGL_DEVICE_COUNT: int | None = None +_OpenVLAOFTWrapper = None +_OpenVLAOFTL1Wrapper = None +_GripperPostProcessTransform = None +_get_peft_model = None +_LoraConfig = None LOG_PROBS_KEY = ("vla_action", "log_probs") +_TORCH_DTYPES = { + "bfloat16": torch.bfloat16, + "float16": torch.float16, + "float32": torch.float32, +} + + +def _openvla_wrapper_cls(mode: str): + """Lazy optional import for the OpenVLA backend.""" + global _OpenVLAOFTL1Wrapper, _OpenVLAOFTWrapper + if not _has_openvla: + raise ImportError("The openvla backend requires the local openvla module.") + if mode == "tokens": + if _OpenVLAOFTWrapper is None: + from openvla import OpenVLAOFTWrapper + + _OpenVLAOFTWrapper = OpenVLAOFTWrapper + return _OpenVLAOFTWrapper + if mode == "l1": + if _OpenVLAOFTL1Wrapper is None: + from openvla import OpenVLAOFTL1Wrapper + + _OpenVLAOFTL1Wrapper = OpenVLAOFTL1Wrapper + return _OpenVLAOFTL1Wrapper + raise ValueError(f"policy.mode must be 'tokens' or 'l1', got {mode!r}.") + + +def _gripper_postprocess_transform_cls(): + """Lazy optional import for OpenVLA gripper post-processing.""" + global _GripperPostProcessTransform + if not _has_openvla: + raise ImportError("The openvla backend requires the local openvla module.") + if _GripperPostProcessTransform is None: + from openvla import GripperPostProcessTransform + + _GripperPostProcessTransform = GripperPostProcessTransform + return _GripperPostProcessTransform + + +def _peft_lora_tools(): + """Lazy optional import for LoRA fine-tuning.""" + global _LoraConfig, _get_peft_model + if not _has_peft: + raise ImportError("policy.lora_rank requires the peft package.") + if _get_peft_model is None or _LoraConfig is None: + from peft import get_peft_model, LoraConfig + + _get_peft_model = get_peft_model + _LoraConfig = LoraConfig + return _get_peft_model, _LoraConfig + + +def _cfg_get(section, key: str, default=None): + """Read an optional key from a config section. + + Works for mappings (``DictConfig``, ``dict``) through ``.get`` and for + attribute-style sections (``SimpleNamespace``, dataclasses) through + ``getattr`` with a default. + """ + getter = getattr(section, "get", None) + if callable(getter): + return getter(key, default) + return getattr(section, key, default) + def candidate_group_size(cfg) -> int: """Number of rollout candidates collected for each GRPO group.""" - collector_get = getattr( - cfg.collector, - "get", - lambda key, default=None: getattr(cfg.collector, key, default), + return int(cfg.collector.candidate_group_size or cfg.collector.group_size) + + +def auto_device(device_spec) -> torch.device: + """Resolve an explicit device or default to cuda:0 when available.""" + if device_spec: + return torch.device(device_spec) + return torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + + +def make_logger(cfg): + """Build the configured experiment logger.""" + if not cfg.logger.backend: + return None + exp_name = generate_exp_name("VLA-GRPO", cfg.logger.exp_name) + return get_logger( + logger_type=cfg.logger.backend, + logger_name="vla_grpo_logging", + experiment_name=exp_name, + wandb_kwargs={ + "mode": cfg.logger.mode, + "config": dict(cfg), + "project": cfg.logger.project_name, + "group": cfg.logger.group_name, + }, ) - return int(collector_get("candidate_group_size", None) or cfg.collector.group_size) def _configure_mujoco_rendering(cfg) -> None: - render_backend = cfg.env.get("render_backend", None) + render_backend = cfg.env.render_backend if render_backend is None: return render_backend = str(render_backend) @@ -82,14 +186,20 @@ def _configure_mujoco_rendering(cfg) -> None: def _worker_render_gpu_device_id( - cfg, worker_idx: int, *, eval_mode: bool = False + cfg, + worker_idx: int, + *, + eval_mode: bool = False, + render_gpu_device_id: int | None = None, ) -> int | None: - render_backend = cfg.env.get("render_backend", None) + if render_gpu_device_id is not None: + return int(render_gpu_device_id) + render_backend = cfg.env.render_backend if render_backend is not None and str(render_backend) != "egl": return None - render_gpu_ids = cfg.env.get("eval_render_gpu_ids", None) if eval_mode else None + render_gpu_ids = cfg.env.eval_render_gpu_ids if eval_mode else None if render_gpu_ids is None: - render_gpu_ids = cfg.env.get("render_gpu_ids", None) + render_gpu_ids = cfg.env.render_gpu_ids if render_gpu_ids is None: return None render_gpu_ids = [int(device_id) for device_id in render_gpu_ids] @@ -98,15 +208,74 @@ def _worker_render_gpu_device_id( return render_gpu_ids[worker_idx % len(render_gpu_ids)] +def _robosuite_egl_device_count() -> int | None: + global _ROBOSUITE_EGL_DEVICE_COUNT + if not _has_robosuite: + return None + if _ROBOSUITE_EGL_DEVICE_COUNT is None: + try: + import robosuite.renderers.context.egl_context as egl_context + except Exception: + _ROBOSUITE_EGL_DEVICE_COUNT = -1 + else: + try: + _ROBOSUITE_EGL_DEVICE_COUNT = len(egl_context.EGL.eglQueryDevicesEXT()) + except Exception: + _ROBOSUITE_EGL_DEVICE_COUNT = -1 + if _ROBOSUITE_EGL_DEVICE_COUNT < 0: + return None + return _ROBOSUITE_EGL_DEVICE_COUNT + + +@contextmanager +def _worker_render_gpu_context(cfg, render_gpu_device_id: int | None): + """Map the requested render GPU to EGL device 0 if EGL is CVD-scoped.""" + if render_gpu_device_id is None or not bool( + cfg.env.render_gpu_device_zero_fallback + ): + yield render_gpu_device_id + return + + egl_device_count = _robosuite_egl_device_count() + if egl_device_count is None or render_gpu_device_id < egl_device_count: + yield render_gpu_device_id + return + + old_mujoco_egl_device_id = os.environ.get("MUJOCO_EGL_DEVICE_ID") + old_cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") + try: + # Some pyxis images expose one EGL device relative to the current + # CUDA_VISIBLE_DEVICES value. Narrowing CVD in the env worker maps the + # requested physical render GPU to local EGL id 0, instead of sending + # every worker to physical GPU 0. + # + # robosuite 1.4 also asserts that MUJOCO_EGL_DEVICE_ID appears as a + # literal entry in CUDA_VISIBLE_DEVICES, even though MuJoCo interprets + # it as the local EGL ordinal. Keep physical GPU 0 visible as a dummy + # second entry for nonzero render GPUs so the assertion accepts local + # EGL id 0 while CUDA still maps local device 0 to render_gpu_device_id. + cuda_visible_devices = str(render_gpu_device_id) + if render_gpu_device_id != 0: + cuda_visible_devices = f"{render_gpu_device_id},0" + os.environ["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices + os.environ["MUJOCO_EGL_DEVICE_ID"] = "0" + yield 0 + finally: + if old_mujoco_egl_device_id is None: + os.environ.pop("MUJOCO_EGL_DEVICE_ID", None) + else: + os.environ["MUJOCO_EGL_DEVICE_ID"] = old_mujoco_egl_device_id + if old_cuda_visible_devices is None: + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + else: + os.environ["CUDA_VISIBLE_DEVICES"] = old_cuda_visible_devices + + def _libero_worker_assignment( cfg, worker_idx: int, *, group_repeats=None, eval_mode=False ) -> tuple[int, int | None, int]: task_ids = list(cfg.env.task_ids) - if ( - eval_mode - or group_repeats is None - or not cfg.env.get("parallel_group_repeats", False) - ): + if eval_mode or group_repeats is None or not cfg.env.parallel_group_repeats: return ( task_ids[worker_idx % len(task_ids)], group_repeats, @@ -129,7 +298,12 @@ def _libero_worker_assignment( ) -def make_policy(cfg, device: torch.device) -> VLAWrapperBase: +def make_policy( + cfg, + device: torch.device, + *, + policy_micro_batch_size: int | None = None, +) -> VLAWrapperBase: # ratio_level="token" gives one importance ratio per action token (the # SimpleVLA-RL / DAPO semantics: the clip thresholds are per-token); # "sequence" sums the chunk's log-probs into a single ratio per decision. @@ -142,34 +316,44 @@ def make_policy(cfg, device: torch.device) -> VLAWrapperBase: vocab_size=cfg.tokenizer.vocab_size, hidden_dim=cfg.policy.hidden_dim, log_probs_mode=log_probs_mode, + return_vla_action_container=False, device=device, ) if cfg.policy.backend == "openvla": - # local import: pulls in transformers/timm via the vendored modeling - from openvla import OpenVLAOFTWrapper + OpenVLAOFTWrapper = _openvla_wrapper_cls(cfg.policy.mode) - policy = OpenVLAOFTWrapper.from_pretrained( - cfg.policy.checkpoint, - torch_dtype=getattr(torch, cfg.policy.dtype), - device=device, - unnorm_key=cfg.policy.unnorm_key, - dataset_statistics=cfg.policy.get("dataset_statistics", None), - temperature=cfg.policy.temperature, - top_k=cfg.policy.get("top_k", None), - log_probs_mode=log_probs_mode, - use_wrist_image=cfg.policy.use_wrist_image, - center_crop=cfg.policy.center_crop, - image_backend=cfg.policy.get("image_backend", "torchvision"), - gripper_binarize=cfg.policy.get("gripper_binarize", False), - gripper_binarize_threshold=cfg.policy.get( - "gripper_binarize_threshold", 0.0 - ), - gripper_invert=cfg.policy.get("gripper_invert", False), - ) + kwargs = { + "torch_dtype": _TORCH_DTYPES[cfg.policy.dtype], + "device": device, + "unnorm_key": cfg.policy.unnorm_key, + "dataset_statistics": cfg.policy.dataset_statistics, + "use_wrist_image": cfg.policy.use_wrist_image, + "center_crop": cfg.policy.center_crop, + "image_backend": cfg.policy.image_backend, + "gripper_binarize": cfg.policy.gripper_binarize, + "gripper_binarize_threshold": cfg.policy.gripper_binarize_threshold, + "gripper_invert": cfg.policy.gripper_invert, + "return_vla_action_container": False, + } + if cfg.policy.mode == "tokens": + kwargs.update( + temperature=cfg.policy.temperature, + top_k=cfg.policy.top_k, + micro_batch_size=policy_micro_batch_size, + log_probs_mode=log_probs_mode, + ) + else: + kwargs.update( + action_head_file=cfg.policy.action_head_file, + proprio_projector_file=cfg.policy.proprio_projector_file, + use_proprio=cfg.policy.use_proprio, + num_images_in_input=cfg.policy.num_images_in_input, + ) + policy = OpenVLAOFTWrapper.from_pretrained(cfg.policy.checkpoint, **kwargs) if cfg.policy.lora_rank: # de-risk fallback to full fine-tuning (RL4VLA shows LoRA r=32 # works); validate on the target hardware - from peft import get_peft_model, LoraConfig + get_peft_model, LoraConfig = _peft_lora_tools() policy.model = get_peft_model( policy.model, @@ -183,8 +367,10 @@ def make_policy(cfg, device: torch.device) -> VLAWrapperBase: raise ValueError(f"Unknown policy backend {cfg.policy.backend!r}.") -def make_action_tokenizer(cfg, policy: VLAWrapperBase) -> ActionTokenizerBase: +def make_action_tokenizer(cfg, policy: VLAWrapperBase) -> ActionTokenizerBase | None: if cfg.policy.backend == "openvla": + if cfg.policy.mode == "l1": + return None # the codec lives in the checkpoint (vocab-tail mapping + norm_stats) if policy.action_tokenizer is None: raise RuntimeError( @@ -196,60 +382,207 @@ def make_action_tokenizer(cfg, policy: VLAWrapperBase) -> ActionTokenizerBase: return UniformActionTokenizer(cfg.tokenizer.vocab_size, low=-1.0, high=1.0) -def _chunk_transform(cfg, tokenizer: ActionTokenizerBase) -> Compose: - # The compose order is load-bearing: the inverse (action-input) path runs - # in reverse, so the tokenizer decode happens before MultiAction unbinds - # the chunk; on the step path SuccessReward and StepCounter run after - # MultiAction, i.e. once per outer (decision) step. stack_rewards=False - # keeps the outer transition dense when an episode ends inside a chunk - # (the decision reward comes from the outer success flag instead). - return Compose( - MultiAction(stack_rewards=False), - ActionTokenizerTransform(tokenizer), - SuccessReward(), - StepCounter(max_steps=cfg.env.max_outer_steps), +def _chunk_transform( + cfg, + tokenizer: ActionTokenizerBase | None, + *, + decode_actions_in_env: bool = True, +) -> Compose: + # The compose order is load-bearing: when token decoding lives in the env, + # the inverse (action-input) path runs in reverse, so tokenizer decode + # happens before MultiAction unbinds the chunk. When token decoding lives + # in the policy, the env receives the continuous chunk directly and + # MultiAction is still responsible for unbinding it. On the step path, + # SuccessReward and StepCounter run after MultiAction, i.e. once per outer + # (decision) step. stack_rewards=False keeps the outer transition dense + # when an episode ends inside a chunk (the decision reward comes from the + # outer success flag instead). + if decode_actions_in_env and tokenizer is not None: + transforms = [MultiAction(stack_rewards=False)] + if cfg.policy.backend == "openvla" and cfg.policy.mode == "tokens": + GripperPostProcessTransform = _gripper_postprocess_transform_cls() + transforms.append( + GripperPostProcessTransform( + action_key="action", + rescale=True, + binarize=cfg.policy.gripper_binarize, + threshold=cfg.policy.gripper_binarize_threshold, + invert=cfg.policy.gripper_invert, + ) + ) + transforms.append(ActionTokenizerTransform(tokenizer)) + else: + transforms = [MultiAction.from_vla(stack_rewards=False)] + transforms.extend( + [ + SuccessReward(), + StepCounter(max_steps=cfg.env.max_outer_steps), + ] ) + return Compose(*transforms) def _make_libero_worker( - cfg, worker_idx: int, *, group_repeats=None, eval_mode=False, from_pixels=False + cfg, + worker_idx: int, + *, + group_repeats=None, + eval_mode=False, + from_pixels=False, + worker_idx_offset: int = 0, + render_gpu_device_id: int | None = None, ): _configure_mujoco_rendering(cfg) + worker_idx = int(worker_idx) + int(worker_idx_offset) task_id, worker_group_repeats, group_id_offset = _libero_worker_assignment( cfg, worker_idx, group_repeats=group_repeats, eval_mode=eval_mode ) parallel_group_repeats = ( not eval_mode and group_repeats is not None - and bool(cfg.env.get("parallel_group_repeats", False)) + and bool(cfg.env.parallel_group_repeats) ) - env_kwargs = dict(cfg.env.get("env_kwargs", None) or {}) + env_kwargs = dict(cfg.env.env_kwargs or {}) render_gpu_device_id = _worker_render_gpu_device_id( - cfg, worker_idx, eval_mode=eval_mode + cfg, + worker_idx, + eval_mode=eval_mode, + render_gpu_device_id=render_gpu_device_id, ) - if render_gpu_device_id is not None: - env_kwargs["render_gpu_device_id"] = render_gpu_device_id - return LiberoEnv( - cfg.env.task_suite, - task_id=task_id, - camera_height=cfg.env.camera_height, - camera_width=cfg.env.camera_width, - env_kwargs=env_kwargs, - wrist_camera="robot0_eye_in_hand" if cfg.policy.use_wrist_image else None, - from_pixels=from_pixels, - max_episode_steps=cfg.env.max_env_steps, - init_state_mode=( - "cycle" if eval_mode else cfg.env.get("train_init_state_mode", "random") + with _worker_render_gpu_context(cfg, render_gpu_device_id) as egl_device_id: + if egl_device_id is not None: + env_kwargs["render_gpu_device_id"] = egl_device_id + return LiberoEnv( + cfg.env.task_suite, + task_id=task_id, + camera_height=cfg.env.camera_height, + camera_width=cfg.env.camera_width, + env_kwargs=env_kwargs, + wrist_camera="robot0_eye_in_hand" if cfg.policy.use_wrist_image else None, + from_pixels=from_pixels, + max_episode_steps=cfg.env.max_env_steps, + init_state_mode=("cycle" if eval_mode else cfg.env.train_init_state_mode), + init_state_id=int(_cfg_get(cfg.env, "train_init_state_id", 0)), + group_repeats=worker_group_repeats, + group_id_offset=group_id_offset, + group_id_mode="init_state" if parallel_group_repeats else "episode", + ) + + +def _num_envs_from_cfg(cfg, *, eval_mode: bool = False) -> int: + if cfg.env.backend == "toy": + return 1 + return int(_cfg_get(cfg.env, "eval_num_envs" if eval_mode else "num_envs", 1)) + + +def _validate_libero_env_count( + cfg, num_envs: int, *, group_repeats=None, eval_mode: bool = False, override=False +) -> None: + task_ids = list(cfg.env.task_ids) + parallel_group_repeats = ( + not eval_mode + and group_repeats is not None + and bool(cfg.env.parallel_group_repeats) + ) + task_coverage_envs = num_envs + if parallel_group_repeats: + group_repeats = int(group_repeats) + candidate_repeats = candidate_group_size(cfg) + if candidate_repeats % group_repeats: + raise ValueError( + "collector.candidate_group_size must be a multiple of " + "collector.group_size when env.parallel_group_repeats=true " + f"({candidate_repeats=} and {group_repeats=})." + ) + if num_envs % group_repeats: + raise ValueError( + "env.num_envs must be a multiple of collector.group_size " + "when env.parallel_group_repeats=true so every parallel " + f"group has exactly {group_repeats} workers ({num_envs=})." + ) + if cfg.env.train_init_state_mode == "random": + raise ValueError( + "env.parallel_group_repeats=true requires " + "env.train_init_state_mode='cycle' or 'fixed'. Random " + "init-state sampling is local to each worker, so workers " + "sharing a group id would not necessarily share the same " + "initial state." + ) + task_coverage_envs = num_envs // group_repeats + if not override and task_coverage_envs < len(task_ids): + raise ValueError( + f"{'eval_num_envs' if eval_mode else 'num_envs'} ({num_envs}) " + f"must cover task_ids ({len(task_ids)} tasks): each worker is " + "bound to one task; fewer workers would silently drop tasks. " + "With env.parallel_group_repeats=true, task coverage is " + f"num_envs / collector.group_size ({task_coverage_envs})." + ) + if not override and task_coverage_envs % len(task_ids): + warnings.warn( + f"effective task workers ({task_coverage_envs}) is not a " + f"multiple of the number of tasks ({len(task_ids)}): tasks " + "will be sampled unevenly." + ) + + +def _make_env_worker( + cfg, + tokenizer: ActionTokenizerBase | None, + worker_idx: int, + *, + group_repeats: int | None = None, + seed: int | None = None, + device: torch.device | None = None, + eval_mode: bool = False, + from_pixels: bool = False, + decode_actions_in_env: bool = True, + worker_idx_offset: int = 0, + render_gpu_device_id: int | None = None, +) -> TransformedEnv: + worker_idx_with_offset = int(worker_idx) + int(worker_idx_offset) + worker_seed = None if seed is None else int(seed) + worker_idx_with_offset + if cfg.env.backend == "toy": + base = ToyVLAEnv( + action_dim=cfg.env.action_dim, + state_dim=cfg.env.state_dim, + image_shape=tuple(cfg.env.image_shape), + from_pixels=from_pixels, + render_size=cfg.env.render_size, + success_steps=cfg.env.success_steps, + success_tol=cfg.env.success_tol, + group_repeats=group_repeats, + group_id_offset=worker_idx_with_offset * GROUP_ID_OFFSET, + batch_size=[], + seed=worker_seed, + device=device, + ) + elif cfg.env.backend == "libero": + base = _make_libero_worker( + cfg, + worker_idx, + group_repeats=group_repeats, + eval_mode=eval_mode, + from_pixels=from_pixels, + worker_idx_offset=worker_idx_offset, + render_gpu_device_id=render_gpu_device_id, + ) + if worker_seed is not None: + base.set_seed(worker_seed) + else: + raise ValueError(f"Unknown env backend {cfg.env.backend!r}.") + return TransformedEnv( + base, + _chunk_transform( + cfg, + tokenizer, + decode_actions_in_env=decode_actions_in_env, ), - group_repeats=worker_group_repeats, - group_id_offset=group_id_offset, - group_id_mode="init_state" if parallel_group_repeats else "episode", ) def make_env( cfg, - tokenizer: ActionTokenizerBase, + tokenizer: ActionTokenizerBase | None, *, group_repeats: int | None = None, seed: int | None = None, @@ -257,6 +590,9 @@ def make_env( eval_mode: bool = False, from_pixels: bool = False, num_envs: int | None = None, + decode_actions_in_env: bool = True, + worker_idx_offset: int = 0, + render_gpu_device_id: int | None = None, ) -> TransformedEnv: if cfg.env.backend == "toy": base = ToyVLAEnv( @@ -264,7 +600,7 @@ def make_env( state_dim=cfg.env.state_dim, image_shape=tuple(cfg.env.image_shape), from_pixels=from_pixels, - render_size=cfg.env.get("render_size", 64), + render_size=cfg.env.render_size, success_steps=cfg.env.success_steps, success_tol=cfg.env.success_tol, group_repeats=group_repeats, @@ -278,55 +614,14 @@ def make_env( # to the train/eval envs sized from the config override = num_envs is not None if num_envs is None: - num_envs = cfg.env.eval_num_envs if eval_mode else cfg.env.num_envs - task_ids = list(cfg.env.task_ids) - parallel_group_repeats = ( - not eval_mode - and group_repeats is not None - and bool(cfg.env.get("parallel_group_repeats", False)) + num_envs = _num_envs_from_cfg(cfg, eval_mode=eval_mode) + _validate_libero_env_count( + cfg, + num_envs, + group_repeats=group_repeats, + eval_mode=eval_mode, + override=override, ) - # each worker hosts ONE MuJoCo task for its whole lifetime: fewer - # workers than tasks would silently drop tasks from the run - task_coverage_envs = num_envs - if parallel_group_repeats: - group_repeats = int(group_repeats) - candidate_repeats = candidate_group_size(cfg) - if candidate_repeats % group_repeats: - raise ValueError( - "collector.candidate_group_size must be a multiple of " - "collector.group_size when env.parallel_group_repeats=true " - f"({candidate_repeats=} and {group_repeats=})." - ) - if num_envs % group_repeats: - raise ValueError( - "env.num_envs must be a multiple of collector.group_size " - "when env.parallel_group_repeats=true so every parallel " - f"group has exactly {group_repeats} workers " - f"({num_envs=})." - ) - if cfg.env.get("train_init_state_mode", "random") == "random": - raise ValueError( - "env.parallel_group_repeats=true requires " - "env.train_init_state_mode='cycle' or 'fixed'. Random " - "init-state sampling is local to each worker, so workers " - "sharing a group id would not necessarily share the same " - "initial state." - ) - task_coverage_envs = num_envs // group_repeats - if not override and task_coverage_envs < len(task_ids): - raise ValueError( - f"{'eval_num_envs' if eval_mode else 'num_envs'} ({num_envs}) " - f"must cover task_ids ({len(task_ids)} tasks): each worker is " - "bound to one task; fewer workers would silently drop tasks. " - "With env.parallel_group_repeats=true, task coverage is " - f"num_envs / collector.group_size ({task_coverage_envs})." - ) - if not override and task_coverage_envs % len(task_ids): - warnings.warn( - f"effective task workers ({task_coverage_envs}) is not a " - f"multiple of the number of tasks ({len(task_ids)}): tasks " - "will be sampled unevenly." - ) base = ParallelEnv( num_envs, [ @@ -337,6 +632,8 @@ def make_env( group_repeats=group_repeats, eval_mode=eval_mode, from_pixels=from_pixels, + worker_idx_offset=worker_idx_offset, + render_gpu_device_id=render_gpu_device_id, ) for worker_idx in range(num_envs) ], @@ -350,7 +647,46 @@ def make_env( base.set_seed(seed) else: raise ValueError(f"Unknown env backend {cfg.env.backend!r}.") - return TransformedEnv(base, _chunk_transform(cfg, tokenizer)) + return TransformedEnv( + base, + _chunk_transform( + cfg, + tokenizer, + decode_actions_in_env=decode_actions_in_env, + ), + ) + + +def _make_collector_env( + cfg, + tokenizer: ActionTokenizerBase | None, + *, + num_envs: int, + group_repeats: int, + seed: int, + device: torch.device, + worker_idx_offset: int, + render_gpu_device_id: int | None, +) -> ParallelEnv: + return ParallelEnv( + num_envs, + [ + partial( + _make_env_worker, + cfg, + tokenizer, + worker_idx, + group_repeats=group_repeats, + seed=seed, + device=device if cfg.env.backend == "toy" else None, + worker_idx_offset=worker_idx_offset, + render_gpu_device_id=render_gpu_device_id, + ) + for worker_idx in range(num_envs) + ], + mp_start_method="spawn", + device=device, + ) def make_replay_buffer( @@ -361,38 +697,30 @@ def make_replay_buffer( # trajectories, the read path samples decisions without replacement. The # advantage transform is returned too so the training loop can flush its # incomplete-group queues at iteration boundaries. - collector_get = getattr( - cfg.collector, - "get", - lambda key, default=None: getattr(cfg.collector, key, default), - ) - max_collect_batches_per_iter = max( - int(collector_get("max_collect_batches_per_iter", 1)), 1 - ) + capacity_group_waves = max(int(cfg.buffer.capacity_group_waves), 1) candidate_size = candidate_group_size(cfg) capacity = ( cfg.collector.groups_per_iter * candidate_size * cfg.env.max_outer_steps - * max_collect_batches_per_iter + * capacity_group_waves ) keep_return_bounds = cfg.advantage.keep_return_bounds if keep_return_bounds is not None: keep_return_bounds = tuple(keep_return_bounds) - advantage_get = getattr( - cfg.advantage, - "get", - lambda key, default=None: getattr(cfg.advantage, key, default), - ) - selector_strategy = advantage_get("candidate_selection", "balanced") - selector_max_combinations = int( - advantage_get("candidate_selection_max_combinations", 100_000) - ) - candidate_selection_min_size = advantage_get("candidate_selection_min_size", None) + selector_strategy = cfg.advantage.candidate_selection + selector_max_combinations = int(cfg.advantage.candidate_selection_max_combinations) + candidate_selection_min_size = cfg.advantage.candidate_selection_min_size + shared_init = bool(cfg.buffer.shared_init) rb = TensorDictReplayBuffer( - storage=LazyTensorStorage(capacity, device=device), - sampler=SamplerWithoutReplacement(drop_last=False), + storage=LazyTensorStorage( + capacity, + device=device, + shared_init=shared_init, + ), batch_size=cfg.loss.mini_batch_size, + consume_after_n_samples=cfg.buffer.consume_after_n_samples, + shared=shared_init, ) advantage = MCAdvantage( grpo_size=cfg.collector.group_size, @@ -411,6 +739,13 @@ def make_replay_buffer( def make_loss_module(cfg, policy: VLAWrapperBase) -> ClipPPOLoss: + if policy.action_head != "tokens": + raise NotImplementedError( + "VLA GRPO training currently expects a token policy with stored " + "action log-probabilities. policy.mode='l1' is available for " + "reference/evaluation rollouts; add a continuous-action loss " + "before using it for training." + ) clip_epsilon = cfg.loss.clip_epsilon if not isinstance(clip_epsilon, float): clip_epsilon = tuple(clip_epsilon) @@ -439,172 +774,224 @@ def make_optimizer(cfg, loss_module: ClipPPOLoss): return optim, scheduler +_WEIGHT_STRATEGY = WeightStrategy(extract_as="tensordict") + + +def policy_weights( + policy: torch.nn.Module, *, trainable_only: bool = False +) -> TensorDictBase: + """Detached CPU TensorDict snapshot of a policy's parameters and buffers. + + With ``trainable_only=True`` the snapshot keeps only the tensors with + ``requires_grad=True`` (e.g. the LoRA adapters), which is enough for + weight syncs to a server policy built from the same checkpoint: the + receiving :class:`~torchrl.weight_update.WeightStrategy` applies a partial + TensorDict in place, leaving frozen weights and buffers untouched. + """ + weights = _WEIGHT_STRATEGY.extract_weights(policy) + if trainable_only: + weights = weights.apply( + lambda tensor: tensor if tensor.requires_grad else None, + filter_empty=True, + ) + return weights.data.detach().clone().cpu() + + +def sync_policy_server( + policy_server: ProcessInferenceServer, + policy: torch.nn.Module, + *, + trainable_only: bool = True, +): + """Push trainer policy weights to the shared process policy server. + + ``trainable_only=True`` (see ``train.weight_sync_trainable_only``) ships + only the ``requires_grad=True`` parameters through the control channel; + with a 7B LoRA policy that is a few hundred MB instead of the ~14 GB full + bf16 parameter set. + """ + return policy_server.update_model_weights( + policy_weights(policy, trainable_only=trainable_only) + ) + + +def apply_policy_weights( + policy: torch.nn.Module, weights: TensorDictBase, device: torch.device +) -> None: + """Apply a TensorDict parameter snapshot to ``policy`` on ``device``.""" + _WEIGHT_STRATEGY.apply_weights(policy, weights.to(device)) + + +def _training_group_repeats(cfg) -> int: + return ( + cfg.collector.group_size + if cfg.env.parallel_group_repeats + else candidate_group_size(cfg) + ) + + +def _render_gpu_for_subcollector(cfg, collector_idx: int) -> int | None: + render_gpu_ids = cfg.env.render_gpu_ids + if render_gpu_ids is None: + return None + render_gpu_ids = [int(device_id) for device_id in render_gpu_ids] + if not render_gpu_ids: + return None + return render_gpu_ids[int(collector_idx) % len(render_gpu_ids)] + + +def _server_config_from_collector(cfg) -> InferenceServerConfig: + return InferenceServerConfig( + max_batch_size=int(cfg.collector.server_max_batch_size), + min_batch_size=int(cfg.collector.server_min_batch_size), + timeout=float(cfg.collector.server_timeout), + collect_stats=bool(cfg.collector.server_collect_stats), + stats_window_size=int(cfg.collector.server_stats_window_size), + ) + + def make_collector( cfg, - env: EnvBase, policy: VLAWrapperBase, device: torch.device, *, - replay_buffer: TensorDictReplayBuffer | None = None, + tokenizer: ActionTokenizerBase | None, + replay_buffer: TensorDictReplayBuffer, post_collect_hook: Callable[[TensorDictBase], None] | None = None, -) -> Collector: - """Endless synchronous collector assembling complete trajectories. - - With no replay buffer, each yielded batch holds exactly one iteration's - worth of complete, done-terminated trajectories, concatenated along time - (``trajs_per_batch`` with ``traj_format="cat"``: flat and unpadded, - episodes delimited by the done flags -- no padding frames for the - image-heavy VLA observations; episodes spanning internal collection - steps are reassembled by the collector, in-flight episodes are held - back). - - With a replay buffer, TorchRL's collector writer path is used instead: - complete trajectories are pushed to the buffer as each internal rollout - batch finishes, and the iterator yields ``None``. This lets the replay - buffer transform keep incomplete GRPO groups across same-policy collection - polls until enough useful decisions have reached storage. - - The policy is held by reference (in-place optimizer updates apply - immediately) and observations/actions are cast between the env's and the - policy's devices by the collector. ``exploration_type=RANDOM`` makes the - collector roll out under a sampling context, which the token policy reads - via :func:`~torchrl.envs.utils.exploration_type` -- no policy mutation, so - this works with any collector (including multi-process workers). +) -> tuple[MultiCollector, ProcessInferenceServer, PolicyClientModule]: + """Build the TorchRL-native VLA rollout stack. + + The stack is always a ``MultiCollector`` whose workers own batched + ``ParallelEnv`` instances and call a shared process policy server through + ``PolicyClientModule``. The collector writes complete trajectories directly + into the replay buffer. """ - num_envs = env.batch_size[0] if env.batch_size else 1 - groups_per_iter = int(cfg.collector.groups_per_iter) + num_collectors = int(cfg.collector.num_collectors) + envs_per_collector = int(cfg.collector.envs_per_collector) group_size = int(cfg.collector.group_size) candidate_size = candidate_group_size(cfg) - env_get = getattr( - cfg.env, - "get", - lambda key, default=None: getattr(cfg.env, key, default), - ) - parallel_group_repeats = bool(env_get("parallel_group_repeats", False)) - group_workers = num_envs - if parallel_group_repeats: + total_envs = num_collectors * envs_per_collector + if cfg.env.parallel_group_repeats: if candidate_size % group_size: raise ValueError( "collector.candidate_group_size must be a multiple of " "collector.group_size when env.parallel_group_repeats=true " f"({candidate_size=} and {group_size=})." ) - if num_envs % group_size: + if envs_per_collector % group_size and not replay_buffer.shared: raise ValueError( - "env.num_envs must be a multiple of collector.group_size when " - "env.parallel_group_repeats=true " - f"({num_envs=} and {group_size=})." + "collector.envs_per_collector must be a multiple of " + "collector.group_size when env.parallel_group_repeats=true " + "and the replay buffer does not share grouped write state " + f"({envs_per_collector=} and {group_size=})." ) - group_workers = num_envs // group_size + + group_workers = ( + total_envs // group_size if cfg.env.parallel_group_repeats else total_envs + ) + if cfg.env.backend == "libero": + _validate_libero_env_count( + cfg, + total_envs, + group_repeats=_training_group_repeats(cfg), + ) + groups_per_iter = int(cfg.collector.groups_per_iter) if groups_per_iter < group_workers: raise ValueError( - "collector.groups_per_iter must be at least the number of parallel " - "group workers. Each worker emits repeated rollouts for its " - "own group ids (or, with env.parallel_group_repeats=true, each " - "logical worker group emits parallel rollouts for its group ids); " - "with fewer groups than workers, no worker can " - f"complete a full GRPO group in one iteration ({groups_per_iter=} " - f"< {group_workers=}), so the dynamic-sampling replay buffer would stay " - "empty. Reduce env.num_envs or increase collector.groups_per_iter." + "collector.groups_per_iter must be at least the number of " + "shared policy-server group workers " + f"({groups_per_iter=} < {group_workers=})." ) if groups_per_iter % group_workers: warnings.warn( - "collector.groups_per_iter is not a multiple of the number of " - "parallel group workers. Some workers will start a partial GRPO " - "group near the iteration boundary; Collector.reset() drops those " - "incomplete groups before the policy update. For best throughput, " - "set the group-worker count to a divisor of collector.groups_per_iter, " - "ideally the same value." + "collector.groups_per_iter is not a multiple of the shared " + "policy-server group-worker count. Some same-policy partial " + "groups can be dropped at the update boundary." ) - if parallel_group_repeats and groups_per_iter != group_workers: - warnings.warn( - "With env.parallel_group_repeats=true, collector.groups_per_iter " - "should usually match the number of logical parallel group workers " - "(env.num_envs / collector.group_size). If each logical worker must " - "advance through multiple group ids inside one collector batch, " - "variable episode lengths can create partial groups that are dropped " - "at the policy-update boundary. To overcollect, prefer additional " - "collector batches via collector.max_collect_batches_per_iter and " - "collector.min_replay_decisions, but note that multiple consecutive " - "batches can still drift without a group barrier; one aligned group " - "wave per update gives the lowest boundary waste." + + env_device = torch.device("cpu") + ctx = mp.get_context("spawn") + transport = MPTransport(ctx=ctx, use_manager=True) + eval_client = transport.client() + rollout_clients = [transport.client() for _ in range(num_collectors)] + policy_micro_batch_size = _cfg_get(cfg.collector, "policy_micro_batch_size", None) + server = ProcessInferenceServer( + policy_factory=partial( + make_policy, + cfg, + device, + policy_micro_batch_size=policy_micro_batch_size, + ), + transport=transport, + server_config=_server_config_from_collector(cfg), + policy_device=device, + output_device=env_device, + mp_context=ctx, + ).start() + + policy_in_keys = policy.in_keys + policy_out_keys = [*policy.out_keys, "policy_version"] + env_factories = [ + partial( + _make_collector_env, + cfg, + tokenizer, + num_envs=envs_per_collector, + group_repeats=_training_group_repeats(cfg), + seed=cfg.env.seed, + device=env_device, + worker_idx_offset=collector_idx * envs_per_collector, + render_gpu_device_id=_render_gpu_for_subcollector(cfg, collector_idx), + ) + for collector_idx in range(num_collectors) + ] + policy_client_factories = [ + partial( + PolicyClientModule, + client, + in_keys=policy_in_keys, + out_keys=policy_out_keys, + max_inflight=cfg.collector.max_inflight_per_env, ) - # Some transformed batched envs (notably LIBERO's ParallelEnv wrapped in - # TransformedEnv) deliberately expose a device-less outer TensorDict even - # though the simulator action path is CPU-only. If only ``policy_device`` - # is passed, Collector keeps CUDA policy outputs on the carrier and hands - # CUDA ``("vla_action", "tokens")`` to the env inverse transforms. The - # rollout still runs, but the decoded CPU simulator actions can silently - # diverge from - # env.rollout(auto_cast_to_device=True). Force the env side to CPU when the - # env does not advertise a device so sampled tokens are copied back before - # MuJoCo/ParallelEnv stepping. - env_device = env.device if env.device is not None else torch.device("cpu") - # With a replay buffer, use one outer step per internal collector poll so - # complete trajectories are handed to the replay-buffer transform as soon - # as they finish instead of waiting for a full max-length rollout from - # every worker. Without a replay buffer, keep the historical full-episode - # polling granularity so direct iteration yields whole group waves. - frames_per_batch = ( - num_envs if replay_buffer is not None else (num_envs * cfg.env.max_outer_steps) - ) - return Collector( - env, - policy, - frames_per_batch=frames_per_batch, + for client in rollout_clients + ] + eval_policy = PolicyClientModule( + eval_client, + in_keys=policy_in_keys, + out_keys=policy_out_keys, + ) + collector = MultiCollector( + env_factories, + policy=None, + policy_factory=policy_client_factories, + frames_per_batch=envs_per_collector, total_frames=-1, - trajs_per_batch=groups_per_iter * candidate_size, + reset_at_each_iter=False, + replay_buffer=replay_buffer, + trajs_per_batch=1, traj_format="cat", exploration_type=ExplorationType.RANDOM, - policy_device=device, + policy_device=env_device, env_device=env_device, - reset_at_each_iter=False, - replay_buffer=replay_buffer, + storing_device=torch.device(cfg.collector.storing_device), + trust_policy=True, + use_buffers=cfg.collector.use_buffers, + sync=False, + num_threads=int(cfg.collector.num_threads), + num_sub_threads=int(cfg.collector.env_sub_threads), post_collect_hook=post_collect_hook, ) + return collector, server, eval_policy -def evaluate(env: TransformedEnv, policy: VLAWrapperBase, cfg) -> float: - """Greedy success rate over (at least) ``cfg.logger.eval_episodes`` episodes. - - One evaluation round = one reset per env row + one episode per row, with - no auto-resets in between (``break_when_all_done`` freezes finished rows - and stops once every row is done). Each LIBERO reset therefore consumes - exactly one cycled initial state, keeping the fixed-trials evaluation - protocol exact, and no collected episode is discarded. Greedy decoding is - requested through the exploration context, not by mutating the policy. - """ - successes = 0.0 - episodes = 0 - with set_exploration_type(ExplorationType.DETERMINISTIC), torch.no_grad(): - while episodes < cfg.logger.eval_episodes: - reset_td = env.reset() - rollout = env.rollout( - cfg.env.max_outer_steps, - policy, - break_when_any_done=False, - break_when_all_done=True, - auto_reset=False, - tensordict=reset_td, - auto_cast_to_device=True, - ) - # one episode per row: success anywhere along the (frozen-once- - # done) row - row_success = rollout["next", "success"].any(-2) - successes += float(row_success.sum()) - episodes += int(row_success.numel()) - return successes / max(episodes, 1) - - -def make_record_env(cfg, tokenizer: ActionTokenizerBase, logger, device): +def make_record_env(cfg, tokenizer: ActionTokenizerBase | None, logger, device): """Single-environment eval recorder feeding a torchrl ``VideoRecorder``. Built with ``from_pixels=True`` so the base env emits a root ``pixels`` frame (``ToyVLAEnv`` renders the tracking scene; ``LiberoEnv`` exposes the camera). A :class:`~torchrl.record.VideoRecorder` transform appended last - collects those frames; :func:`record_eval_video` rolls out greedily and - flushes them to ``logger``. One environment keeps the video a single clean - stream (rather than a tiled grid of workers). + collects those frames during evaluator rollouts. One environment keeps the + video a single clean stream rather than a tiled grid of workers. """ env = make_env( cfg, @@ -628,29 +1015,407 @@ def make_record_env(cfg, tokenizer: ActionTokenizerBase, logger, device): return env, recorder -def record_eval_video(env, recorder, policy: VLAWrapperBase, cfg, step: int) -> None: - """Roll out ``cfg.logger.video_episodes`` greedy episodes and log one video. +def _make_eval_env( + cfg, + tokenizer: ActionTokenizerBase | None, + logger, + device: torch.device, +) -> TransformedEnv: + if logger is not None: + env, _ = make_record_env(cfg, tokenizer, logger, device) + return env + return make_env( + cfg, + tokenizer, + seed=cfg.env.seed + 1, + device=device if cfg.env.backend == "toy" else None, + eval_mode=True, + ) - Each reset consumes one cycled initial state and ``break_when_all_done`` - stops at the episode's natural end -- the same protocol as :func:`evaluate`, - so the recorded episodes mirror the measured ones. ``recorder.dump`` stacks - every frame seen since the last dump into a single clip and writes it to the - logger at ``step``. Greedy decoding is requested through the exploration - context, matching :func:`evaluate`. - """ - with set_exploration_type(ExplorationType.DETERMINISTIC), torch.no_grad(): - for _ in range(max(int(cfg.logger.video_episodes), 1)): - reset_td = env.reset() - env.rollout( - cfg.env.max_outer_steps, - policy, - break_when_any_done=False, - break_when_all_done=True, - auto_reset=False, - tensordict=reset_td, - auto_cast_to_device=True, + +def _policy_module_factory(policy, *args, **kwargs): + return policy + + +def _eval_success_metrics(rollout: TensorDictBase) -> dict[str, float]: + success = rollout.get(("next", "success"), None) + if success is None: + return {"success_rate": float("nan")} + mask = rollout.get(("collector", "mask"), None) + success = success.bool() + if mask is not None: + success = success & mask.unsqueeze(-1).expand_as(success) + success = success.reshape(success.shape[0], -1).any(-1) + else: + traj_ids = rollout.get(("collector", "traj_ids"), None) + if traj_ids is None: + success = success.reshape(1, -1).any(-1) + else: + traj_ids = traj_ids.reshape(-1) + success = torch.stack( + [ + success.reshape(success.shape[0], -1)[traj_ids == traj_id].any() + for traj_id in traj_ids.unique(sorted=True) + ] ) - recorder.dump(step=step) + return {"success_rate": float(success.float().mean())} + + +def make_evaluator( + cfg, + tokenizer: ActionTokenizerBase | None, + policy, + logger, + device: torch.device, +) -> Evaluator: + """Build the TorchRL evaluator used by the VLA GRPO recipe.""" + record_video = logger is not None and cfg.logger.eval_backend == "thread" + if record_video and cfg.env.backend == "libero": + task_ids = list(cfg.env.task_ids) + if len(task_ids) > 1 and not bool( + _cfg_get(cfg.logger, "record_video_single_task", False) + ): + raise ValueError( + "logger.eval_backend='thread' with a logger replaces the eval " + "env with a single-env video recorder bound to task " + f"{task_ids[0]}, so eval/success_rate would cover 1 of the " + f"{len(task_ids)} configured env.task_ids (and " + "env.eval_num_envs would be ignored). Set " + "logger.record_video_single_task=true to opt in to " + "single-task eval video, or keep logger.eval_backend=" + "'process' for suite-wide evaluation without video." + ) + env_factory = partial( + _make_eval_env, + cfg, + tokenizer, + logger if record_video else None, + device, + ) + return Evaluator( + env_factory, + policy=None, + policy_factory=partial(_policy_module_factory, policy), + num_trajectories=cfg.logger.eval_episodes, + max_steps=0, + frames_per_batch=cfg.env.max_outer_steps, + collector_kwargs={"traj_format": "cat"}, + log_prefix="eval", + reward_keys=("next", "reward"), + done_keys=("next", "done"), + device=device, + exploration_type=ExplorationType.DETERMINISTIC, + metrics_fn=_eval_success_metrics, + dump_video=record_video, + busy_policy=cfg.logger.eval_busy_policy, + backend=cfg.logger.eval_backend, + ) + + +def _sync_replay_sampler_writes(replay_buffer: TensorDictReplayBuffer) -> None: + """Mirror shared writer progress into the local consuming sampler state.""" + write_count = int(replay_buffer.write_count) + try: + previous_write_count = int(replay_buffer.__dict__["_vla_synced_write_count"]) + except KeyError: + previous_write_count = 0 + if write_count <= previous_write_count: + return + storage = replay_buffer._storage + storage_capacity = int(storage.max_size) + num_new_writes = min(write_count - previous_write_count, storage_capacity) + start = write_count - num_new_writes + indices = torch.arange(start, write_count, dtype=torch.long) + replay_buffer.mark_update(indices.remainder(storage_capacity)) + replay_buffer.__dict__["_vla_synced_write_count"] = write_count + + +def wait_for_replay( + replay_buffer: TensorDictReplayBuffer, + *, + min_replay_decisions: int, + poll_interval_s: float, + log_interval_s: float, + iteration: int, +) -> dict[str, float | int]: + """Wait until the replay buffer has enough sampleable decisions.""" + polls = 0 + next_log_s = 0.0 + with timeit("replay_wait") as wait_timer: + _sync_replay_sampler_writes(replay_buffer) + while len(replay_buffer) < min_replay_decisions: + time.sleep(poll_interval_s) + polls += 1 + _sync_replay_sampler_writes(replay_buffer) + elapsed = wait_timer.elapsed() + if elapsed >= next_log_s: + torchrl_logger.info( + "waiting for replay iteration %d decisions %d/%d " "elapsed_s %.1f", + iteration, + len(replay_buffer), + min_replay_decisions, + elapsed, + ) + next_log_s = elapsed + log_interval_s + elapsed = wait_timer.elapsed() + return { + "buffer/wait_polls": polls, + "buffer/wait_s": elapsed, + "buffer/decisions_before_update": len(replay_buffer), + } + + +def replay_ready_target(cfg) -> int: + """Number of sampleable decisions required before an update starts.""" + if cfg.collector.min_replay_decisions: + return int(cfg.collector.min_replay_decisions) + return ( + int(cfg.collector.groups_per_iter) + * candidate_group_size(cfg) + * int(cfg.env.max_outer_steps) + ) + + +def update( + replay_buffer: TensorDictReplayBuffer, + loss_module: ClipPPOLoss, + optim: torch.optim.Optimizer, + scheduler, + cfg, + device: torch.device, + *, + logger=None, + iteration: int, +) -> dict[str, float | int]: + """Run one PPO update over currently sampleable replay-buffer decisions.""" + num_decisions = len(replay_buffer) + target_samples = max( + num_decisions * int(cfg.buffer.consume_after_n_samples), + num_decisions, + ) + accumulate = max(int(cfg.loss.accumulate_batches), 1) + losses = [] + clip_fractions = [] + ess = [] + grad_norms = [] + optim_steps = 0 + trained_decisions = 0 + micro_batches = 0 + + def optimizer_step() -> None: + nonlocal optim_steps + grad_norms.append( + torch.nn.utils.clip_grad_norm_( + loss_module.parameters(), cfg.optim.max_grad_norm + ) + ) + optim.step() + optim.zero_grad(set_to_none=True) + scheduler.step() + optim_steps += 1 + + with timeit("train") as train_timer: + while trained_decisions < target_samples and len(replay_buffer): + batch = replay_buffer.sample() + batch = batch.to(device) + trained_decisions += batch.shape[0] + loss_vals = loss_module(batch) + loss = loss_vals["loss_objective"] / accumulate + loss.backward() + micro_batches += 1 + losses.append(loss_vals["loss_objective"].detach()) + clip_fractions.append(loss_vals["clip_fraction"].detach()) + # Token-level ESS retains action feature dimensions. Reduce each + # minibatch before aggregation so partial batches do not make the + # diagnostic shapes heterogeneous. + ess.append(loss_vals["ESS"].detach().mean()) + if micro_batches % accumulate == 0: + optimizer_step() + tail_micro_batches = micro_batches % accumulate + if tail_micro_batches: + # Each micro-batch loss was divided by the full `accumulate`, but + # the tail optimizer step accumulated fewer micro-batches. Rescale + # the accumulated gradients so the tail step averages over the + # micro-batches that actually contributed to it. + for param in loss_module.parameters(): + if param.grad is not None: + param.grad.mul_(accumulate / tail_micro_batches) + optimizer_step() + + train_time = max(train_timer.elapsed(), 1e-9) + metrics: dict[str, float | int] = { + "train/decisions": trained_decisions, + "train/micro_batches": micro_batches, + "train/optim_steps": optim_steps, + "throughput/train_decisions_per_s": trained_decisions / train_time, + "throughput/optim_steps_per_s": optim_steps / train_time, + } + if losses: + metrics.update( + { + "train/loss_objective": torch.stack(losses).mean().item(), + "train/clip_fraction": torch.stack(clip_fractions).mean().item(), + "train/ESS": torch.stack(ess).mean().item(), + "train/grad_norm": torch.stack(grad_norms).mean().item(), + } + ) + log_metrics(logger, metrics, iteration) + return metrics + + +_WORKER_ADVANTAGE_PATH = "replay_buffer._transform[0]" + + +def _advantage_stats( + advantage: MCAdvantage, collector: MultiCollector | None = None +) -> list[dict[str, float | int]]: + if collector is not None and not advantage.is_shared: + return collector.map_fn(f"{_WORKER_ADVANTAGE_PATH}.get_stats") + return [advantage.get_stats()] + + +def reset_advantage_state( + advantage: MCAdvantage, collector: MultiCollector | None = None +) -> None: + """Clear incomplete groups and reset counters at a policy boundary.""" + advantage.clear_queues() + advantage.reset_stats() + if collector is not None and not advantage.is_shared: + collector.map_fn(f"{_WORKER_ADVANTAGE_PATH}.clear_queues") + collector.map_fn(f"{_WORKER_ADVANTAGE_PATH}.reset_stats") + + +def reset_collection_state(advantage: MCAdvantage, collector: MultiCollector) -> None: + """Drop partial trajectories before advancing the behavior policy.""" + reset_advantage_state(advantage, collector) + collector.map_fn("reset") + + +def advantage_metrics( + advantage: MCAdvantage, collector: MultiCollector | None = None +) -> dict[str, float | int]: + """Compact metrics snapshot from GRPO replay-transform writer states.""" + stats = _advantage_stats(advantage, collector) + completed_trajectories = sum( + int(worker["completed_trajectories"]) for worker in stats + ) + return { + "buffer/complete_groups": sum( + int(worker["completed_groups"]) for worker in stats + ), + "buffer/kept_groups": sum(int(worker["written_groups"]) for worker in stats), + "buffer/skipped_groups": sum(int(worker["dropped_groups"]) for worker in stats), + "buffer/rescued_groups": sum(int(worker["rescued_groups"]) for worker in stats), + "buffer/queued_groups": sum(int(worker["queued_groups"]) for worker in stats), + "buffer/queued_trajectories": sum( + int(worker["queued_trajectories"]) for worker in stats + ), + "buffer/max_queued_trajectories_per_group": ( + max(int(worker["max_queued_trajectories_per_group"]) for worker in stats) + ), + "collector/completed_decisions": sum( + int(worker["completed_decisions"]) for worker in stats + ), + "collector/completed_trajectories": completed_trajectories, + "collector/successful_trajectories": sum( + int(worker["successful_trajectories"]) for worker in stats + ), + "collector/trajectory_return_sum": sum( + float(worker["trajectory_return_sum"]) for worker in stats + ), + "collector/trajectory_return_max": ( + max(float(worker["trajectory_return_max"]) for worker in stats) + if completed_trajectories + else 0.0 + ), + } + + +def iteration_metrics( + cfg, + *, + num_decisions: int, + total_episodes: int, + collect_metrics: dict[str, float | int], + group_metrics: dict[str, float | int], + train_metrics: dict[str, float | int], + eval_metrics: dict[str, float | int], + timings: dict[str, float], +) -> tuple[dict[str, float | int], float]: + """Merge per-iteration VLA GRPO metrics into the logger payload.""" + completed_trajectories = int(group_metrics["collector/completed_trajectories"]) + train_success = float(group_metrics["collector/successful_trajectories"]) / max( + completed_trajectories, 1 + ) + collect_time = max(timings["time/collect"], 1e-9) + completed_decisions = int(group_metrics["collector/completed_decisions"]) + metrics: dict[str, float | int] = { + "train/success_rate": train_success, + "train/episodes_total": total_episodes, + "buffer/decisions": num_decisions, + "throughput/inference_env_steps_per_s": ( + completed_decisions * cfg.env.chunk_size / collect_time + ), + "throughput/inference_decisions_per_s": completed_decisions / collect_time, + } + metrics.update(collect_metrics) + metrics.update(group_metrics) + metrics.update(train_metrics) + metrics.update(eval_metrics) + metrics.update(timings) + return metrics, train_success + + +def log_iteration_summary( + iteration: int, + *, + train_success: float, + num_decisions: int, + group_metrics: dict[str, float | int], + timings: dict[str, float], +) -> None: + """One-line training progress summary for stdout logs.""" + torchrl_logger.info( + "iteration %d success %.3f decisions %d kept_groups %d " + "skipped_groups %d queued_trajs %d collect_s %.1f train_s %.1f", + iteration, + train_success, + num_decisions, + int(group_metrics["buffer/kept_groups"]), + int(group_metrics["buffer/skipped_groups"]), + int(group_metrics["buffer/queued_trajectories"]), + timings["time/collect"], + timings["time/train"], + ) + + +def checkpoint_tensordict( + policy: torch.nn.Module, optim, scheduler, iteration: int +) -> TensorDict: + """TensorDict checkpoint payload for the VLA recipe.""" + return TensorDict( + { + "policy": policy_weights(policy), + "iteration": torch.tensor(iteration, dtype=torch.long), + "torch_rng_state": torch.get_rng_state(), + "optim": NonTensorData(optim.state_dict()), + "scheduler": NonTensorData(scheduler.state_dict()), + }, + batch_size=[], + ) + + +def save_checkpoint(path, policy: torch.nn.Module, optim, scheduler, iteration: int): + checkpoint_tensordict(policy, optim, scheduler, iteration).save(path) + + +def load_checkpoint(path, policy: torch.nn.Module, optim, scheduler, device) -> int: + checkpoint = TensorDict.load(path) + apply_policy_weights(policy, checkpoint["policy"], device) + optim.load_state_dict(checkpoint["optim"]) + scheduler.load_state_dict(checkpoint["scheduler"]) + torch.set_rng_state(checkpoint["torch_rng_state"]) + return int(checkpoint["iteration"].item()) + 1 def log_metrics(logger, metrics: dict, step: int) -> None: diff --git a/sota-implementations/vla_grpo/vla-grpo.py b/sota-implementations/vla_grpo/vla-grpo.py index 513d9ca6483..6ed7e3a96ab 100644 --- a/sota-implementations/vla_grpo/vla-grpo.py +++ b/sota-implementations/vla_grpo/vla-grpo.py @@ -6,856 +6,209 @@ This is the SimpleVLA-RL recipe (`arXiv:2509.09674 `_): a token-head VLA -policy emits a whole action chunk per forward (parallel decoding), -trajectories are collected in groups sharing the same initial state, the -advantage is the group-normalized binary success return broadcast to every -chunk decision, degenerate groups are dropped (dynamic sampling), and the -policy is updated with an asymmetric-clip PPO objective (no critic, no -KL-to-reference, no entropy bonus). - -The training-sample unit is the *decision* (one outer step of the -``MultiAction``-transformed environment = one chunk). Per-iteration -accounting: ``groups_per_iter`` initial states x ``candidate_group_size`` -rollout candidates per state, each contributing up to ``max_outer_steps`` -decisions; the dynamic sampling filter and candidate selector decide how many -trajectories are useful for optimization, so the effective batch is variable. - -Two configurations ship: - -- ``vla_grpo_toy.yaml`` (default): TinyVLA on the ToyVLAEnv tracking task, - single device, no simulator dependencies; exercised in the sota CI. -- ``vla_grpo_libero.yaml``: OpenVLA-OFT (token variant, 7B) on LIBERO with - the full SimpleVLA-RL hyper-parameters; parallel MuJoCo workers feed a - single training device. Multi-GPU sharded training (FSDP) of the 7B model - is the documented next step and should be sized on the target hardware; - the LoRA fallback (``policy.lora_rank``) fits a single device. +policy emits a whole action chunk per forward, trajectories are collected in +same-initial-state groups, the trajectory-level binary success return is +normalized within each group, and PPO updates the sampled action tokens with an +asymmetric clipping objective. """ from __future__ import annotations -import math -import multiprocessing as mp import os -import queue -import tempfile import warnings import hydra import torch import tqdm from torchrl._utils import logger as torchrl_logger, timeit -from torchrl.record.loggers import generate_exp_name, get_logger from utils import ( + advantage_metrics, + auto_device, candidate_group_size, - evaluate, + iteration_metrics, + load_checkpoint, + log_iteration_summary, log_metrics, make_action_tokenizer, make_collector, - make_env, + make_evaluator, + make_logger, make_loss_module, make_optimizer, make_policy, - make_record_env, make_replay_buffer, - record_eval_video, + replay_ready_target, + reset_collection_state, + save_checkpoint, + sync_policy_server, + update, + wait_for_replay, ) warnings.filterwarnings("ignore", category=UserWarning, module="tensordict") -def _trainable_policy_state_dict(policy) -> dict[str, torch.Tensor]: - """CPU state dict containing the trainable policy weights.""" - trainable_names = { - name for name, parameter in policy.named_parameters() if parameter.requires_grad - } - state_dict = policy.state_dict() - if not trainable_names: - trainable_names = set(state_dict) - return { - name: tensor.detach().cpu() - for name, tensor in state_dict.items() - if name in trainable_names - } - - -def _load_trainable_policy_state_dict( - policy, state_dict: dict[str, torch.Tensor] -) -> None: - if not state_dict: - raise RuntimeError("Policy checkpoint did not contain trainable weights.") - _, unexpected = policy.load_state_dict(state_dict, strict=False) - if unexpected: - raise RuntimeError( - f"Unexpected policy checkpoint keys: {', '.join(unexpected)}" - ) - - -def _eval_device(cfg) -> torch.device: - return torch.device( - cfg.logger.get("eval_device", None) - or cfg.policy.device - or ("cuda:0" if torch.cuda.is_available() else "cpu") - ) - - -def _evaluator_worker(request_queue, response_queue, cfg) -> None: - eval_env = None - try: - torch.manual_seed(cfg.env.seed + 1) - device = _eval_device(cfg) - policy = make_policy(cfg, device) - tokenizer = make_action_tokenizer(cfg, policy) - eval_env = make_env( - cfg, - tokenizer, - seed=cfg.env.seed + 1, - device=device if cfg.env.backend == "toy" else None, - eval_mode=True, - ) - with torch.no_grad(): - policy(eval_env.fake_tensordict().to(device)) - response_queue.put(("ready", None, None)) - while True: - request = request_queue.get() - if request is None: - break - iteration, state_path = request - try: - state_dict = torch.load(state_path, map_location="cpu") - _load_trainable_policy_state_dict(policy, state_dict) - success = evaluate(eval_env, policy, cfg) - response_queue.put((iteration, success, None)) - except Exception as err: - response_queue.put((iteration, None, repr(err))) - finally: - if os.path.exists(state_path): - os.unlink(state_path) - except Exception as err: - response_queue.put((-1, None, repr(err))) - finally: - if eval_env is not None: - eval_env.close(raise_if_closed=False) - - -class _EvaluatorProcess: - def __init__(self, cfg) -> None: - self._ctx = mp.get_context("spawn") - self._requests = self._ctx.Queue(maxsize=1) - self._responses = self._ctx.Queue(maxsize=1) - self._state_dir = ( - cfg.logger.get("eval_state_dir", None) or tempfile.gettempdir() - ) - self._pending_iteration = None - self._pending_state_path = None - os.makedirs(self._state_dir, exist_ok=True) - self._process = self._ctx.Process( - target=_evaluator_worker, - args=(self._requests, self._responses, cfg), - ) - self._process.start() - ready, _, error = self._recv() - if error is not None: - raise RuntimeError(f"Evaluator failed to initialize: {error}") - if ready != "ready": - raise RuntimeError(f"Unexpected evaluator init message: {ready!r}.") - - @property - def has_pending(self) -> bool: - return self._pending_iteration is not None - - @property - def pending_iteration(self) -> int | None: - return self._pending_iteration - - def submit(self, policy, iteration: int) -> bool: - if self.has_pending: - self._raise_if_failed() - return False - if not self._process.is_alive(): - self._raise_if_failed() - state_path = self._make_state_path(iteration) - torch.save(_trainable_policy_state_dict(policy), state_path) - self._requests.put((iteration, state_path)) - self._pending_iteration = iteration - self._pending_state_path = state_path - return True - - def poll(self) -> tuple[int, float] | None: - if not self.has_pending: - self._raise_if_failed() - return None - try: - response = self._responses.get_nowait() - except queue.Empty: - self._raise_if_failed() - return None - return self._handle_response(response) - - def wait(self) -> tuple[int, float]: - if not self.has_pending: - raise RuntimeError("No evaluator request is pending.") - return self._handle_response(self._recv()) - - def evaluate(self, policy, iteration: int) -> float: - if not self.submit(policy, iteration): - raise RuntimeError( - f"Evaluator already has a pending request for iteration " - f"{self._pending_iteration}." - ) - _, success = self.wait() - return success - - def close(self) -> None: - if self._pending_state_path is not None and os.path.exists( - self._pending_state_path - ): - os.unlink(self._pending_state_path) - self._pending_iteration = None - self._pending_state_path = None - if self._process.is_alive(): - self._requests.put(None) - self._process.join(timeout=30) - if self._process.is_alive(): - self._process.terminate() - self._process.join(timeout=30) - - def _make_state_path(self, iteration: int) -> str: - return os.path.join( - self._state_dir, - f"vla_grpo_eval_state_{os.getpid()}_{iteration}.pt", - ) - - def _handle_response(self, response) -> tuple[int, float]: - eval_iteration, success, error = response - state_path = self._pending_state_path - pending_iteration = self._pending_iteration - self._pending_iteration = None - self._pending_state_path = None - if state_path is not None and os.path.exists(state_path): - os.unlink(state_path) - if error is not None: - raise RuntimeError( - f"Evaluator failed at iteration {eval_iteration}: {error}" - ) - if eval_iteration != pending_iteration: - raise RuntimeError( - f"Evaluator returned iteration {eval_iteration}, " - f"expected {pending_iteration}." - ) - return eval_iteration, success - - def _recv(self): - while True: - try: - return self._responses.get(timeout=1.0) - except queue.Empty: - self._raise_if_failed() - - def _raise_if_failed(self) -> None: - if self._process.is_alive(): - return - try: - _, _, error = self._responses.get_nowait() - except queue.Empty as err: - raise RuntimeError("Evaluator process exited without a result.") from err - raise RuntimeError(f"Evaluator process failed to initialize: {error}") - - -def save_checkpoint(path, policy, optim, scheduler, iteration): - torch.save( - { - "policy": policy.state_dict(), - "optim": optim.state_dict(), - "scheduler": scheduler.state_dict(), - "iteration": iteration, - "torch_rng_state": torch.get_rng_state(), - }, - path, - ) - - -def load_checkpoint(path, policy, optim, scheduler) -> int: - checkpoint = torch.load(path, weights_only=False) - policy.load_state_dict(checkpoint["policy"]) - optim.load_state_dict(checkpoint["optim"]) - scheduler.load_state_dict(checkpoint["scheduler"]) - torch.set_rng_state(checkpoint["torch_rng_state"]) - return checkpoint["iteration"] + 1 - - @hydra.main(config_path="config", config_name="vla_grpo_toy", version_base="1.1") def main(cfg): # noqa: F821 torch.manual_seed(cfg.env.seed) - device = torch.device( - cfg.policy.device - if cfg.policy.device - else ("cuda:0" if torch.cuda.is_available() else "cpu") + train_device = auto_device(cfg.policy.device) + rollout_device = ( + torch.device(cfg.collector.policy_device) + if cfg.collector.policy_device + else train_device ) - - # Logger - logger = None - if cfg.logger.backend: - exp_name = generate_exp_name("VLA-GRPO", cfg.logger.exp_name) - logger = get_logger( - logger_type=cfg.logger.backend, - logger_name="vla_grpo_logging", - experiment_name=exp_name, - wandb_kwargs={ - "mode": cfg.logger.mode, - "config": dict(cfg), - "project": cfg.logger.project_name, - "group": cfg.logger.group_name, - }, - ) - - # Policy first (the LIBERO action codec lives in the checkpoint), then - # the environments: training rollouts are grouped (the same initial state - # is replayed group_size times and stamped with a group_id); evaluation - # uses fresh (cycled) initial states and greedy decoding. - policy = make_policy(cfg, device) - tokenizer = make_action_tokenizer(cfg, policy) - collection_group_size = candidate_group_size(cfg) - env_get = getattr( - cfg.env, - "get", - lambda key, default=None: getattr(cfg.env, key, default), - ) - train_group_repeats = ( - cfg.collector.group_size - if env_get("parallel_group_repeats", False) - else collection_group_size + eval_device = ( + torch.device(cfg.logger.eval_device) + if cfg.logger.eval_device + else rollout_device ) - train_env = make_env( - cfg, - tokenizer, - group_repeats=train_group_repeats, - seed=cfg.env.seed, - device=device if cfg.env.backend == "toy" else None, + buffer_device = ( + torch.device(cfg.buffer.device) if cfg.buffer.device else train_device ) - eval_process = bool(cfg.logger.get("eval_process", False)) - eval_env = None - if not eval_process: - eval_env = make_env( - cfg, - tokenizer, - seed=cfg.env.seed + 1, - device=device if cfg.env.backend == "toy" else None, - eval_mode=True, - ) - # materialize lazy layers (TinyVLA) on a spec-shaped fake observation: a - # real reset would consume a grouped init (train env) or a cycled eval - # state (eval env) - with torch.no_grad(): - fake_env = eval_env if eval_env is not None else train_env - fake_td = fake_env.fake_tensordict() - policy(fake_td.to(device)) - - rollout_device = torch.device(cfg.collector.get("policy_device", None) or device) - if rollout_device == device: - rollout_policy = policy - else: - rollout_policy = make_policy(cfg, rollout_device) - with torch.no_grad(): - rollout_policy(fake_td.to(rollout_device)) - _load_trainable_policy_state_dict( - rollout_policy, _trainable_policy_state_dict(policy) - ) - - # optional eval-video recorder: a separate single-env rollout rendered to - # pixels and written through a torchrl VideoRecorder (logged to wandb / - # tensorboard / csv). Off by default and only built when a logger exists. - record_env = recorder = None - if cfg.logger.record_video and eval_process: - warnings.warn( - "logger.record_video is disabled when logger.eval_process=true; " - "the evaluator process returns scalar success metrics only." - ) - elif cfg.logger.record_video and logger is not None: - record_env, recorder = make_record_env(cfg, tokenizer, logger, device) - evaluator = _EvaluatorProcess(cfg) if eval_process else None - async_eval = bool(eval_process and cfg.logger.get("eval_async", False)) - - buffer_device = torch.device(cfg.buffer.device) if cfg.buffer.device else device + logger = make_logger(cfg) + policy = make_policy(cfg, train_device) + tokenizer = make_action_tokenizer(cfg, policy) replay_buffer, advantage_transform = make_replay_buffer(cfg, buffer_device) loss_module = make_loss_module(cfg, policy) optim, scheduler = make_optimizer(cfg, loss_module) - collector = make_collector( + collector, policy_server, eval_policy = make_collector( cfg, - train_env, - rollout_policy, + policy, rollout_device, + tokenizer=tokenizer, replay_buffer=replay_buffer, ) - collector_iter = iter(collector) + evaluator = make_evaluator( + cfg, + tokenizer, + eval_policy, + logger, + eval_device, + ) start_iter = 0 if cfg.checkpoint.resume: - start_iter = load_checkpoint(cfg.checkpoint.resume, policy, optim, scheduler) - if rollout_policy is not policy: - _load_trainable_policy_state_dict( - rollout_policy, _trainable_policy_state_dict(policy) - ) + start_iter = load_checkpoint( + cfg.checkpoint.resume, policy, optim, scheduler, train_device + ) torchrl_logger.info( - f"Resumed from {cfg.checkpoint.resume} at iteration {start_iter}." + "Resumed from %s at iteration %d.", cfg.checkpoint.resume, start_iter ) - episodes_per_iter = cfg.collector.groups_per_iter * collection_group_size - max_collect_batches_per_iter = max( - int(cfg.collector.get("max_collect_batches_per_iter", 1)), 1 - ) - max_same_policy_collect_attempts = max( - int(cfg.collector.get("max_same_policy_collect_attempts", 1)), 1 - ) - min_replay_decisions = int(cfg.collector.get("min_replay_decisions", 0) or 0) - num_envs = train_env.batch_size[0] if train_env.batch_size else 1 - collector_frames_per_poll = max( - int(getattr(collector, "requested_frames_per_batch", num_envs)), 1 - ) - collect_polls_per_group_wave = max( - math.ceil( - episodes_per_iter * int(cfg.env.max_outer_steps) / collector_frames_per_poll - ), - 1, - ) - max_collect_polls_per_iter = ( - collect_polls_per_group_wave * max_collect_batches_per_iter - ) - accumulate = max(int(cfg.loss.accumulate_batches), 1) - pbar = tqdm.tqdm(total=cfg.collector.total_iters, initial=start_iter) + target_replay_decisions = replay_ready_target(cfg) + episodes_per_iter = cfg.collector.groups_per_iter * candidate_group_size(cfg) total_episodes = start_iter * episodes_per_iter + weight_sync_trainable_only = bool(cfg.train.weight_sync_trainable_only) + sync_policy_server(policy_server, policy, trainable_only=weight_sync_trainable_only) + collector.start() - iteration = start_iter - retry_same_policy = False - same_policy_collect_polls = 0 - same_policy_collect_time = 0.0 - same_policy_collect_attempts = 0 - same_policy_safety_cap_hits = 0 - while iteration < cfg.collector.total_iters: - # Collect under a fixed rollout policy until enough useful decisions - # have reached the replay buffer (or the safety cap is hit). The - # collector writer pushes complete trajectories directly to the replay - # buffer as internal rollout polls finish; MCAdvantage is the replay - # buffer transform, so incomplete GRPO groups stay queued across these - # same-policy polls instead of being discarded at a collector-batch - # boundary. Rollouts sample (rather than argmax) because the collector - # runs under exploration_type=RANDOM (set in make_collector); the policy - # reads that context, so the script never mutates it. - with timeit("collect") as collect_timer: - if not retry_same_policy: - advantage_transform.reset_stats() - same_policy_collect_polls = 0 - same_policy_collect_time = 0.0 - same_policy_collect_attempts = 0 - same_policy_safety_cap_hits = 0 - retry_same_policy = False - same_policy_collect_attempts += 1 - collect_polls = 0 - safety_cap_hit = False - progress_log_interval = max(collect_polls_per_group_wave // 4, 1) - while collect_polls < max_collect_polls_per_iter: - next(collector_iter) - collect_polls += 1 - replay_decisions = len(replay_buffer) - reached_replay_target = ( - min_replay_decisions > 0 - and replay_decisions >= min_replay_decisions - ) - should_log_progress = ( - collect_polls % progress_log_interval == 0 - or reached_replay_target - or collect_polls == max_collect_polls_per_iter + pbar = tqdm.tqdm(total=cfg.collector.total_iters, initial=start_iter) + try: + for iteration in range(start_iter, cfg.collector.total_iters): + with timeit("collect"): + collect_metrics = wait_for_replay( + replay_buffer, + min_replay_decisions=target_replay_decisions, + poll_interval_s=cfg.collector.replay_wait_s, + log_interval_s=cfg.collector.replay_log_s, + iteration=iteration, ) - if should_log_progress: - torchrl_logger.info( - "collection progress iteration %d polls %d/%d " - "waves %.2f replay_decisions %d/%d raw_decisions %d " - "completed_trajs %d kept_groups %d skipped_groups %d " - "rescued_groups %d queued_trajs %d max_queue %d " - "elapsed_s %.1f total_waves %.2f attempts %d", - iteration, - collect_polls, - max_collect_polls_per_iter, - collect_polls / collect_polls_per_group_wave, - replay_decisions, - min_replay_decisions, - advantage_transform.completed_decisions, - advantage_transform.completed_trajectories, - advantage_transform.written_groups, - advantage_transform.dropped_groups, - advantage_transform.rescued_groups, - advantage_transform.queued_trajectories, - advantage_transform.max_queued_trajectories_per_group, - collect_timer.elapsed(), - (same_policy_collect_polls + collect_polls) - / collect_polls_per_group_wave, - same_policy_collect_attempts, - ) - if min_replay_decisions > 0: - if reached_replay_target: - break - elif collect_polls >= collect_polls_per_group_wave: - break - else: - safety_cap_hit = True - if min_replay_decisions > 0 and len(replay_buffer) < min_replay_decisions: - safety_cap_hit = True + with collector.pause(): + num_decisions = len(replay_buffer) + train_metrics = update( + replay_buffer, + loss_module, + optim, + scheduler, + cfg, + train_device, + logger=logger, + iteration=iteration, + ) + group_metrics = advantage_metrics(advantage_transform, collector) + reset_collection_state(advantage_transform, collector) + sync_policy_server( + policy_server, policy, trainable_only=weight_sync_trainable_only + ) - attempt_collect_time = collect_timer.elapsed() - same_policy_collect_time += attempt_collect_time - same_policy_collect_polls += collect_polls - same_policy_safety_cap_hits += int(safety_cap_hit) - collect_group_waves = ( - same_policy_collect_polls / collect_polls_per_group_wave - ) - completed_trajectories = advantage_transform.completed_trajectories - completed_decisions = advantage_transform.completed_decisions - group_metrics = { - "buffer/collect_batches": collect_group_waves, - "buffer/collect_group_waves": collect_group_waves, - "buffer/collect_polls": same_policy_collect_polls, - "buffer/current_collect_polls": collect_polls, - "buffer/collect_attempts": same_policy_collect_attempts, - "buffer/collect_safety_cap_hit": float(same_policy_safety_cap_hits > 0), - "buffer/collect_safety_cap_hits": same_policy_safety_cap_hits, - "buffer/current_collect_safety_cap_hit": float(safety_cap_hit), - "buffer/max_same_policy_collect_attempts": ( - max_same_policy_collect_attempts - ), - "buffer/min_replay_decisions": min_replay_decisions, - "buffer/grpo_group_size": cfg.collector.group_size, - "buffer/candidate_group_size": collection_group_size, - "buffer/complete_groups": advantage_transform.completed_groups, - "buffer/complete_groups_written": advantage_transform.written_groups, - "buffer/kept_groups": advantage_transform.written_groups, - "buffer/dropped_dynamic_sampling_groups": ( - advantage_transform.dropped_groups - ), - "buffer/dropped_complete_groups": advantage_transform.dropped_groups, - "buffer/skipped_groups": advantage_transform.dropped_groups, - "buffer/rescued_oversampled_groups": ( - advantage_transform.rescued_groups - ), - "buffer/selected_trajectories": ( - advantage_transform.selected_trajectories - ), - "buffer/unselected_candidate_trajectories": ( - advantage_transform.unselected_trajectories - ), - "buffer/partial_groups": advantage_transform.queued_groups, - "buffer/queued_groups": advantage_transform.queued_groups, - "buffer/queued_incomplete_groups": advantage_transform.queued_groups, - "buffer/queued_trajectories": (advantage_transform.queued_trajectories), - "buffer/queued_incomplete_trajectories": ( - advantage_transform.queued_trajectories - ), - "buffer/max_queued_trajectories_per_group": ( - advantage_transform.max_queued_trajectories_per_group - ), - "collector/raw_decisions": completed_decisions, - "collector/completed_decisions": completed_decisions, - "collector/raw_trajectories": completed_trajectories, - "collector/completed_trajectories": completed_trajectories, - "collector/successful_trajectories": ( - advantage_transform.successful_trajectories - ), - "collector/trajectory_return_sum": ( - advantage_transform.trajectory_return_sum - ), - "collector/trajectory_return_max": ( - advantage_transform.trajectory_return_max - if completed_trajectories - else 0.0 - ), - } - # PPO update over the decisions that survived dynamic sampling, with - # gradient accumulation (micro-batches of mini_batch_size decisions) - num_decisions = len(replay_buffer) - insufficient_replay = ( - min_replay_decisions > 0 and num_decisions < min_replay_decisions - ) - if ( - insufficient_replay - and advantage_transform.queued_trajectories - and same_policy_collect_attempts < max_same_policy_collect_attempts - ): + eval_metrics = {} + if iteration % cfg.logger.eval_iter == 0: + if cfg.logger.eval_async: + evaluator.trigger_eval(weights=None, step=iteration) + else: + eval_metrics = evaluator.evaluate(weights=None, step=iteration) + log_metrics(logger, eval_metrics, iteration) + if cfg.logger.eval_async: + ready_eval = evaluator.poll(timeout=0) + if ready_eval is not None: + eval_metrics = ready_eval + log_metrics(logger, eval_metrics, iteration) + + timings = timeit.todict(prefix="time") timeit.erase() - torchrl_logger.warning( - "iteration %d hit the collection safety cap with only %d/%d " - "replay decisions and still has %d queued trajectories across " - "%d groups (max_queue=%d/%d). Keeping the replay buffer and " - "MCAdvantage queues, then continuing collection under the same " - "policy instead of treating this as a policy-update boundary. " - "attempt %d/%d collect_s %.1f", - iteration, - num_decisions, - min_replay_decisions, - advantage_transform.queued_trajectories, - advantage_transform.queued_groups, - advantage_transform.max_queued_trajectories_per_group, - collection_group_size, - same_policy_collect_attempts, - max_same_policy_collect_attempts, - same_policy_collect_time, + completed_trajectories = int( + group_metrics["collector/completed_trajectories"] ) - retry_same_policy = True - continue - if insufficient_replay and advantage_transform.queued_trajectories: - torchrl_logger.warning( - "iteration %d reached the same-policy collection attempt cap " - "with only %d/%d replay decisions and %d queued trajectories " - "across %d groups (max_queue=%d/%d). Proceeding with an " - "undersized policy update and clearing in-flight groups at the " - "policy boundary. attempts %d/%d collect_s %.1f", - iteration, - num_decisions, - min_replay_decisions, - advantage_transform.queued_trajectories, - advantage_transform.queued_groups, - advantage_transform.max_queued_trajectories_per_group, - collection_group_size, - same_policy_collect_attempts, - max_same_policy_collect_attempts, - same_policy_collect_time, + total_episodes += completed_trajectories + metrics, train_success = iteration_metrics( + cfg, + num_decisions=num_decisions, + total_episodes=total_episodes, + collect_metrics=collect_metrics, + group_metrics=group_metrics, + train_metrics=train_metrics, + eval_metrics=eval_metrics, + timings=timings, ) - total_episodes += completed_trajectories - losses = [] - clip_fractions = [] - ess = [] - grad_norms = [] - optim_steps = 0 - train_decisions = 0 - - def optimizer_step(grad_norms=grad_norms): - nonlocal optim_steps - grad_norms.append( - torch.nn.utils.clip_grad_norm_( - loss_module.parameters(), cfg.optim.max_grad_norm - ) + metrics["train/lr"] = scheduler.get_last_lr()[0] + log_metrics(logger, metrics, iteration) + log_iteration_summary( + iteration, + train_success=train_success, + num_decisions=num_decisions, + group_metrics=group_metrics, + timings=timings, ) - optim.step() - optim.zero_grad(set_to_none=True) - scheduler.step() - optim_steps += 1 - - micro_batches = 0 - with timeit("train"): - for _ in range(cfg.loss.ppo_epochs): - if not num_decisions: - break - for batch in replay_buffer: - batch = batch.to(device) - train_decisions += batch.shape[0] - # ratio_level="token" gives per-token importance ratios; - # ClipPPOLoss broadcasts the per-decision advantage over the - # token dims itself, so the script passes it through as-is. - loss_vals = loss_module(batch) - loss = loss_vals["loss_objective"] / accumulate - loss.backward() - micro_batches += 1 - losses.append(loss_vals["loss_objective"].detach()) - clip_fractions.append(loss_vals["clip_fraction"]) - # Token-level ESS retains action feature dimensions. Reduce - # each minibatch before aggregation so partial batches do - # not make the diagnostic shapes heterogeneous. - ess.append(loss_vals["ESS"].detach().mean()) - if micro_batches % accumulate == 0: - optimizer_step() - if micro_batches % accumulate: - optimizer_step() - replay_buffer.empty() - # Synchronous policy boundary: old-policy in-flight episodes and - # incomplete GRPO groups must not leak into the next rollout policy. - advantage_transform.queues.clear() - collector.reset() - if rollout_policy is not policy: - _load_trainable_policy_state_dict( - rollout_policy, _trainable_policy_state_dict(policy) + pbar.update(1) + pbar.set_description( + f"success {train_success:.2f} decisions {num_decisions}" ) - eval_success = None - eval_source_iteration = None - if async_eval and evaluator is not None: - eval_result = evaluator.poll() - if eval_result is not None: - eval_source_iteration, eval_success = eval_result - torchrl_logger.info( - "async eval completed for iteration %d success %.3f", - eval_source_iteration, - eval_success, + if ( + cfg.checkpoint.save_iter + and (iteration + 1) % cfg.checkpoint.save_iter == 0 + ): + save_checkpoint( + os.path.join(os.getcwd(), "checkpoint_latest"), + policy, + optim, + scheduler, + iteration, ) - if iteration % cfg.logger.eval_iter == 0: - if evaluator is None: - with timeit("eval"): - eval_success = evaluate(eval_env, policy, cfg) - eval_source_iteration = iteration - elif async_eval: - with timeit("eval_submit"): - eval_result = evaluator.poll() - if eval_result is not None: - eval_source_iteration, eval_success = eval_result - torchrl_logger.info( - "async eval completed for iteration %d success %.3f", - eval_source_iteration, - eval_success, - ) - if evaluator.submit(policy, iteration): - torchrl_logger.info( - "submitted async eval for iteration %d", iteration - ) - else: - torchrl_logger.warning( - "skipped async eval for iteration %d because " - "iteration %d is still pending", - iteration, - evaluator.pending_iteration, - ) - else: - with timeit("eval"): - eval_success = evaluator.evaluate(policy, iteration) - eval_source_iteration = iteration - if recorder is not None: - record_eval_video(record_env, recorder, policy, cfg, iteration) - timings = timeit.todict(prefix="time") - timings["time/collect_current_attempt"] = timings.get("time/collect", 0.0) - timings["time/collect"] = same_policy_collect_time - timeit.erase() - collect_time = max(timings.get("time/collect", 0.0), 1e-9) - train_time = max(timings.get("time/train", 0.0), 1e-9) - num_decisions_collected = int(group_metrics["collector/raw_decisions"]) - completed_episode_decisions = int( - group_metrics["collector/completed_decisions"] - ) - completed_trajectories = int(group_metrics["collector/completed_trajectories"]) - env_steps = num_decisions_collected * cfg.env.chunk_size - success_rate = float(group_metrics["collector/successful_trajectories"]) / max( - completed_trajectories, 1 - ) - episode_decisions = completed_episode_decisions / max(completed_trajectories, 1) - reward_mean = float(group_metrics["collector/trajectory_return_sum"]) / max( - completed_trajectories, 1 - ) - reward_max = float(group_metrics["collector/trajectory_return_max"]) - metrics = { - "train/success_rate": success_rate, - "train/episode_decisions": episode_decisions, - "train/episodes_total": total_episodes, - "train/lr": scheduler.get_last_lr()[0], - # reward curves: per-episode return (binary-success reward summed - # over the episode), averaged and best-of-batch - "train/reward_mean": reward_mean, - "train/reward_max": reward_max, - "buffer/decisions": num_decisions, - "buffer/useful_replay_decisions": num_decisions, - "buffer/kept_fraction": num_decisions / max(1, num_decisions_collected), - "collector/completed_episode_decisions": completed_episode_decisions, - "throughput/replay_decisions_per_s": num_decisions / collect_time, - # inference throughput: env steps / decisions generated per second - # during collection (policy rollout) - "throughput/inference_env_steps_per_s": env_steps / collect_time, - "throughput/inference_decisions_per_s": num_decisions_collected - / collect_time, - # training throughput: decisions consumed / optimizer steps taken - # per second during the PPO update - "throughput/train_decisions_per_s": train_decisions / train_time, - "throughput/optim_steps_per_s": optim_steps / train_time, - } - metrics.update(group_metrics) - metrics.update(timings) - if eval_success is not None: - metrics["eval/success_rate"] = eval_success - metrics["eval/source_iteration"] = eval_source_iteration - if losses: - metrics.update( - { - "train/loss_objective": torch.stack(losses).mean().item(), - "train/clip_fraction": torch.stack(clip_fractions).mean().item(), - "train/ESS": torch.stack(ess).mean().item(), - "train/grad_norm": torch.stack(grad_norms).mean().item(), - } - ) - torchrl_logger.info( - "iteration %d success %.3f decisions %d raw_decisions %d " - "raw_trajs %d collect_waves %.2f polls %d kept_groups %d " - "skipped_groups %d rescued_groups %d queued_trajs %d " - "max_queue %d collect_s %.1f train_s %.1f", - iteration, - metrics["train/success_rate"], - num_decisions, - num_decisions_collected, - int(metrics.get("collector/raw_trajectories", 0)), - float(metrics.get("buffer/collect_group_waves", 0.0)), - int(metrics.get("buffer/collect_polls", 0)), - int(metrics.get("buffer/kept_groups", 0)), - int(metrics.get("buffer/dropped_dynamic_sampling_groups", 0)), - int(metrics.get("buffer/rescued_oversampled_groups", 0)), - int(metrics.get("buffer/queued_trajectories", 0)), - int(metrics.get("buffer/max_queued_trajectories_per_group", 0)), - timings.get("time/collect", 0.0), - timings.get("time/train", 0.0), - ) - log_metrics(logger, metrics, iteration) - pbar.update(1) - pbar.set_description( - f"success {metrics['train/success_rate']:.2f} decisions {num_decisions}" - ) - - if cfg.checkpoint.save_iter and (iteration + 1) % cfg.checkpoint.save_iter == 0: + if cfg.checkpoint.save_iter: save_checkpoint( - os.path.join(os.getcwd(), "checkpoint_latest.pt"), + os.path.join(os.getcwd(), "checkpoint_latest"), policy, optim, scheduler, - iteration, + cfg.collector.total_iters - 1, ) - iteration += 1 - - if cfg.checkpoint.save_iter: - save_checkpoint( - os.path.join(os.getcwd(), "checkpoint_latest.pt"), - policy, - optim, - scheduler, - cfg.collector.total_iters - 1, - ) - if async_eval and evaluator is not None and evaluator.has_pending: - eval_source_iteration, final_success = evaluator.wait() + if cfg.logger.eval_async and evaluator.pending: + final_async = evaluator.wait(timeout=cfg.logger.eval_timeout_s) + if final_async is not None: + log_metrics(logger, final_async, cfg.collector.total_iters) + final_eval = evaluator.evaluate(weights=None, step=cfg.collector.total_iters) + log_metrics(logger, final_eval, cfg.collector.total_iters) torchrl_logger.info( - "async eval completed for iteration %d success %.3f", - eval_source_iteration, - final_success, - ) - if logger is not None: - log_metrics( - logger, - { - "eval/success_rate": final_success, - "eval/source_iteration": eval_source_iteration, - }, - cfg.collector.total_iters, - ) - pbar.close() - if logger is not None: - if evaluator is None: - final_success = evaluate(eval_env, policy, cfg) - else: - final_success = evaluator.evaluate(policy, cfg.collector.total_iters) - log_metrics( - logger, {"eval/success_rate": final_success}, cfg.collector.total_iters + "Final greedy success rate: %.3f", final_eval["eval/success_rate"] ) - torchrl_logger.info(f"Final greedy success rate: {final_success:.3f}") - if recorder is not None: - record_eval_video( - record_env, recorder, policy, cfg, cfg.collector.total_iters - ) - collector.shutdown() - train_env.close(raise_if_closed=False) - if eval_env is not None: - eval_env.close(raise_if_closed=False) - if record_env is not None: - record_env.close(raise_if_closed=False) - if evaluator is not None: - evaluator.close() + finally: + pbar.close() + evaluator.shutdown() + collector.shutdown() + policy_server.shutdown() + policy_server.transport.close() if __name__ == "__main__": diff --git a/sota-implementations/vla_grpo/vla-grpo.sbatch b/sota-implementations/vla_grpo/vla-grpo.sbatch index 143b7de8c59..c6f9004257f 100644 --- a/sota-implementations/vla_grpo/vla-grpo.sbatch +++ b/sota-implementations/vla_grpo/vla-grpo.sbatch @@ -1,12 +1,19 @@ #!/bin/bash -# Single-node launcher for the LIBERO configuration of the VLA GRPO recipe. -# The MuJoCo workers are CPU processes; the policy trains on one GPU (LoRA -# fallback) -- size cpus-per-task to env.num_envs + headroom. For multi-GPU -# EGL rendering, request more GPUs and override env.render_gpu_ids accordingly. +# Single-node, single-GPU launcher for the LIBERO configuration of the VLA +# GRPO recipe. The MuJoCo workers are CPU processes; the policy trains on one +# GPU (LoRA fallback) -- size cpus-per-task to +# collector.num_collectors * collector.envs_per_collector + headroom. +# The launch command below overrides the multi-GPU H100 defaults of +# vla_grpo_libero.yaml (train GPU 0, rollout server GPU 1, render GPUs 2-7) +# down to one GPU, as documented in the README hardware notes. With +# env.parallel_group_repeats=true the collector env total must stay at +# len(env.task_ids) * collector.group_size = 80 to cover every task, hence +# one collector with 80 envs. For multi-GPU EGL rendering, request more GPUs +# and override env.render_gpu_ids accordingly. #SBATCH --job-name=vla-grpo #SBATCH --nodes=1 #SBATCH --ntasks-per-node=1 -#SBATCH --cpus-per-task=32 +#SBATCH --cpus-per-task=96 #SBATCH --gres=gpu:1 #SBATCH --output=logs/vla-grpo-%j.out #SBATCH --error=logs/vla-grpo-%j.err @@ -22,5 +29,11 @@ export ROBOT_PLATFORM=LIBERO python sota-implementations/vla_grpo/vla-grpo.py \ --config-name vla_grpo_libero \ + policy.device=null \ + collector.policy_device=null \ + collector.num_collectors=1 \ + collector.envs_per_collector=80 \ + 'env.render_gpu_ids=[0]' \ + 'env.eval_render_gpu_ids=[0]' \ logger.group_name="${SLURM_JOB_ID}" \ "$@" diff --git a/test/collectors/test_evaluator.py b/test/collectors/test_evaluator.py index 2376c2fb7bd..5a447e493eb 100644 --- a/test/collectors/test_evaluator.py +++ b/test/collectors/test_evaluator.py @@ -14,7 +14,11 @@ from tensordict.nn import TensorDictModule from torch import nn from torchrl.collectors import Evaluator -from torchrl.collectors._evaluator import _freeze_vecnorm, _wrap_env_factory_frozen +from torchrl.collectors._evaluator import ( + _extract_metrics_from_trajectories, + _freeze_vecnorm, + _wrap_env_factory_frozen, +) from torchrl.envs import SerialEnv, TransformedEnv from torchrl.envs.env_creator import EnvCreator from torchrl.envs.transforms import Compose, RewardSum, StepCounter, VecNormV2 @@ -718,6 +722,25 @@ def test_fallback_without_reward_sum(self): finally: evaluator.shutdown() + def test_cat_traj_format_metrics(self): + """Evaluator metrics support flat trajectory batches.""" + env = _make_batched_env(num_envs=4, max_steps=5) + policy = _make_policy(env) + evaluator = Evaluator( + env, + policy, + max_steps=20, + num_trajectories=3, + collector_kwargs={"traj_format": "cat"}, + ) + try: + metrics = evaluator.evaluate(step=0) + assert metrics["eval/num_episodes"] == 3 + assert "eval/reward" in metrics + assert "eval/episode_length" in metrics + finally: + evaluator.shutdown() + def test_episode_length_from_step_count(self): """With StepCounter, episode_length should come from step_count at done.""" env = _make_batched_env(num_envs=4, max_steps=5) @@ -730,6 +753,28 @@ def test_episode_length_from_step_count(self): finally: evaluator.shutdown() + def test_missing_traj_info_fallback_invokes_callback(self): + """Without mask/traj_ids the batch is one trajectory and the caller is told.""" + num_steps = 6 + batch = TensorDict( + { + ("next", "reward"): torch.ones(num_steps, 1), + ("next", "done"): torch.zeros(num_steps, 1, dtype=torch.bool), + }, + batch_size=[num_steps], + ) + calls = [] + metrics = _extract_metrics_from_trajectories( + batch, + ("next", "reward"), + ("next", "done"), + None, + on_missing_traj_info=lambda: calls.append(1), + ) + assert len(calls) == 1 + assert metrics["num_episodes"] == 1 + assert metrics["reward"] == pytest.approx(float(num_steps)) + def test_batched_async_metrics(self): """Async eval with batched env produces correct metrics.""" env = _make_batched_env(num_envs=4, max_steps=5) diff --git a/test/data/test_vla.py b/test/data/test_vla.py index aa5200817b4..d60d9bc739e 100644 --- a/test/data/test_vla.py +++ b/test/data/test_vla.py @@ -30,6 +30,7 @@ ) _has_pil = importlib.util.find_spec("PIL") is not None +_has_tensorflow = importlib.util.find_spec("tensorflow") is not None _has_torchvision = importlib.util.find_spec("torchvision") is not None @@ -296,6 +297,17 @@ def test_torchvision_backend_matches_reference_shape(self): assert fast.dtype == fast_again.dtype == ref.dtype == torch.uint8 assert (fast.float() - ref.float()).abs().mean() < 25.0 + @pytest.mark.skipif(not _has_tensorflow, reason="TensorFlow not found") + def test_tensorflow_reference_backend(self): + proc = OpenVLAImagePreprocessor(size=32, backend="tensorflow", center_crop=True) + images = torch.arange(3 * 24 * 40, dtype=torch.int64).reshape(1, 3, 24, 40) + images = images.remainder(256).to(torch.uint8) + first = proc(images) + second = proc(images) + assert first.shape == torch.Size([1, 3, 32, 32]) + assert first.dtype == torch.uint8 + torch.testing.assert_close(first, second) + @pytest.mark.skipif(not _has_pil, reason="Pillow not found") def test_normalization(self): proc = OpenVLAImagePreprocessor( @@ -305,6 +317,19 @@ def test_normalization(self): assert out.shape == torch.Size([3, 16, 16]) assert out.dtype == torch.float32 + @pytest.mark.skipif(not _has_pil, reason="Pillow not found") + def test_fused_backbone_normalization(self): + proc = OpenVLAImagePreprocessor( + size=16, + backend="pil", + mean=[[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]], + std=[[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]], + ) + out = proc(torch.zeros(3, 16, 16, dtype=torch.uint8)) + assert out.shape == torch.Size([6, 16, 16]) + torch.testing.assert_close(out[:3], torch.zeros_like(out[:3])) + torch.testing.assert_close(out[3:], -torch.ones_like(out[3:])) + class TestActionTokenizer: def test_encode_values(self): @@ -421,6 +446,58 @@ def test_norm_stats_unnormalization(self): # gripper dim decodes near the un-normalized [-1, 1] poles assert (recon[:, 2] - actions[:, 2]).abs().max() <= 2.0 / 255 + 1e-5 + def test_cpu_decode_matches_openvla_numpy_reference_exactly(self): + norm_low = [-0.7454732114076613, -0.6616071462631226] + norm_high = [0.9375, 0.8758928775787354] + tok = VocabTailActionTokenizer( + 256, + full_vocab_size=32064, + norm_low=norm_low, + norm_high=norm_high, + norm_mask=torch.ones(2, dtype=torch.bool), + ) + tokens = torch.tensor([[32063, 31936], [31808, 32000]]) + bins = np.linspace(-1.0, 1.0, 256) + centers = (bins[:-1] + bins[1:]) / 2.0 + indices = np.clip(32064 - tokens.numpy() - 1, 0, len(centers) - 1) + normalized = centers[indices] + expected = 0.5 * (normalized + 1.0) * ( + np.asarray(norm_high) - np.asarray(norm_low) + 1e-8 + ) + np.asarray(norm_low) + decoded = tok.decode(tokens) + # the affine runs in NumPy float64 for reference parity, then the + # result is cast back to the tokenizer's float32 working dtype + assert decoded.dtype == torch.float32 + np.testing.assert_array_equal(decoded.numpy(), expected.astype(np.float32)) + + def test_decode_norm_stats_output_dtype_and_device(self): + # the NumPy float64 detour must not leak float64 (or a device change) + # into the returned action chunk + norm_low = torch.tensor([-0.5, 0.0]) + norm_high = torch.tensor([0.5, 2.0]) + tok = VocabTailActionTokenizer(256, norm_low=norm_low, norm_high=norm_high) + tokens = tok.encode(torch.tensor([[0.1, 1.0], [-0.2, 0.5]])) + decoded = tok.decode(tokens) + assert decoded.dtype == torch.float32 + assert decoded.device == tokens.device + + @pytest.mark.gpu + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + def test_decode_norm_stats_cuda(self): + norm_low = torch.tensor([-0.5, 0.0]) + norm_high = torch.tensor([0.5, 2.0]) + tok = VocabTailActionTokenizer( + 256, norm_low=norm_low, norm_high=norm_high + ).cuda() + actions = torch.tensor([[0.1, 1.0], [-0.2, 0.5]], device="cuda") + tokens = tok.encode(actions) + assert tokens.device.type == "cuda" + decoded = tok.decode(tokens) + assert decoded.device.type == "cuda" + assert decoded.dtype == torch.float32 + ref = VocabTailActionTokenizer(256, norm_low=norm_low, norm_high=norm_high) + torch.testing.assert_close(decoded.cpu(), ref.decode(tokens.cpu())) + def test_from_norm_stats(self): norm_stats = { "libero_spatial_no_noops": { @@ -440,9 +517,9 @@ def test_from_norm_stats(self): VocabTailActionTokenizer.from_norm_stats(norm_stats, "bad_key") def test_gripper_binarize_threshold(self): - # SimpleVLA-RL's LIBERO gripper post-processing normalizes a [0, 1] - # gripper to [-1, 1] before taking the sign, which is equivalent to - # thresholding the decoded gripper at 0.5. + # Unmasked OpenVLA gripper dims are already decoded in the normalized + # [-1, 1] convention; SimpleVLA-RL signs that value before inverting + # for LIBERO. norm_low = torch.tensor([0.0]) norm_high = torch.tensor([1.0]) mask = torch.tensor([False]) @@ -455,11 +532,14 @@ def test_gripper_binarize_threshold(self): norm_high=norm_high, norm_mask=mask, gripper_binarize=True, - gripper_binarize_threshold=0.5, + gripper_binarize_threshold=0.0, gripper_invert=True, ) - tokens = base.encode(torch.tensor([[0.25], [0.75]])) - torch.testing.assert_close(tok.decode(tokens), torch.tensor([[1.0], [-1.0]])) + tokens = base.encode(torch.tensor([[-0.25], [0.25]])) + decoded = tok.decode(tokens) + torch.testing.assert_close( + decoded, torch.tensor([[1.0], [-1.0]], dtype=decoded.dtype) + ) def test_validation(self): with pytest.raises(ValueError, match="num_bins"): diff --git a/test/llm/test_llm_objectives.py b/test/llm/test_llm_objectives.py index 125cea91102..da617deff3e 100644 --- a/test/llm/test_llm_objectives.py +++ b/test/llm/test_llm_objectives.py @@ -190,6 +190,25 @@ def test_mc_advantage_dynamic_sampling(self): assert adv_t.completed_trajectories == 0 assert adv_t.completed_decisions == 0 + def test_mc_advantage_clear_queues_preserves_stats(self): + adv_t = MCAdvantage( + grpo_size=2, + prompt_key="group_id", + trajectory_return="sum", + ) + assert adv_t.inv(_make_group_traj(0, [1.0])) is None + assert adv_t.queued_groups == 1 + assert adv_t.queued_trajectories == 1 + assert adv_t.completed_trajectories == 1 + stats = adv_t.get_stats() + assert stats["queued_groups"] == 1 + assert stats["queued_trajectories"] == 1 + assert stats["completed_trajectories"] == 1 + adv_t.clear_queues() + assert adv_t.queued_groups == 0 + assert adv_t.queued_trajectories == 0 + assert adv_t.completed_trajectories == 1 + def test_mc_advantage_candidate_selection_rescues_dynamic_sampling_group(self): selector = MCAdvantageSelector() assert selector.select( @@ -359,6 +378,32 @@ def test_mc_advantage_multi_done_flat_batch(self): assert sum(len(q) for q in adv_t.queues.values()) == 1 assert adv_t.max_queued_trajectories_per_group == 1 + def test_mc_advantage_mixed_lazy_plain_trajectories_return_plain_tensordict(self): + # Collector/replay paths can hand MCAdvantage a mix of lazy/view + # trajectories and plain TensorDicts. MCAdvantage should normalize its + # queued trajectories and return a consistent concrete TensorDict. + adv_t = MCAdvantage(grpo_size=2, prompt_key="group_id", trajectory_return="sum") + assert adv_t.inv(_make_group_traj(0, [1.0])) is None + lazy_traj = lazy_stack([_make_group_traj(1, [0.0])], 0).unbind(0)[0] + assert adv_t.inv(lazy_traj) is None + assert isinstance(adv_t._queue_list(0)[0], TensorDict) + assert isinstance(adv_t._queue_list(1)[0], TensorDict) + + flat = torch.cat( + [ + _make_group_traj(0, [0.0]), + _make_group_traj(1, [1.0]), + ], + 0, + ) + out = adv_t.inv(flat) + + assert isinstance(out, TensorDict) + assert out.shape == (4,) + assert out["advantage"].shape == out["next", "reward"].shape + assert adv_t.queued_groups == 0 + assert adv_t.max_queued_trajectories_per_group == 0 + def test_mc_advantage_share_memory(self): adv_t = MCAdvantage(grpo_size=2, prompt_key="group_id", trajectory_return="sum") assert not adv_t.is_shared @@ -380,6 +425,12 @@ def test_mc_advantage_share_memory(self): adv_t.queues.clear() assert adv_t.queued_groups == 0 + def test_mc_advantage_local_queues(self, monkeypatch): + monkeypatch.setenv("TORCHRL_MC_ADVANTAGE_LOCAL_QUEUES", "1") + adv_t = MCAdvantage(grpo_size=2, prompt_key="group_id", trajectory_return="sum") + assert adv_t.share_memory_() is adv_t + assert not adv_t.is_shared + def test_mc_advantage_share_memory_multiprocessing(self): start_method = "fork" if "fork" in mp.get_all_start_methods() else "spawn" ctx = mp.get_context(start_method) diff --git a/test/modules/test_actor.py b/test/modules/test_actor.py index e46fba8c56c..086bd42393b 100644 --- a/test/modules/test_actor.py +++ b/test/modules/test_actor.py @@ -981,6 +981,17 @@ def test_continuous(self): assert policy.out_keys == [("vla_action", "chunk")] assert ("observation", "image") in policy.in_keys + def test_plain_vla_action_fields(self): + policy = TinyVLA( + action_dim=7, + chunk_size=4, + return_vla_action_container=False, + ) + out = policy(_make_obs_td()) + assert not isinstance(out["vla_action"], VLAAction) + assert out["vla_action", "chunk"].shape == torch.Size([2, 4, 7]) + assert policy.out_keys == [("vla_action", "chunk")] + def test_multistep_actor_uses_vla_chunk_as_cache(self): policy = TinyVLA(action_dim=3, chunk_size=2) actor = MultiStepActorWrapper(policy, n_steps=2) diff --git a/test/rb/test_storages.py b/test/rb/test_storages.py index f77dec13030..c0816299bf9 100644 --- a/test/rb/test_storages.py +++ b/test/rb/test_storages.py @@ -870,6 +870,18 @@ def test_shared_storage_multiprocess(self, storage_cls, use_tmpdir, tmpdir): assert expected.issubset(values) assert len(storage) >= 8 + def test_shared_init_reconciles_non_cpu_device(self): + """Shared init installs a CPU memmap backing; a non-cpu storage device + must be reconciled (else samplers build indices on the wrong device).""" + storage = LazyTensorStorage(max_size=8, device="meta", shared_init=True) + rb = ReplayBuffer(storage=storage, batch_size=2) + data = TensorDict({"x": torch.arange(4, dtype=torch.float32)}, batch_size=(4,)) + with pytest.warns(UserWarning, match="cannot be honored"): + rb.extend(data) + assert torch.device(storage.device).type == "cpu" + sample = rb.sample(2) + assert sample["x"].device.type == "cpu" + def prioritized_collector_worker(self, rb, worker_id, queue): data = TensorDict( { diff --git a/test/test_collectors.py b/test/test_collectors.py index 8f24f4f22b0..03d2ae57f24 100644 --- a/test/test_collectors.py +++ b/test/test_collectors.py @@ -6113,6 +6113,13 @@ def test_start_multi(self, total_frames, cls): break else: raise RuntimeError("RB is empty") + with collector.pause(): + worker_lengths = collector.map_fn("replay_buffer.__len__") + worker_write_counts = collector.get_distant_attr( + "replay_buffer.write_count" + ) + assert worker_lengths == [len(rb)] * 2 + assert worker_write_counts == [rb.write_count] * 2 finally: collector.async_shutdown() del collector @@ -7018,6 +7025,51 @@ class TestTrajsPerBatchReplayBuffer: * **SliceSampler integration**: sampled slices respect episode boundaries. """ + def test_replay_buffer_tensordict_device_metadata_normalization(self): + rb = ReplayBuffer(storage=LazyTensorStorage(16, device="cpu")) + data0 = TensorDict({"obs": torch.zeros(4, 3)}, [4]) + data1 = TensorDict({"obs": torch.ones(4, 3)}, [4]) + assert data0.device is None + assert data1.device is None + + rb.extend(_maybe_normalize_replay_buffer_tensordict_device(data0, rb)) + rb.extend(_maybe_normalize_replay_buffer_tensordict_device(data1, rb)) + + stored = rb.storage.get(slice(0, len(rb))) + assert stored.device == torch.device("cpu") + assert len(rb) == 8 + + def test_replay_buffer_tensordict_device_metadata_uses_storage_payload(self): + rb = ReplayBuffer(storage=LazyTensorStorage(16, device="cpu")) + rb.extend(TensorDict({"obs": torch.zeros(4, 3)}, [4])) + rb.storage.device = None + + data = TensorDict({"obs": torch.ones(4, 3)}, [4]) + data = _maybe_normalize_replay_buffer_tensordict_device(data, rb) + + assert data.device == torch.device("cpu") + rb.extend(data) + assert len(rb) == 8 + + def test_replay_buffer_tensordict_device_metadata_normalizes_nested(self): + rb = ReplayBuffer(storage=LazyTensorStorage(16, device="cpu")) + data = TensorDict( + { + "obs": torch.zeros(4, 3), + "next": TensorDict({"obs": torch.ones(4, 3)}, [4]), + }, + [4], + device="cpu", + ) + data["next"].clear_device_() + + data = _maybe_normalize_replay_buffer_tensordict_device(data, rb) + + assert data.device == torch.device("cpu") + assert data["next"].device == torch.device("cpu") + rb.extend(data) + assert len(rb) == 4 + @staticmethod def _make_env_and_policy(max_steps=4): env_fn = lambda: TransformedEnv( # noqa: E731 @@ -7707,6 +7759,29 @@ def test_trajs_per_batch_multi_async_collector_rb(self): assert ("collector", "traj_ids") in rb.sample(2).keys(True) self._assert_rb_trajectories_complete(rb) + def test_trajs_per_batch_multi_async_collector_empty_rb_ack(self): + """MultiAsyncCollector handles replay-buffer ACKs before any complete traj.""" + max_steps = 10 + env_fn, policy = self._make_env_and_policy(max_steps) + rb = ReplayBuffer(storage=LazyTensorStorage(200), shared=True) + collector = MultiAsyncCollector( + [env_fn], + policy, + replay_buffer=rb, + frames_per_batch=1, + total_frames=1, + trajs_per_batch=1, + cat_results="stack", + ) + try: + collector_iter = iter(collector) + out = next(collector_iter) + finally: + collector.shutdown() + + assert out is None + assert len(rb) == 0 + def test_multi_async_probabilistic_actor_writes_log_prob_to_rb(self): max_steps = 4 num_trajs = 2 diff --git a/test/test_custom_envs.py b/test/test_custom_envs.py index 8558a6f6c1f..7de89c33cd0 100644 --- a/test/test_custom_envs.py +++ b/test/test_custom_envs.py @@ -397,6 +397,7 @@ def test_tracking_grouped_inits(self, batch_size): state_dim=4, success_steps=2, group_repeats=3, + group_id_offset=10, batch_size=batch_size, seed=0, ) @@ -405,7 +406,7 @@ def test_tracking_grouped_inits(self, batch_size): td = env.reset() group_ids.append(td["group_id"].reshape(()).item()) targets.append(td["observation", "state"][..., 2:4].clone()) - assert group_ids == [0, 0, 0, 1, 1, 1] + assert group_ids == [10, 10, 10, 11, 11, 11] torch.testing.assert_close(targets[1], targets[0]) torch.testing.assert_close(targets[2], targets[0]) torch.testing.assert_close(targets[4], targets[3]) @@ -413,7 +414,7 @@ def test_tracking_grouped_inits(self, batch_size): assert not torch.allclose(targets[3], targets[0]) # the group id rides every step of the episode td["action"] = td["observation", "state"][..., 2:4] - assert env.step(td)["next", "group_id"].reshape(()).item() == 1 + assert env.step(td)["next", "group_id"].reshape(()).item() == 11 def test_tracking_group_repeats_validation(self): with pytest.raises(ValueError, match="success_steps"): diff --git a/test/test_inference_server.py b/test/test_inference_server.py index 17a50dad695..cee16b4aeb3 100644 --- a/test/test_inference_server.py +++ b/test/test_inference_server.py @@ -173,6 +173,30 @@ def test_batch_of_requests(self): assert "action" in r.keys() assert r["action"].shape == (2,) + def test_batch_of_requests_with_mixed_root_device_metadata(self): + transport = _MockTransport() + policy = _make_policy() + cpu_item = TensorDict( + {"observation": torch.randn(4), "next": TensorDict({}, [])}, + [], + device="cpu", + ) + metadata_free_item = TensorDict( + {"observation": torch.randn(4), "next": TensorDict({}, [])}, + [], + ) + with InferenceServer(policy, transport, max_batch_size=2, min_batch_size=2): + futures = [ + transport.submit(cpu_item), + transport.submit(metadata_free_item), + ] + results = [f.result(timeout=5.0) for f in futures] + + assert len(results) == 2 + for result in results: + assert "action" in result.keys() + assert result["action"].shape == (2,) + def test_collate_fn_is_called(self): calls = [] @@ -388,6 +412,29 @@ def test_policy_version_is_returned(self): result = client(TensorDict({"observation": torch.randn(4)})) assert result["meta", "policy_version"].item() == 12 + def test_update_model_increments_policy_version(self): + transport = ThreadingTransport() + policy = _make_policy() + with InferenceServer(policy, transport) as server: + client = transport.client() + observation = torch.randn(4) + before = client(TensorDict({"observation": observation.clone()})) + + def update(model): + with torch.no_grad(): + model.module.weight.zero_() + model.module.bias.fill_(1.0) + + server.update_model(update) + after = client(TensorDict({"observation": observation.clone()})) + stats = server.stats() + + assert before["policy_version"].item() == 0 + assert after["policy_version"].item() == 1 + assert stats["policy_version"] == 1 + assert stats["weight_updates"] == 1 + torch.testing.assert_close(after["action"], torch.ones(2)) + @pytest.mark.gpu @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") def test_cuda_policy_cpu_output(self): @@ -775,6 +822,22 @@ def test_single_request_in_process(self): assert "action" in result.keys() assert result["action"].shape == (2,) + @pytest.mark.slow + def test_single_request_with_manager_queues(self): + """Manager-backed MPTransport works from the parent process.""" + ctx = mp.get_context("spawn") + transport = MPTransport(ctx=ctx, use_manager=True) + client = transport.client() + policy = _make_policy() + try: + with InferenceServer(policy, transport, max_batch_size=4): + td = TensorDict({"observation": torch.randn(4)}) + result = client(td) + assert "action" in result.keys() + assert result["action"].shape == (2,) + finally: + transport.close() + @pytest.mark.slow def test_cross_process_actors(self): """Actors in separate processes get correct results.""" diff --git a/test/transforms/test_action_transforms.py b/test/transforms/test_action_transforms.py index 7a45a7fe731..1e27d53cd98 100644 --- a/test/transforms/test_action_transforms.py +++ b/test/transforms/test_action_transforms.py @@ -25,7 +25,7 @@ ) from packaging import version from tensordict import NonTensorData, TensorDict, TensorDictBase -from tensordict.nn import TensorDictModuleBase, WrapModule +from tensordict.nn import TensorDictModuleBase, TensorDictSequential, WrapModule from torch import nn from torchrl.data import ( @@ -3036,6 +3036,34 @@ def test_transform_model(self): ) assert td[ACTION_TOKENS_KEY].tolist() == [[0, 128, 255]] + def test_transform_decode_model(self): + tok = UniformActionTokenizer(256, low=-1.0, high=1.0) + t = ActionTokenizerTransform(tok, mode="decode") + action = torch.tensor([[-0.5, 0.25, 0.5]]) + tokens = tok.encode(action) + td = t(TensorDict({ACTION_TOKENS_KEY: tokens.clone()}, batch_size=[1])) + torch.testing.assert_close(td["action"], tok.decode(tokens)) + assert t.in_keys == [ACTION_TOKENS_KEY] + assert t.out_keys == ["action"] + inv = t.inv(TensorDict({"action": action.clone()}, batch_size=[1])) + torch.testing.assert_close(inv[ACTION_TOKENS_KEY], tok.encode(action)) + + def test_transform_decode_tensordict_sequential(self): + tok = UniformActionTokenizer(256, low=-1.0, high=1.0) + action = torch.tensor([[-0.5, 0.25, 0.5]]) + tokens = tok.encode(action) + policy = TensorDictSequential(ActionTokenizerTransform(tok, mode="decode")) + td = policy(TensorDict({ACTION_TOKENS_KEY: tokens.clone()}, batch_size=[1])) + torch.testing.assert_close(td["action"], tok.decode(tokens)) + + def test_decode_mode_does_not_rewrite_env_action_spec(self): + tok = UniformActionTokenizer(16, low=-1.0, high=1.0) + env = TransformedEnv( + ContinuousActionVecMockEnv(), ActionTokenizerTransform(tok, mode="decode") + ) + assert "action" in env.full_action_spec.keys(True, True) + assert ACTION_TOKENS_KEY not in env.full_action_spec.keys(True, True) + def test_transform_rb(self): tok = UniformActionTokenizer(16, low=-1.0, high=1.0) rb = TensorDictReplayBuffer( @@ -3128,6 +3156,11 @@ def test_requires_tokenizer(self): with pytest.raises(TypeError, match="ActionTokenizerBase"): ActionTokenizerTransform(object()) + def test_invalid_mode(self): + tok = UniformActionTokenizer(256, low=-1.0, high=1.0) + with pytest.raises(ValueError, match="mode"): + ActionTokenizerTransform(tok, mode="invalid") + if __name__ == "__main__": args, unknown = argparse.ArgumentParser().parse_known_args() diff --git a/torchrl/collectors/_evaluator.py b/torchrl/collectors/_evaluator.py index 80efb08ed64..f5c774d5968 100644 --- a/torchrl/collectors/_evaluator.py +++ b/torchrl/collectors/_evaluator.py @@ -85,6 +85,7 @@ from tensordict import TensorDict, TensorDictBase from tensordict.nn import TensorDictModuleBase +from torchrl._utils import logger as torchrl_logger from torchrl.envs import EnvBase from torchrl.envs.utils import ExplorationType, set_exploration_type from torchrl.weight_update.weight_sync_schemes import WeightStrategy @@ -703,28 +704,45 @@ def _extract_metrics_from_trajectories( done_keys: NestedKey, metrics_fn: Callable[[TensorDictBase], dict[str, float]] | None, eval_time: float | None = None, + on_missing_traj_info: Callable[[], None] | None = None, ) -> dict[str, Any]: """Extract evaluation metrics from a trajectory batch produced by a collector. - *traj_batch* has shape ``(num_trajectories, max_traj_len)`` with a - ``("collector", "mask")`` boolean field marking valid timesteps. + *traj_batch* can be padded with a ``("collector", "mask")`` boolean field + or flat (``traj_format="cat"``) with ``("collector", "traj_ids")``. + When neither field is present the whole batch is treated as a single + trajectory and *on_missing_traj_info* (if provided) is invoked so the + caller can warn about potentially wrong episode metrics. """ - mask = traj_batch.get(("collector", "mask")) # [N, T] - num_trajectories = traj_batch.shape[0] - episode_rewards = [] episode_lengths = [] total_frames = 0 - ep_reward_td = traj_batch.get(_EPISODE_REWARD_KEY, None) - step_count_td = traj_batch.get(("next", "step_count"), None) - reward_td = traj_batch.get(reward_keys, None) - - for i in range(num_trajectories): - traj_mask = mask[i] # [T] - if traj_mask.ndim > 1: - traj_mask = traj_mask.squeeze(-1) - valid_len = traj_mask.sum().item() + mask = traj_batch.get(("collector", "mask"), None) + if mask is not None: + trajectories = [] + for i in range(traj_batch.shape[0]): + traj_mask = mask[i] + if traj_mask.ndim > 1: + traj_mask = traj_mask.squeeze(-1) + valid_len = int(traj_mask.sum().item()) + if valid_len: + trajectories.append(traj_batch[i, :valid_len]) + else: + traj_ids = traj_batch.get(("collector", "traj_ids"), None) + if traj_ids is None: + if on_missing_traj_info is not None: + on_missing_traj_info() + trajectories = [traj_batch] + else: + traj_ids = traj_ids.reshape(-1) + trajectories = [ + traj_batch[traj_ids == traj_id] + for traj_id in traj_ids.unique(sorted=True) + ] + + for trajectory in trajectories: + valid_len = trajectory.shape[0] if valid_len == 0: continue total_frames += int(valid_len) @@ -732,21 +750,25 @@ def _extract_metrics_from_trajectories( # Last valid index last_idx = int(valid_len) - 1 + ep_reward_td = trajectory.get(_EPISODE_REWARD_KEY, None) + step_count_td = trajectory.get(("next", "step_count"), None) + reward_td = trajectory.get(reward_keys, None) + if ep_reward_td is not None: # Prefer episode_reward from RewardSum (cumulative return) - r = ep_reward_td[i, last_idx] + r = ep_reward_td[last_idx] if r.ndim > 0: r = r.squeeze(-1) episode_rewards.append(r.item()) elif reward_td is not None: # Fallback: sum raw rewards over valid trajectory steps - valid_rewards = reward_td[i, : int(valid_len)] + valid_rewards = reward_td if valid_rewards.ndim > 1: valid_rewards = valid_rewards.squeeze(-1) episode_rewards.append(valid_rewards.sum().item()) if step_count_td is not None: - ep_len = step_count_td[i, last_idx] + ep_len = step_count_td[last_idx] if ep_len.ndim > 0: ep_len = ep_len.squeeze(-1) episode_lengths.append(ep_len.item()) @@ -980,6 +1002,7 @@ def __init__( # Collector (created lazily) self._collector = None self._collector_iter = None # persistent iterator for multi-collector + self._warned_missing_traj_info = False # Threading state self._lock = threading.Lock() @@ -1206,6 +1229,19 @@ def _run_eval( self._done_keys, self._metrics_fn, eval_time=time.perf_counter() - eval_start, + on_missing_traj_info=self._warn_missing_traj_info, + ) + + def _warn_missing_traj_info(self) -> None: + """Warn (once per instance) that trajectory boundaries are unknown.""" + if self._warned_missing_traj_info: + return + self._warned_missing_traj_info = True + torchrl_logger.warning( + "Evaluator: the collected batch carries neither ('collector', 'mask') " + "nor ('collector', 'traj_ids'); treating the whole batch as a single " + "trajectory. Episode metrics (reward, episode_length, num_episodes) " + "may be wrong if the batch actually contains several trajectories." ) def dump_video(self, step: int | None = None) -> None: diff --git a/torchrl/collectors/_multi_async.py b/torchrl/collectors/_multi_async.py index db77c4367aa..0a2e1068304 100644 --- a/torchrl/collectors/_multi_async.py +++ b/torchrl/collectors/_multi_async.py @@ -212,7 +212,7 @@ def _get_from_queue(self, timeout=None) -> tuple[int, int, TensorDictBase]: else: idx = new_data out = self.out_tensordicts[idx] - if not self.replay_buffer and (j == 0 or use_buffers): + if self.replay_buffer is None and (j == 0 or use_buffers): # we clone the data to make sure that we'll be working with a fixed copy out = out.clone() return idx, j, out diff --git a/torchrl/collectors/_multi_base.py b/torchrl/collectors/_multi_base.py index 04be8456eaa..f28345368ac 100644 --- a/torchrl/collectors/_multi_base.py +++ b/torchrl/collectors/_multi_base.py @@ -1071,6 +1071,18 @@ def _enable_replay_buffer_worker_init(storage): ) if getattr(storage, "shared_init", False): return + storage_device = getattr(storage, "device", None) + if ( + storage_device is not None + and storage_device != "auto" + and torch.device(storage_device).type != "cpu" + ): + warnings.warn( + f"Worker-initialized replay buffers store data in a CPU " + f"memory-mapped tensordict; the storage device " + f"({storage_device}) cannot be honored and will be reset to " + f"'cpu' at initialization time." + ) storage.shared_init = True storage._init_lock = mp.Lock() storage._init_event = mp.Event() diff --git a/torchrl/collectors/_runner.py b/torchrl/collectors/_runner.py index a611b63c1ce..91ac5122151 100644 --- a/torchrl/collectors/_runner.py +++ b/torchrl/collectors/_runner.py @@ -247,11 +247,33 @@ def _main_async_collector( run_free = False if msg == "pause": queue_out.put((idx, "paused"), timeout=_TIMEOUT) - while not pipe_child.poll(1e-2): - continue - data_in, msg = pipe_child.recv() - if msg != "restart": - raise RuntimeError(f"Expected msg='restart', got {msg=}") + while True: + while not pipe_child.poll(1e-2): + continue + data_in, msg = pipe_child.recv() + if msg == "restart": + break + if msg == "cascade_execute": + attr_path, args, kwargs = data_in + try: + result = inner_collector.cascade_execute( + attr_path, *args, **kwargs + ) + pipe_child.send((result, "cascade_execute")) + except Exception as err: + pipe_child.send((err, "cascade_execute")) + continue + if msg == "get_distant_attr": + try: + result = inner_collector.get_distant_attr(data_in) + pipe_child.send((result, "get_distant_attr")) + except Exception as err: + pipe_child.send((err, "get_distant_attr")) + continue + raise RuntimeError( + "Expected a paused-worker control message or " + f"msg='restart', got {msg=}." + ) msg = "continue" else: data_in = None diff --git a/torchrl/collectors/utils.py b/torchrl/collectors/utils.py index 6e1ea9887ba..f8b2b8433c9 100644 --- a/torchrl/collectors/utils.py +++ b/torchrl/collectors/utils.py @@ -458,7 +458,7 @@ def _map_weight( def _traj_chunk_ends_done(chunk: TensorDictBase) -> bool: """Return ``True`` if the last step of *chunk* carries a done/terminated signal.""" - for key in (("next", "done"), ("next", "terminated")): + for key in (("next", "done"), ("next", "terminated"), ("next", "truncated")): signal = chunk.get(key, None) if signal is not None and signal[-1].any().item(): return True diff --git a/torchrl/data/replay_buffers/storages.py b/torchrl/data/replay_buffers/storages.py index 557b24453f3..8110db7e597 100644 --- a/torchrl/data/replay_buffers/storages.py +++ b/torchrl/data/replay_buffers/storages.py @@ -1540,6 +1540,7 @@ def _init_coordinator( ) temp_memmap_storage._init_standard(data) self._storage = temp_memmap_storage._storage + self._reconcile_shared_init_device() return def _wait_for_init(self) -> None: @@ -1547,9 +1548,23 @@ def _wait_for_init(self) -> None: self._init_event.wait() storage = TensorDict.load_memmap(self._init_directory) self._storage = storage + self._reconcile_shared_init_device() self.initialized = True return + def _reconcile_shared_init_device(self) -> None: + # Shared init swaps the backing for a CPU memory-mapped tensordict; + # a stale non-cpu self.device would make samplers build indices on + # the wrong device (RuntimeError at the first sample). + device = self.device + if device not in (None, "auto") and torch.device(device).type != "cpu": + warnings.warn( + f"LazyTensorStorage(shared_init=True) stores data in a CPU " + f"memory-mapped tensordict; the requested storage device " + f"({device}) cannot be honored and is reset to 'cpu'." + ) + self.device = torch.device("cpu") + # Read blocks def get(self, indices: slice) -> TensorDictBase | torch.Tensor | Any: if not self.initialized and self.shared_init: diff --git a/torchrl/data/vla/preprocessing.py b/torchrl/data/vla/preprocessing.py index 113d7eaf652..6570c490e88 100644 --- a/torchrl/data/vla/preprocessing.py +++ b/torchrl/data/vla/preprocessing.py @@ -7,14 +7,17 @@ import importlib.util import io +from collections.abc import Sequence from typing import Literal +import numpy as np import torch import torch.nn.functional as F from torchrl._utils import implement_for _has_pil = importlib.util.find_spec("PIL") is not None +_has_tensorflow = importlib.util.find_spec("tensorflow") is not None _has_torchvision = importlib.util.find_spec("torchvision") is not None __all__ = ["OpenVLAImagePreprocessor"] @@ -54,20 +57,24 @@ def _torchvision_jpeg_roundtrip( # noqa: F811 class OpenVLAImagePreprocessor: """OpenVLA-style image resize, JPEG round-trip and optional center crop. - The operation order mirrors the OpenVLA-OFT evaluation path: resize to a - square image, JPEG encode/decode at the requested quality, optionally apply - a 0.9-area center crop, and resize back. The default ``"torchvision"`` - backend keeps data as tensors and uses ``torchvision.io`` JPEG codecs; - ``"pil"`` is the reference/debugging backend. + The ``"tensorflow"`` backend mirrors the OpenVLA-OFT evaluation path: + JPEG encode/decode at the requested quality, resize with Lanczos3, + optionally apply a 0.9-area center crop, and resize back. The default + ``"torchvision"`` backend keeps data as tensors and uses + ``torchvision.io`` JPEG codecs; ``"pil"`` is a lightweight debugging + backend. Args: size (int): Square output size. Defaults to ``224``. jpeg_quality (int): JPEG quality. Defaults to ``95``. center_crop (bool): Whether to apply the OpenVLA 0.9-area center crop. Defaults to ``False``. - backend (str): ``"torchvision"`` or ``"pil"``. Defaults to - ``"torchvision"``. + backend (str): ``"torchvision"``, ``"pil"`` or ``"tensorflow"``. + Defaults to ``"torchvision"``. mean (torch.Tensor | sequence, optional): Per-channel normalization mean. + A two-dimensional sequence applies multiple normalizations to the + same image and concatenates the results along the channel axis, + as required by fused OpenVLA vision backbones. std (torch.Tensor | sequence, optional): Per-channel normalization std. .. note:: @@ -90,9 +97,9 @@ def __init__( size: int = 224, jpeg_quality: int = 95, center_crop: bool = False, - backend: Literal["torchvision", "pil"] = "torchvision", - mean: torch.Tensor | list[float] | tuple[float, ...] | None = None, - std: torch.Tensor | list[float] | tuple[float, ...] | None = None, + backend: Literal["torchvision", "pil", "tensorflow"] = "torchvision", + mean: torch.Tensor | Sequence[float] | Sequence[Sequence[float]] | None = None, + std: torch.Tensor | Sequence[float] | Sequence[Sequence[float]] | None = None, ) -> None: if size < 1: raise ValueError(f"size must be >= 1, got {size}.") @@ -100,9 +107,10 @@ def __init__( raise ValueError( f"jpeg_quality must be between 1 and 100, got {jpeg_quality}." ) - if backend not in ("torchvision", "pil"): + if backend not in ("torchvision", "pil", "tensorflow"): raise ValueError( - f"backend must be 'torchvision' or 'pil', got {backend!r}." + "backend must be 'torchvision', 'pil' or 'tensorflow', " + f"got {backend!r}." ) if backend == "torchvision" and not _has_torchvision: raise ImportError( @@ -111,6 +119,10 @@ def __init__( ) if backend == "pil" and not _has_pil: raise ImportError("OpenVLAImagePreprocessor backend='pil' requires Pillow.") + if backend == "tensorflow" and not _has_tensorflow: + raise ImportError( + "OpenVLAImagePreprocessor backend='tensorflow' requires TensorFlow." + ) if (mean is None) != (std is None): raise ValueError("mean and std must be provided together.") self.size = int(size) @@ -119,6 +131,17 @@ def __init__( self.backend = backend self.mean = None if mean is None else torch.as_tensor(mean, dtype=torch.float32) self.std = None if std is None else torch.as_tensor(std, dtype=torch.float32) + if self.mean is not None: + if self.mean.shape != self.std.shape: + raise ValueError( + "mean and std must have matching shapes, got " + f"{tuple(self.mean.shape)} and {tuple(self.std.shape)}." + ) + if self.mean.ndim not in (1, 2) or self.mean.shape[-1] != 3: + raise ValueError( + "mean and std must have shape [3] or [N, 3], got " + f"{tuple(self.mean.shape)}." + ) def __call__(self, images: torch.Tensor) -> torch.Tensor: """Preprocess one image ``[C, H, W]`` or a batch ``[..., C, H, W]``.""" @@ -136,19 +159,27 @@ def __call__(self, images: torch.Tensor) -> torch.Tensor: flat = self._to_uint8(flat) if self.backend == "torchvision": processed = self._torchvision(flat) + elif self.backend == "tensorflow": + processed = self._tensorflow(flat) else: processed = self._pil(flat) processed = processed.reshape(*original_shape, *processed.shape[-3:]) if self.mean is not None: processed = processed.to(torch.float32).div_(255.0) - view_shape = *((1,) * (processed.ndim - 3)), -1, 1, 1 - mean = self.mean.to(device=processed.device, dtype=processed.dtype).view( - view_shape + means = self.mean.reshape(-1, self.mean.shape[-1]).to( + device=processed.device, dtype=processed.dtype ) - std = self.std.to(device=processed.device, dtype=processed.dtype).view( - view_shape + stds = self.std.reshape(-1, self.std.shape[-1]).to( + device=processed.device, dtype=processed.dtype + ) + view_shape = *((1,) * (processed.ndim - 3)), -1, 1, 1 + processed = torch.cat( + [ + processed.sub(mean.view(view_shape)).div(std.view(view_shape)) + for mean, std in zip(means, stds, strict=True) + ], + dim=-3, ) - processed = processed.sub_(mean).div_(std) return processed @staticmethod @@ -195,8 +226,57 @@ def _torchvision(self, images: torch.Tensor) -> torch.Tensor: decoded = self._resize(decoded) return decoded + def _tensorflow(self, images: torch.Tensor) -> torch.Tensor: + import tensorflow as tf + + out = [] + resize_size = (self.size, self.size) + for image in images.cpu(): + array = image.permute(1, 2, 0).numpy().astype("uint8") + if array.shape[-1] == 1: + array = np.repeat(array, 3, axis=-1) + tf_image = tf.convert_to_tensor(array) + tf_image = tf.image.encode_jpeg(tf_image, quality=self.jpeg_quality) + tf_image = tf.io.decode_image( + tf_image, expand_animations=False, dtype=tf.uint8 + ) + tf_image = tf.image.resize( + tf_image, resize_size, method="lanczos3", antialias=True + ) + tf_image = tf.cast(tf.clip_by_value(tf.round(tf_image), 0, 255), tf.uint8) + if self.center_crop: + expanded = tf.expand_dims( + tf.image.convert_image_dtype(tf_image, tf.float32), axis=0 + ) + crop_size = tf.reshape( + tf.clip_by_value(tf.sqrt(tf.constant(_CROP_AREA_SCALE)), 0, 1), + shape=(1,), + ) + offsets = (1 - crop_size) / 2 + boxes = tf.stack( + [ + offsets, + offsets, + offsets + crop_size, + offsets + crop_size, + ], + axis=1, + ) + tf_image = tf.image.crop_and_resize( + expanded, boxes, tf.range(1), resize_size + )[0] + tf_image = tf.clip_by_value(tf_image, 0, 1) + tf_image = tf.image.convert_image_dtype( + tf_image, tf.uint8, saturate=True + ) + out.append( + torch.from_numpy( + np.asarray(tf_image.numpy(), dtype="uint8").copy() + ).permute(2, 0, 1) + ) + return torch.stack(out, 0).to(images.device) + def _pil(self, images: torch.Tensor) -> torch.Tensor: - import numpy as np from PIL import Image out = [] diff --git a/torchrl/data/vla/tokenizers.py b/torchrl/data/vla/tokenizers.py index 2674f6b81c5..05e07bbea89 100644 --- a/torchrl/data/vla/tokenizers.py +++ b/torchrl/data/vla/tokenizers.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING +import numpy as np import torch from torch import nn @@ -74,7 +75,8 @@ class UniformActionTokenizer(ActionTokenizerBase): >>> tokens = tok.encode(torch.tensor([-1.0, 0.0, 1.0])) >>> tokens tensor([ 0, 128, 255]) - >>> torch.allclose(tok.decode(tokens), torch.tensor([-0.998, 0.002, 0.998]), atol=1e-2) + >>> expected = torch.tensor([-0.998, 0.002, 0.998]) + >>> torch.allclose(tok.decode(tokens), expected, atol=1e-2) True >>> tok.vocab_size 256 @@ -139,7 +141,7 @@ def decode(self, tokens: torch.Tensor) -> torch.Tensor: def from_metadata( cls, metadata: RobotDatasetMetadata, num_bins: int ) -> UniformActionTokenizer: - """Build from the ``action_low``/``action_high`` of a :class:`~torchrl.data.vla.RobotDatasetMetadata`.""" + """Build from the action bounds of a dataset metadata object.""" if metadata.action_low is None or metadata.action_high is None: raise ValueError( f"metadata {metadata.dataset_id!r} has no action bounds " @@ -173,9 +175,12 @@ class VocabTailActionTokenizer(ActionTokenizerBase): Optionally, dataset statistics (the ``norm_stats`` shipped with OpenVLA checkpoints) un-normalize decoded actions to the environment's action space -- and normalize actions before encoding -- via the affine q01/q99 - map ``a_env = 0.5 * (a + 1) * (q99 - q01) + q01`` applied to the + map ``a_env = 0.5 * (a + 1) * (q99 - q01 + 1e-8) + q01`` applied to the dimensions selected by ``mask`` (the gripper dimension is typically - excluded). See :meth:`from_norm_stats`. + excluded). To reproduce the reference MuJoCo action path exactly, decoding + with normalization statistics is performed in NumPy float64 on CPU; the + result is cast back to ``float32`` on the input tokens' device. See + :meth:`from_norm_stats` and :meth:`decode`. Args: num_bins (int): number of bin edges per action dimension (the OpenVLA @@ -258,8 +263,11 @@ def __init__( self.register_buffer("bins", bins) self.register_buffer("bin_centers", (bins[:-1] + bins[1:]) / 2.0) if norm_low is not None: - norm_low = torch.as_tensor(norm_low, dtype=torch.float32) - norm_high = torch.as_tensor(norm_high, dtype=torch.float32) + # The reference action path performs this affine in NumPy float64. + # Preserve the checkpoint's JSON precision for exact CPU decode; + # device-side encode/decode casts these buffers to the action dtype. + norm_low = torch.as_tensor(norm_low, dtype=torch.float64) + norm_high = torch.as_tensor(norm_high, dtype=torch.float64) if norm_mask is None: norm_mask = torch.ones_like(norm_low, dtype=torch.bool) else: @@ -286,10 +294,12 @@ def _digitize(self, actions: torch.Tensor) -> torch.Tensor: def encode(self, actions: torch.Tensor) -> torch.Tensor: if self.norm_low is not None: - norm_low = self.norm_low.to(actions.device) - norm_high = self.norm_high.to(actions.device) + norm_low = self.norm_low.to(device=actions.device, dtype=actions.dtype) + norm_high = self.norm_high.to(device=actions.device, dtype=actions.dtype) norm_mask = self.norm_mask.to(actions.device) - scale = (norm_high - norm_low).clamp_min(1e-8) + # same affine convention as decode / the OpenVLA reference: + # a = 2 * (a_env - q01) / (q99 - q01 + 1e-8) - 1 + scale = norm_high - norm_low + 1e-8 normalized = 2.0 * (actions - norm_low) / scale - 1.0 actions = torch.where(norm_mask, normalized, actions) digitized = self._digitize(actions) @@ -298,38 +308,64 @@ def encode(self, actions: torch.Tensor) -> torch.Tensor: return self.num_bins - digitized def decode(self, tokens: torch.Tensor) -> torch.Tensor: + """Map token ids back to continuous actions ``[..., action_dim]``. + + When normalization statistics are set (see :meth:`from_norm_stats`), + the de-tokenization and the q01/q99 un-normalization are computed in + NumPy float64 on CPU for bit-exact parity with the OpenVLA-OFT + reference implementation. This incurs a device-to-host round-trip + (and a host sync) on every call when ``tokens`` lives on an + accelerator; the result is moved back to ``tokens.device`` and cast + to the tokenizer's working dtype (``float32``). + + Without statistics, decoding runs entirely on ``tokens.device`` and + returns bin centers in ``[-1, 1]``. + """ + if self.norm_low is not None: + # LIBERO executes CPU actions. Reproduce SimpleVLA's NumPy float64 + # de-tokenization before handing those values to MuJoCo; float32 + # stats or a missing range epsilon make fragile closed-loop + # rollouts diverge after several chunks. + tokens_np = tokens.detach().cpu().numpy() + if self.full_vocab_size is not None: + digitized = self.full_vocab_size - tokens_np + else: + digitized = self.num_bins - tokens_np + bins = np.linspace(-1.0, 1.0, self.num_bins) + bin_centers = (bins[:-1] + bins[1:]) / 2.0 + index = np.clip(digitized - 1, a_min=0, a_max=bin_centers.shape[0] - 1) + actions = bin_centers[index] + norm_low = self.norm_low.detach().cpu().numpy() + norm_high = self.norm_high.detach().cpu().numpy() + norm_mask = self.norm_mask.detach().cpu().numpy().astype(bool) + unnormalized = ( + 0.5 * (actions + 1.0) * (norm_high - norm_low + 1e-8) + norm_low + ) + actions = np.where(norm_mask, unnormalized, actions) + if self.gripper_binarize or self.gripper_invert: + gripper = ~norm_mask + if self.gripper_binarize: + binary = (actions > self.gripper_binarize_threshold).astype( + actions.dtype + ) * 2.0 - 1.0 + actions = np.where(gripper, binary, actions) + if self.gripper_invert: + actions = np.where(gripper, -actions, actions) + # ride the input device and the tensor-path output dtype so the + # float64 CPU detour does not leak into downstream action chunks + return torch.from_numpy(np.ascontiguousarray(actions)).to( + device=tokens.device, dtype=self.bin_centers.dtype + ) + # operate on the tokens' device (policy and env may differ): move the - # small lookup/normalization buffers there so the output rides the - # input device + # small lookup buffer there so the output rides the input device bin_centers = self.bin_centers.to(tokens.device) if self.full_vocab_size is not None: digitized = self.full_vocab_size - tokens else: digitized = self.num_bins - tokens index = (digitized - 1).clamp(0, bin_centers.shape[0] - 1) - actions = bin_centers[index] - if self.norm_low is not None: - norm_low = self.norm_low.to(tokens.device) - norm_high = self.norm_high.to(tokens.device) - norm_mask = self.norm_mask.to(tokens.device) - unnormalized = 0.5 * (actions + 1.0) * (norm_high - norm_low) + norm_low - actions = torch.where(norm_mask, unnormalized, actions) - if ( - self.gripper_binarize or self.gripper_invert - ) and self.norm_mask is not None: - # the gripper dims are the *unmasked* ones (not min-max normalized). - # Robots expect a firm open/close, not the continuous bin value the - # model emits: binarize to +/-1 by sign (optionally invert the - # open/close convention). - gripper = ~self.norm_mask.to(tokens.device) - if self.gripper_binarize: - binval = (actions > self.gripper_binarize_threshold).to( - actions.dtype - ) * 2.0 - 1.0 - actions = torch.where(gripper, binval, actions) - if self.gripper_invert: - actions = torch.where(gripper, -actions, actions) - return actions + return bin_centers[index] @classmethod def from_norm_stats( @@ -371,8 +407,8 @@ def from_norm_stats( return cls( num_bins, full_vocab_size=full_vocab_size, - norm_low=torch.as_tensor(stats["q01"], dtype=torch.float32), - norm_high=torch.as_tensor(stats["q99"], dtype=torch.float32), + norm_low=torch.as_tensor(stats["q01"], dtype=torch.float64), + norm_high=torch.as_tensor(stats["q99"], dtype=torch.float64), norm_mask=torch.as_tensor(mask, dtype=torch.bool) if mask is not None else None, diff --git a/torchrl/envs/custom/vla.py b/torchrl/envs/custom/vla.py index 451a6991615..2b368afb9ff 100644 --- a/torchrl/envs/custom/vla.py +++ b/torchrl/envs/custom/vla.py @@ -95,6 +95,10 @@ class ToyVLAEnv(EnvBase): advantages require (n rollouts per initial state, e.g. grouped by :class:`~torchrl.objectives.llm.MCAdvantage`). Defaults to ``None`` (a fresh target every episode, no ``group_id`` entry). + group_id_offset (int, optional): offset added to grouped rollout ids. + This lets several single ToyVLAEnv instances collect grouped + rollouts in parallel without mixing unrelated initial states in the + downstream group-advantage transform. Defaults to ``0``. batch_size (torch.Size, optional): number of vectorized copies. Defaults to ``torch.Size([])`` (a single environment). device (torch.device, optional): device of the specs. @@ -150,6 +154,7 @@ def __init__( success_steps: int | None = None, success_tol: float = 0.25, group_repeats: int | None = None, + group_id_offset: int = 0, batch_size: torch.Size | None = None, device: torch.device | None = None, seed: int | None = None, @@ -195,6 +200,7 @@ def __init__( self.success_steps = int(success_steps) if success_steps is not None else None self.success_tol = float(success_tol) self.group_repeats = int(group_repeats) if group_repeats is not None else None + self.group_id_offset = int(group_id_offset) if self.group_repeats is not None and self.batch_size.numel() > 1: raise ValueError( "group_repeats only supports a single environment " @@ -368,7 +374,8 @@ def _reset(self, tensordict: TensorDictBase | None = None, **kwargs) -> TensorDi if self._episode_count % self.group_repeats == 0: self._target = self._sample_target() self._group_id = torch.full_like( - self._group_id, self._episode_count // self.group_repeats + self._group_id, + self.group_id_offset + self._episode_count // self.group_repeats, ) self._episode_count += 1 else: diff --git a/torchrl/envs/transforms/_action.py b/torchrl/envs/transforms/_action.py index 2f8a3e25abd..de7300609f1 100644 --- a/torchrl/envs/transforms/_action.py +++ b/torchrl/envs/transforms/_action.py @@ -11,7 +11,7 @@ from copy import copy from enum import IntEnum from textwrap import indent -from typing import Any, TYPE_CHECKING +from typing import Any, Literal, TYPE_CHECKING import torch @@ -1988,24 +1988,32 @@ class ActionTokenizerTransform(Transform): Like any TorchRL transform it plugs onto a replay buffer or an environment interchangeably: - - **forward** (``encode``): maps the continuous action (or action chunk) at - ``in_key`` to discrete token ids at ``out_key`` -- e.g. building the token - training target for an autoregressive (RT-2 / OpenVLA-style) token VLA on - the replay-buffer sample path. - - **inverse** (``decode``): maps token ids at ``out_key`` back to a continuous - action at ``in_key`` -- e.g. decoding the tokens a token-head policy emits, - on the environment action-input path, before the base env consumes them. - On a replay buffer the inverse is a no-op when the token entry is - absent, so extending with raw (untokenized) data is safe; attached to an - environment, missing tokens on the step path raise instead. - - When attached to an environment, the policy-facing action spec is rewritten - to a :class:`~torchrl.data.Categorical` over the tokenizer's vocabulary, so - the env advertises the token interface the policy is expected to produce - (the decoded continuous action is consumed by the base env internally). - Using the same tokenizer instance on the replay buffer (encode) and on the - env (decode) guarantees that training targets and execution share the exact - same binning. + - **forward encode mode** (``mode="encode"``, the default): maps the + continuous action (or action chunk) at ``in_key`` to discrete token ids at + ``out_key`` -- e.g. building the token training target for an + autoregressive (RT-2 / OpenVLA-style) token VLA on the replay-buffer + sample path. + - **inverse encode mode**: maps token ids at ``out_key`` back to a + continuous action at ``in_key`` -- e.g. decoding the tokens a token-head + policy emits, on the environment action-input path, before the base env + consumes them. + - **forward decode mode** (``mode="decode"``): maps token ids at + ``out_key`` to continuous actions at ``in_key``. This is useful on the + policy side, for instance as a module after a token VLA policy in a + :class:`~tensordict.nn.TensorDictSequential`, so CPU environments can + receive decoded actions without owning tokenizer buffers. + + On a replay buffer the inverse is a no-op when the token entry is absent, + so extending with raw (untokenized) data is safe; attached to an + environment, missing tokens on the step path raise instead. + + When attached to an environment in encode mode, the policy-facing action + spec is rewritten to a :class:`~torchrl.data.Categorical` over the + tokenizer's vocabulary, so the env advertises the token interface the + policy is expected to produce (the decoded continuous action is consumed by + the base env internally). Using the same tokenizer instance on the replay + buffer (encode) and on the env (decode through the inverse path) guarantees + that training targets and execution share the exact same binning. Args: tokenizer (ActionTokenizerBase): the tokenizer to apply. @@ -2015,6 +2023,9 @@ class ActionTokenizerTransform(Transform): out_key (NestedKey): the discrete token ids. Defaults to ``("vla_action", "tokens")``. Pass ``"action_tokens"`` for the flat compatibility key. + mode (str, optional): ``"encode"`` makes :meth:`forward` encode + actions into tokens and :meth:`inv` decode tokens into actions. + ``"decode"`` swaps these directions. Defaults to ``"encode"``. Examples: >>> import torch @@ -2030,6 +2041,11 @@ class ActionTokenizerTransform(Transform): >>> back = t.inv(TensorDict({("vla_action", "tokens"): td["vla_action", "tokens"]}, batch_size=[1])) >>> back["action"].shape torch.Size([1, 3]) + >>> # policy-side decode: token policy -> decoded continuous action + >>> decode = ActionTokenizerTransform(tok, mode="decode") + >>> policy_td = TensorDict({("vla_action", "tokens"): td["vla_action", "tokens"]}, batch_size=[1]) + >>> decode(policy_td)["action"].shape + torch.Size([1, 3]) >>> # on a replay buffer: raw actions written through extend are stored >>> # as-is and tokenized on the sample path >>> from torchrl.data import LazyTensorStorage, TensorDictReplayBuffer @@ -2070,46 +2086,82 @@ def __init__( *, in_key: NestedKey = ACTION_KEY, out_key: NestedKey = ACTION_TOKENS_KEY, + mode: Literal["encode", "decode"] = "encode", ) -> None: if not isinstance(tokenizer, ActionTokenizerBase): raise TypeError( f"tokenizer must be an ActionTokenizerBase, got {type(tokenizer)}." ) + if mode not in ("encode", "decode"): + raise ValueError(f"mode must be either 'encode' or 'decode', got {mode!r}.") + action_key = unravel_key(in_key) + token_key = unravel_key(out_key) # ``forward`` is fully overridden (encode in_key -> out_key on the data # path), so no forward keys are declared: the token entry only exists # on the data path, never in the env's output specs. The inverse # direction reads ``out_keys_inv`` (the tokens) and writes # ``in_keys_inv`` (the action passed to the base env). + if mode == "encode": + in_keys = [] + out_keys = [] + in_keys_inv = [action_key] + out_keys_inv = [token_key] + else: + in_keys = [token_key] + out_keys = [action_key] + in_keys_inv = [token_key] + out_keys_inv = [action_key] super().__init__( - in_keys=[], - out_keys=[], - in_keys_inv=[in_key], - out_keys_inv=[out_key], + in_keys=in_keys, + out_keys=out_keys, + in_keys_inv=in_keys_inv, + out_keys_inv=out_keys_inv, ) self.tokenizer = tokenizer + self.mode = mode + self._action_key = action_key + self._token_key = token_key @property def in_key(self) -> NestedKey: - return self.in_keys_inv[0] + return self._action_key @property def out_key(self) -> NestedKey: - return self.out_keys_inv[0] + return self._token_key def forward(self, tensordict: TensorDictBase) -> TensorDictBase: - action = tensordict.get(self.in_key, default=None) - if action is None: + if self.mode == "encode": + source_key = self.in_key + dest_key = self.out_key + transform = self.tokenizer.encode + else: + source_key = self.out_key + dest_key = self.in_key + transform = self.tokenizer.decode + value = tensordict.get(source_key, default=None) + if value is None: if self.missing_tolerance: return tensordict raise KeyError( - f"{type(self).__name__}: '{self.in_key}' not found in tensordict " + f"{type(self).__name__}: '{source_key}' not found in tensordict " f"{tensordict}." ) - tensordict.set(self.out_key, self.tokenizer.encode(action)) + tensordict.set(dest_key, transform(value)) return tensordict def _inv_apply_transform(self, tokens: torch.Tensor) -> torch.Tensor: - return self.tokenizer.decode(tokens) + if self.mode == "encode": + return self.tokenizer.decode(tokens) + return self.tokenizer.encode(tokens) + + def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase: + # This transform acts on data/replay-buffer samples through ``forward`` + # and on env actions through ``inv``. When used as a policy-side decode + # module, ``in_keys``/``out_keys`` are populated for + # TensorDictSequential introspection, but an attached env should not try + # to decode observations on reset/step. + return next_tensordict def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase: # Without a parent env (replay-buffer ``extend``), raw (untokenized) @@ -2118,11 +2170,14 @@ def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase: # them; a policy writing the raw action key by mistake should not # silently bypass the decode). The env reset path never reaches the # inverse: ``enable_inv_on_reset`` defaults to False. - if self.parent is None and tensordict.get(self.out_key, default=None) is None: + expected_key = self.out_key if self.mode == "encode" else self.in_key + if self.parent is None and tensordict.get(expected_key, default=None) is None: return tensordict return super()._inv_call(tensordict) def transform_input_spec(self, input_spec: Composite) -> Composite: + if self.mode == "decode": + return input_spec # Expose the token interface to the policy: replace the base env's # continuous action spec with a Categorical over the tokenizer # vocabulary. The continuous spec is removed rather than moved to the @@ -2158,3 +2213,8 @@ def transform_input_spec(self, input_spec: Composite) -> Composite: if token_key != action_key: del input_spec["full_action_spec", action_key] return input_spec + + def transform_output_spec(self, output_spec: Composite) -> Composite: + if self.mode == "decode": + return output_spec + return super().transform_output_spec(output_spec) diff --git a/torchrl/envs/transforms/_reward.py b/torchrl/envs/transforms/_reward.py index 04989243695..de1244c4c41 100644 --- a/torchrl/envs/transforms/_reward.py +++ b/torchrl/envs/transforms/_reward.py @@ -995,13 +995,12 @@ def _apply_transform(self, reward: Tensor) -> TensorDictBase: class SuccessReward(Transform): - """Sparse 0/1 success reward for reinforcement fine-tuning. + """Sparse reward from a binary success signal. Reads a boolean (or 0/1) success signal and writes a sparse reward - (``scale`` on success, ``0`` otherwise). This is the trajectory-level - success reward used by SimpleVLA-RL / RL4VLA-style VLA RL, where a binary - task-completion signal is the only reward, but it is a general transform: - sparse task-completion rewards are ubiquitous in goal-conditioned RL. + (``scale`` on success, ``0`` otherwise). This transform is useful for + environments that expose task completion separately from reward, such as + goal-conditioned tasks or sparse robotics benchmarks. It is a standard leaf transform: it can be appended to a :class:`~torchrl.envs.TransformedEnv` (it overwrites the step reward from the @@ -1010,6 +1009,10 @@ class SuccessReward(Transform): :class:`~torchrl.data.Bounded` spec over ``{0, scale}`` (shaped like the success entry); the reward is written at step time only, never at reset. + .. note:: + This transform can be used in any setup with a binary completion + signal, including VLA-style robot manipulation recipes. + Args: success_key (NestedKey): the boolean success signal to read. Defaults to ``"success"``. diff --git a/torchrl/modules/inference_server/_mp.py b/torchrl/modules/inference_server/_mp.py index ef512e0b674..351d37b2cff 100644 --- a/torchrl/modules/inference_server/_mp.py +++ b/torchrl/modules/inference_server/_mp.py @@ -23,6 +23,10 @@ class MPTransport(QueueBasedTransport): Args: ctx: a multiprocessing context (e.g. ``mp.get_context("spawn")``). Defaults to ``mp.get_context("spawn")``. + use_manager (bool, optional): if ``True``, back the request and + response queues with a multiprocessing manager. This is useful + when clients are forwarded through another spawned process. + Defaults to ``False``. Example: >>> import multiprocessing as mp @@ -32,19 +36,62 @@ class MPTransport(QueueBasedTransport): >>> p.start() # queue inherited """ - def __init__(self, ctx: mp.context.BaseContext | None = None): + def __init__( + self, ctx: mp.context.BaseContext | None = None, *, use_manager: bool = False + ): super().__init__() self._ctx = ctx if ctx is not None else mp.get_context("spawn") - self._request_queue: mp.Queue = self._ctx.Queue() - self._response_queues: dict[int, mp.Queue] = {} + self._use_manager = bool(use_manager) + self._manager = self._ctx.Manager() if use_manager else None + if self._manager is None: + self._request_queue: mp.Queue = self._ctx.Queue() + self._response_queues: dict[int, mp.Queue] = {} + else: + self._request_queue = self._manager.Queue() + self._response_queues = self._manager.dict() def _make_response_queue(self) -> mp.Queue: + if self._manager is not None: + return self._manager.Queue() + if self._use_manager: + # The manager handle is dropped by __getstate__: an unpickled copy + # cannot create manager-backed response queues, and falling back to + # a raw mp.Queue would fail later (and confusingly) when registered + # in the manager-backed response-queue dict. + raise RuntimeError( + "Cannot create a client from an unpickled manager-backed " + "MPTransport: the manager handle is not serialized. Clients of " + "a manager-backed MPTransport must be created with " + "transport.client() in the owning process before " + "pickling/sending the transport." + ) return self._ctx.Queue() + def __getstate__(self) -> dict: + state = super().__getstate__() + state["_manager"] = None + return state + + def close(self) -> None: + """Release transport resources. + + Shuts down the multiprocessing manager backing the request/response + queues when the transport was built with ``use_manager=True`` + (a no-op otherwise). The process that owns the transport must call + this once the server and all clients are done with it: + :class:`~torchrl.modules.inference_server.ProcessInferenceServer` + does not close the transport on ``shutdown()``. + """ + if self._manager is not None: + self._manager.shutdown() + def client(self) -> _QueueInferenceClient: """Create an actor-side client with a dedicated response queue. Must be called in the parent process **before** spawning children. + In particular, a manager-backed transport (``use_manager=True``) + loses its manager handle when pickled, so calling this on an + unpickled copy raises a :class:`RuntimeError`. Returns: A :class:`_QueueInferenceClient` that can be passed to a child diff --git a/torchrl/modules/inference_server/_queue_transport.py b/torchrl/modules/inference_server/_queue_transport.py index 496736144df..28cbe4eac81 100644 --- a/torchrl/modules/inference_server/_queue_transport.py +++ b/torchrl/modules/inference_server/_queue_transport.py @@ -216,6 +216,8 @@ def drain_with_timing( item = self._request_queue.get(block=False) except Exception: break + if item is None: + continue if len(item) == 4: actor_id, req_id, item_submitted_at, td = item else: @@ -231,9 +233,11 @@ def wait_for_work(self, timeout: float) -> None: if self._peeked is not None: return try: - self._peeked = self._request_queue.get(timeout=timeout) + peeked = self._request_queue.get(timeout=timeout) except Exception: - pass + return + if peeked is not None: + self._peeked = peeked def resolve(self, callback: tuple[int, int], result: TensorDictBase) -> None: """Route the result to the correct actor's response queue.""" diff --git a/torchrl/modules/inference_server/_server.py b/torchrl/modules/inference_server/_server.py index 52ba30604be..5614b67f650 100644 --- a/torchrl/modules/inference_server/_server.py +++ b/torchrl/modules/inference_server/_server.py @@ -14,7 +14,7 @@ from concurrent.futures import Future from multiprocessing.synchronize import Event as MPEvent from statistics import mean -from typing import Literal +from typing import Any, Literal import torch from tensordict import lazy_stack @@ -32,7 +32,11 @@ InferenceServerConfig, ) from torchrl.modules.inference_server._transport import InferenceTransport -from torchrl.weight_update import SharedMemWeightSyncScheme, WeightSyncScheme +from torchrl.weight_update import ( + SharedMemWeightSyncScheme, + WeightStrategy, + WeightSyncScheme, +) _CODE_TO_INTERACTION_TYPE = { 0: InteractionType.MODE, @@ -43,6 +47,48 @@ } +def _normalize_tensordict_device_metadata(data: TensorDictBase) -> TensorDictBase: + target_device = data.device + if target_device is not None: + target_device = torch.device(target_device) + for value in data.values(include_nested=True, leaves_only=True): + value_device = getattr(value, "device", None) + if value_device is None: + continue + value_device = torch.device(value_device) + if target_device is None: + target_device = value_device + elif value_device != target_device: + return data + + if target_device is None: + return data + + needs_normalization = data.device != target_device + if not needs_normalization: + for value in data.values(include_nested=True, leaves_only=False): + if isinstance(value, TensorDictBase) and value.device != target_device: + needs_normalization = True + break + if not needs_normalization: + return data + + data = data.copy() + data.clear_device_() + return data.to(target_device) + + +def _default_collate(items: list[TensorDictBase]) -> TensorDictBase: + return lazy_stack( + [ + _normalize_tensordict_device_metadata(item) + if isinstance(item, TensorDictBase) + else item + for item in items + ] + ) + + class InferenceServer: """Auto-batching inference server. @@ -203,7 +249,7 @@ def __init__( self.max_batch_size = max_batch_size self.min_batch_size = min_batch_size self.timeout = timeout - self.collate_fn = collate_fn if collate_fn is not None else lazy_stack + self.collate_fn = collate_fn if collate_fn is not None else _default_collate policy_device = device if policy_device is None else policy_device self.policy_device = ( torch.device(policy_device) if policy_device is not None else None @@ -334,6 +380,30 @@ def update_policy_weights_(self, model_id=None, policy_or_weights=None, **kwargs """ self._mark_weight_update() + def update_model( + self, + update_fn: Callable[[nn.Module], Any], + *, + mark_weight_update: bool = True, + ) -> Any: + """Apply an in-place update to the served model under the model lock. + + Args: + update_fn (Callable): function called with ``self.model`` while + inference is blocked by the server's model lock. + mark_weight_update (bool, optional): if ``True``, increment the + behavior-policy version and weight-update counter after + ``update_fn`` succeeds. Defaults to ``True``. + + Returns: + The value returned by ``update_fn``. + """ + with self._model_lock: + result = update_fn(self.model) + if mark_weight_update: + self._mark_weight_update() + return result + # -- lifecycle ------------------------------------------------------------ def start(self) -> InferenceServer: @@ -402,10 +472,10 @@ def _poll_weight_update(self) -> None: return with self._model_lock: weights = ws.receive(timeout=0.0) - if weights is not None and getattr(ws, "context", None) is not self: - # When the server is the scheme context, receive() already - # cascaded into update_policy_weights_; do not count twice. - self._mark_weight_update() + if weights is not None and getattr(ws, "context", None) is not self: + # When the server is the scheme context, receive() already + # cascaded into update_policy_weights_; do not count twice. + self._mark_weight_update() def _set_policy_version(self, result_batch: TensorDictBase) -> TensorDictBase: """Annotate inference outputs with the behavior policy version.""" @@ -611,6 +681,22 @@ def _process_server_entry( "alive": server.is_alive, "policy_version": server.policy_version, } + elif verb == "update_model_weights": + weights = payload_in["weights"] + mark_weight_update = payload_in.get("mark_weight_update", True) + with server._model_lock: + if hasattr(server.model, "load_policy_weights"): + server.model.load_policy_weights(weights) + else: + WeightStrategy(extract_as="tensordict").apply_weights( + server.model, + weights.to(server.policy_device) + if server.policy_device is not None + else weights, + ) + if mark_weight_update: + server._mark_weight_update() + payload = {"accepted": True} elif verb == "shutdown": shutdown_event.set() payload = {"accepted": True} @@ -849,7 +935,7 @@ def start(self) -> ProcessInferenceServer: def _request_control( self, - verb: Literal["stats", "health", "shutdown"], + verb: Literal["stats", "health", "update_model_weights", "shutdown"], payload: dict | None = None, timeout: float = 5.0, ): @@ -939,6 +1025,31 @@ def stats( """ return self._request_control("stats", {"reset": reset}, timeout=timeout) + def update_model_weights( + self, + weights: TensorDictBase, + *, + mark_weight_update: bool = True, + timeout: float = 300.0, + ) -> dict[str, bool]: + """Apply TensorDict weights to the model hosted by the child process. + + This is a blocking control-plane round trip; large models can take a + while to transfer and apply, hence the generous default timeout. + + Args: + weights (TensorDictBase): weights to apply to the child's model. + mark_weight_update (bool, optional): whether to bump the child's + behavior-policy version. Defaults to ``True``. + timeout (float, optional): seconds to wait for the child to apply + the weights. Defaults to ``300.0``. + """ + return self._request_control( + "update_model_weights", + {"weights": weights, "mark_weight_update": mark_weight_update}, + timeout=timeout, + ) + def health(self, *, timeout: float = 5.0) -> dict[str, int | bool | None]: """Return a lightweight child-process health snapshot. diff --git a/torchrl/modules/vla/common.py b/torchrl/modules/vla/common.py index fe66c735efc..99f3fa5553a 100644 --- a/torchrl/modules/vla/common.py +++ b/torchrl/modules/vla/common.py @@ -56,6 +56,10 @@ class VLAWrapperBase(TensorDictModuleBase): output_mode (str, optional): ``"chunk"``, ``"tokens"`` or ``"both"``. Defaults to ``"chunk"`` for continuous heads and ``"tokens"`` for token heads. + return_vla_action_container (bool): whether to write the structured + :class:`~torchrl.data.vla.VLAAction` object at the VLA action root + key. When ``False``, only its plain TensorDict fields are written. + Defaults to ``True``. vocab_size (int, optional): Number of action-token bins, required for token heads. action_tokenizer (ActionTokenizerBase, optional): Token/chunk codec used @@ -124,6 +128,7 @@ def __init__( processor=None, input_mode: InputMode = "canonical", output_mode: OutputMode | None = None, + return_vla_action_container: bool = True, action_dim: int, chunk_size: int, action_head: ActionHead = "continuous", @@ -199,6 +204,7 @@ def __init__( self.action_head = action_head self.input_mode = input_mode self.output_mode = output_mode + self.return_vla_action_container = bool(return_vla_action_container) self.vocab_size = None if vocab_size is None else int(vocab_size) self.action_tokenizer = action_tokenizer self.action_kwargs = {} if action_kwargs is None else dict(action_kwargs) @@ -429,34 +435,40 @@ def _set_action( log_probs: torch.Tensor | None = None, mask: torch.Tensor | None = None, ) -> None: - action = VLAAction( - chunk=chunk, - tokens=tokens, - logits=logits, - log_probs=log_probs, - mask=mask, - batch_size=out.batch_size, - device=out.device, - ) - out.set(self._tensor_keys.vla_action, action) - if chunk is not None and not self._is_vla_field_key( - self._tensor_keys.action_chunk, "chunk" + if self.return_vla_action_container: + action = VLAAction( + chunk=chunk, + tokens=tokens, + logits=logits, + log_probs=log_probs, + mask=mask, + batch_size=out.batch_size, + device=out.device, + ) + out.set(self._tensor_keys.vla_action, action) + if chunk is not None and ( + not self.return_vla_action_container + or not self._is_vla_field_key(self._tensor_keys.action_chunk, "chunk") ): out.set(self._tensor_keys.action_chunk, chunk) - if tokens is not None and not self._is_vla_field_key( - self._tensor_keys.action_tokens, "tokens" + if tokens is not None and ( + not self.return_vla_action_container + or not self._is_vla_field_key(self._tensor_keys.action_tokens, "tokens") ): out.set(self._tensor_keys.action_tokens, tokens) - if logits is not None and not self._is_vla_field_key( - self._tensor_keys.action_logits, "logits" + if logits is not None and ( + not self.return_vla_action_container + or not self._is_vla_field_key(self._tensor_keys.action_logits, "logits") ): out.set(self._tensor_keys.action_logits, logits) - if log_probs is not None and not self._is_vla_field_key( - self._tensor_keys.log_probs, "log_probs" + if log_probs is not None and ( + not self.return_vla_action_container + or not self._is_vla_field_key(self._tensor_keys.log_probs, "log_probs") ): out.set(self._tensor_keys.log_probs, log_probs) - if mask is not None and not self._is_vla_field_key( - self._tensor_keys.action_mask, "mask" + if mask is not None and ( + not self.return_vla_action_container + or not self._is_vla_field_key(self._tensor_keys.action_mask, "mask") ): out.set(self._tensor_keys.action_mask, mask) diff --git a/torchrl/modules/vla/models.py b/torchrl/modules/vla/models.py index 6cd552ec832..83fb8c97290 100644 --- a/torchrl/modules/vla/models.py +++ b/torchrl/modules/vla/models.py @@ -68,6 +68,9 @@ class TinyVLA(VLAWrapperBase): mode (str, optional): backward-compatible alias for ``default_interaction_type``. Defaults to ``None``. device (DEVICE_TYPING, optional): device to move the parameters to. + return_vla_action_container (bool): whether to write the structured + VLAAction container in addition to its TensorDict fields. Defaults + to ``True``. Examples: >>> import torch @@ -105,6 +108,7 @@ def __init__( device: DEVICE_TYPING | None = None, input_mode: InputMode = "canonical", output_mode: OutputMode | None = None, + return_vla_action_container: bool = True, action_tokenizer: ActionTokenizerBase | None = None, return_log_probs: bool | None = None, return_logits: bool = False, @@ -120,6 +124,7 @@ def __init__( use_state=use_state, input_mode=input_mode, output_mode=output_mode, + return_vla_action_container=return_vla_action_container, action_tokenizer=action_tokenizer, default_interaction_type=default_interaction_type, log_probs_mode=log_probs_mode, @@ -157,13 +162,25 @@ def _hash_text(self, strings: list[str], device: torch.device) -> torch.Tensor: ] return torch.tensor(indices, dtype=torch.long, device=device) + def _flatten_instruction_strings(self, data) -> list[str]: + if isinstance(data, str): + return [data] + if isinstance(data, list): + result = [] + for item in data: + result.extend(self._flatten_instruction_strings(item)) + return result + return [str(data)] + def _instruction_strings(self, tensordict: TensorDictBase, batch: int) -> list[str]: instruction = self._get_instruction(tensordict) data = getattr(instruction, "tolist", lambda: instruction)() if isinstance(data, str): data = [data] * batch - elif not isinstance(data, list): - data = [str(data)] * batch + else: + data = self._flatten_instruction_strings(data) + if len(data) == 1 and batch != 1: + data = data * batch return data def _features(self, tensordict: TensorDictBase) -> torch.Tensor: @@ -172,12 +189,17 @@ def _features(self, tensordict: TensorDictBase) -> torch.Tensor: image = image.float() / 255.0 else: image = image.float() - batch = image.shape[0] - feats = [self.image_encoder(image)] + batch_shape = image.shape[:-3] + batch = int(torch.tensor(batch_shape).prod().item()) + image = image.reshape(batch, *image.shape[-3:]) + feats = [self.image_encoder(image).reshape(*batch_shape, -1)] if self.use_state: - feats.append(self.state_encoder(self._get_state(tensordict).float())) + state = self._get_state(tensordict).float() + state = state.reshape(batch, state.shape[-1]) + feats.append(self.state_encoder(state).reshape(*batch_shape, -1)) strings = self._instruction_strings(tensordict, batch) - feats.append(self.text_embedding(self._hash_text(strings, image.device))) + text = self.text_embedding(self._hash_text(strings, image.device)) + feats.append(text.reshape(*batch_shape, -1)) return self.trunk(torch.cat(feats, dim=-1)) def _predict(self, tensordict: TensorDictBase) -> torch.Tensor: diff --git a/torchrl/objectives/llm/grpo.py b/torchrl/objectives/llm/grpo.py index ecedc94ce79..81d6945beca 100644 --- a/torchrl/objectives/llm/grpo.py +++ b/torchrl/objectives/llm/grpo.py @@ -7,6 +7,7 @@ import contextlib import importlib.util import multiprocessing as mp +import os from collections import defaultdict, deque from dataclasses import dataclass @@ -1062,6 +1063,18 @@ class MCAdvantage(Transform): their transform inside the replay-buffer actor; use :class:`RayMCAdvantage` when only the grouping transform state should be centralized in Ray. + .. note:: Setting the environment variable + ``TORCHRL_MC_ADVANTAGE_LOCAL_QUEUES=1`` turns :meth:`share_memory_` + into a no-op: the grouping queues and counters stay process-local + instead of being moved to a multiprocessing manager. Use this to skip + the manager round-trips when a single process writes to the replay + buffer, or when every trajectory of a group is guaranteed to be + produced by the same writer process. Caveat: with multiple writer + processes, a GRPO group whose trajectories span several workers can + never complete -- each worker only sees its local fraction of the + group, so those trajectories are held in memory forever and never + written to the buffer. + .. warning:: This transform will flatten the input tensordicts and therefore is not compatible yet with replay buffers hosting storages of more than one dimension. @@ -1238,6 +1251,14 @@ def share_memory_(self) -> MCAdvantage: """Move replay-buffer write state to multiprocessing-shared objects.""" if self.is_shared: return self + if os.environ.get("TORCHRL_MC_ADVANTAGE_LOCAL_QUEUES") == "1": + torchrl_logger.info( + "MCAdvantage.share_memory_ skipped because " + "TORCHRL_MC_ADVANTAGE_LOCAL_QUEUES=1: grouping queues stay " + "process-local, so GRPO groups whose trajectories span several " + "writer processes will never complete." + ) + return self manager = mp.Manager() queues = {group: list(queue) for group, queue in self.queues.items()} self.queues = _MCAdvantageSharedQueues( @@ -1274,6 +1295,24 @@ def reset_stats(self) -> None: self.selected_trajectories = 0 self.unselected_trajectories = 0 + def clear_queues(self) -> None: + """Clear incomplete Monte-Carlo trajectory groups.""" + with self._state_lock: + self.queues.clear() + + def get_stats(self) -> dict[str, float | int]: + """Return a serializable snapshot of counters and pending queues.""" + with self._state_lock: + stats = {name: self._get_stat(name) for name in _MCADVANTAGE_STATS} + stats.update( + queued_groups=self.queued_groups, + queued_trajectories=self.queued_trajectories, + max_queued_trajectories_per_group=( + self.max_queued_trajectories_per_group + ), + ) + return stats + @property def queued_groups(self) -> int: """Number of incomplete groups currently held in memory.""" @@ -1306,7 +1345,41 @@ def _queue_delete(self, group) -> None: def forward(self, tensordict: TensorDictBase) -> GRPOLossOutput: return tensordict - def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase: + @staticmethod + def _concrete_if_possible(tensordict: TensorDictBase) -> TensorDictBase: + """Materialize lazy inputs unless their tensor leaves are ragged. + + Ragged leaves (e.g. variable-length token sequences in a lazy stack) + cannot be stacked into a contiguous tensordict; returning the lazy + input unchanged is always safe, so any ``RuntimeError`` raised by + ``contiguous()`` falls back to it rather than pattern-matching the + (version-dependent) error message. + """ + try: + return tensordict.contiguous() + except RuntimeError as err: + torchrl_logger.debug( + f"MCAdvantage: keeping lazy tensordict; contiguous() failed with: {err}" + ) + return tensordict + + @staticmethod + def _cat_tensordicts( + tensordicts: list[TensorDictBase], dim: int = 0 + ) -> TensorDictBase: + """Concatenate trajectory TensorDicts with a concrete output type. + + ``MCAdvantage`` may receive view/lazy TensorDicts from replay storage and + plain TensorDicts from direct collector writes. Queued and returned + trajectory groups should nevertheless have a consistent concrete + representation, both to avoid backend-specific lazy-stack assumptions + and to keep replay-buffer storage writes predictable. + """ + return torch.cat( + [MCAdvantage._concrete_if_possible(td) for td in tensordicts], dim + ) + + def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase | None: if self.verbose: torchrl_logger.info( f"Invoking MCAdvantage.\nData size: {tensordict.shape}.\nCurrent queue size: {len(self.queues)}.\nTotal queue content: {sum(len(q) for q in self.queues.values())}" @@ -1322,7 +1395,7 @@ def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase: tensordicts = tensordict.split(splits.tolist()) tensordicts = [self._inv_call(td) for td in tensordicts] tensordicts = [td for td in tensordicts if td is not None] - return torch.cat(tensordicts) if tensordicts else None + return self._cat_tensordicts(tensordicts) if tensordicts else None # Then we have a single trajectory with self._state_lock: return self._inv_call_single(tensordict) @@ -1338,7 +1411,7 @@ def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase: continue result.append(td_out) if result: - return torch.cat(result, 0) + return self._cat_tensordicts(result, 0) return def _inv_call_single(self, tensordict: TensorDictBase) -> TensorDictBase | None: @@ -1354,6 +1427,7 @@ def _inv_call_single(self, tensordict: TensorDictBase) -> TensorDictBase | None: ) if not tensordict[-1][self.done_key].all(): raise RuntimeError("Expected the trajectory to be done.") + tensordict = self._concrete_if_possible(tensordict) group = tensordict[0][self.prompt_key] if isinstance(group, torch.Tensor): # tensor group identifiers (e.g. an integer group id stamped @@ -1411,7 +1485,7 @@ def _inv_call_single(self, tensordict: TensorDictBase) -> TensorDictBase | None: self._queue_delete(group) self.completed_groups += 1 # Cat is the most robust way to combine the trajs - tds = torch.cat(trajs, -1) + tds = self._cat_tensordicts(trajs, -1) # Collect rewards reward = tds.get(self.rewards_key, as_nested_tensor=True) reward_mean = reward.values().mean() @@ -1493,7 +1567,7 @@ def _trajectory_advantage( torchrl_logger.info(f"Group returns: {returns=} {advantage=}") for traj, reward, adv in zip(trajs, rewards, advantage.unbind(0)): traj.set(self.advantage_key, adv.expand(reward.shape).clone()) - return torch.cat(trajs, -1) + return self._cat_tensordicts(trajs, -1) class _MCAdvantageRayActor: