From 5679bf1cd1028c6c0a18209c9bcabe646b523121 Mon Sep 17 00:00:00 2001 From: Vincent Moens Date: Mon, 6 Jul 2026 09:58:26 +0100 Subject: [PATCH] Update [ghstack-poisoned] --- sota-implementations/vla_grpo/README.md | 237 ++- .../vla_grpo/compare_simplevla_rollouts.py | 885 ++++++++++ .../vla_grpo/config/vla_grpo_libero.yaml | 105 +- .../vla_grpo/config/vla_grpo_toy.yaml | 57 +- sota-implementations/vla_grpo/openvla.py | 1358 ++++++++++++---- sota-implementations/vla_grpo/test_openvla.py | 800 ++++++--- sota-implementations/vla_grpo/utils.py | 1432 ++++++++++------- sota-implementations/vla_grpo/vla-grpo.py | 954 ++--------- test/collectors/test_evaluator.py | 19 + test/data/test_vla.py | 62 +- test/llm/test_llm_objectives.py | 51 + test/modules/test_actor.py | 11 + test/test_collectors.py | 75 + test/test_inference_server.py | 89 + test/transforms/test_action_transforms.py | 35 +- torchrl/collectors/_evaluator.py | 49 +- torchrl/collectors/_multi_async.py | 2 +- torchrl/collectors/_runner.py | 32 +- torchrl/collectors/utils.py | 2 +- torchrl/data/vla/preprocessing.py | 118 +- torchrl/data/vla/tokenizers.py | 63 +- torchrl/envs/transforms/_action.py | 122 +- torchrl/envs/transforms/_reward.py | 13 +- torchrl/modules/inference_server/_client.py | 29 + torchrl/modules/inference_server/_mp.py | 30 +- .../inference_server/_queue_transport.py | 8 +- torchrl/modules/inference_server/_server.py | 158 +- torchrl/modules/vla/common.py | 52 +- torchrl/modules/vla/models.py | 34 +- torchrl/objectives/llm/grpo.py | 58 +- 30 files changed, 4724 insertions(+), 2216 deletions(-) create mode 100644 sota-implementations/vla_grpo/compare_simplevla_rollouts.py diff --git a/sota-implementations/vla_grpo/README.md b/sota-implementations/vla_grpo/README.md index bc06d071f2e..2357e79a564 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 @@ -180,6 +180,27 @@ policy = OpenVLAOFTWrapper.from_pretrained( tokenizer = policy.action_tokenizer # decode tokens -> env actions ``` +The wrapper also has an explicit `policy.mode=l1` path for the official +continuous OpenVLA-OFT reference checkpoints. That mode loads the released +`action_head--150000_checkpoint.pt` and +`proprio_projector--150000_checkpoint.pt` components, uses two images +(`agentview` plus wrist) and the 8-D OpenVLA proprio vector, and writes a +continuous `("vla_action", "chunk")`. It is meant to validate the environment, +image preprocessing, proprio normalization, and evaluator/collector path +against the supervised reference policy: + +```text +policy.mode=l1 +policy.checkpoint=moojink/openvla-7b-oft-finetuned-libero-spatial +policy.use_wrist_image=true +policy.use_proprio=true +policy.num_images_in_input=2 +policy.lora_rank=0 +``` + +The PPO update still expects token log-probabilities, so this mode is for +reference/evaluation probes until a continuous-action GRPO loss is added. + Before spending RL compute on a checkpoint, validate the loading path by evaluating the SFT checkpoint greedily on its LIBERO suite through `torchrl.envs.LiberoEnv` (`init_state_mode="cycle"`, 50 trials/task) and compare @@ -204,47 +225,37 @@ throughput split into collection and optimization: - `throughput/train_decisions_per_s` - `throughput/optim_steps_per_s` -The collector can also be switched between synchronous and async execution -paths for throughput experiments: +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 -# fully synchronous baseline -python sota-implementations/vla_grpo/vla-grpo.py \ - collector.async_env=false collector.async_policy=false - -# asynchronous env slots, but no policy auto-batching -python sota-implementations/vla_grpo/vla-grpo.py \ - collector.async_env=true collector.async_policy=false env.num_envs=8 - -# asynchronous env slots plus auto-batched policy inference -python sota-implementations/vla_grpo/vla-grpo.py \ - collector.async_env=true collector.async_policy=true env.num_envs=8 \ - collector.server_max_batch_size=8 collector.server_timeout=0.01 +python sota-implementations/vla_grpo/vla-grpo.py --config-name vla_grpo_libero ``` -`collector.async_env=true` uses `AsyncBatchedCollector` so faster environment -slots do not wait at a global step barrier. `collector.async_policy=true` routes -policy calls through an inference server; with multiple async env slots this -enables auto-batching and logs `policy_server/*` counters such as average batch -size, request rate, and queue/forward latency. The `false/true` combination is -available as a policy-server plumbing ablation, but policy auto-batching is most -meaningful when several env slots submit requests concurrently. +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. -Eval rollouts can also be rendered to video (`logger.record_video=true`, on by -default). A dedicated single-environment recorder is built with +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 -`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`. - -Checkpointing is shared by the toy and LIBERO configs. `checkpoint_latest.pt` -is written to the hydra run directory every `checkpoint.save_iter` iterations; +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 @@ -267,47 +278,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: @@ -321,31 +338,77 @@ 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. - 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/compare_simplevla_rollouts.py b/sota-implementations/vla_grpo/compare_simplevla_rollouts.py new file mode 100644 index 00000000000..9e12b5744fe --- /dev/null +++ b/sota-implementations/vla_grpo/compare_simplevla_rollouts.py @@ -0,0 +1,885 @@ +#!/usr/bin/env python +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +"""Compare TorchRL VLA rollouts against the SimpleVLA-RL/VeRL rollout. + +The script loads one OpenVLA-OFT token policy and evaluates the same +``(task_id, LIBERO init-state id)`` trajectories through two rollout stacks: + +* TorchRL: ``LiberoEnv`` + VLA GRPO transforms + TorchRL policy wrapper. +* SimpleVLA-RL: the reference ``verl.workers.rollout.RobHFRollout`` trajectory + generator from the paper codebase. + +The output is a JSON file with per-trajectory records and side-by-side success +rates. Run it from a checkout with both LIBERO and SimpleVLA-RL available, for +example: + +.. code-block:: bash + + export SIMPLEVLA_RL_ROOT=/path/to/SimpleVLA-RL + export LIBERO_ROOT=/path/to/LIBERO + export PYTHONPATH="${SIMPLEVLA_RL_ROOT}:${LIBERO_ROOT}:${PYTHONPATH:-}" + python sota-implementations/vla_grpo/compare_simplevla_rollouts.py \ + --task-ids 0 --trial-ids 0 --policy-device cuda:0 --render-gpu 0 +""" +from __future__ import annotations + +import argparse +import importlib.util +import json +import os +import random +import subprocess +import sys +from collections.abc import Iterable +from dataclasses import asdict, dataclass +from datetime import datetime, UTC +from pathlib import Path +from typing import Any, Literal + +os.environ.setdefault("LIBERO_CONFIG_PATH", os.path.expanduser("~/.libero")) +os.environ.setdefault("ROBOT_PLATFORM", "LIBERO") +os.environ.setdefault("MUJOCO_GL", "egl") +os.environ.setdefault("PYOPENGL_PLATFORM", "egl") +os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") +os.environ.setdefault("WANDB_MODE", "disabled") + +_SIMPLEVLA_ROOT_ENV = os.environ.get("SIMPLEVLA_RL_ROOT") +_LIBERO_ROOT_ENV = os.environ.get("LIBERO_ROOT") +for _path in (_SIMPLEVLA_ROOT_ENV, _LIBERO_ROOT_ENV): + if _path and _path not in sys.path: + sys.path.insert(0, _path) + +import numpy as np +import torch + +import utils as vla_utils +from omegaconf import DictConfig, OmegaConf +from tensordict import TensorDict +from tensordict.nn import InteractionType, set_interaction_type +from torchrl._utils import logger as torchrl_logger +from torchrl.data.vla import ACTION_TOKENS_KEY +from torchrl.envs import LiberoEnv, TransformedEnv +from torchrl.envs.utils import step_mdp + +_has_verl = importlib.util.find_spec("verl") is not None +_DataProto = None +_RobHFRollout = None +_ORIGINAL_TORCH_LOAD = torch.load + + +def _torch_load_trusted_init_state(*args, **kwargs): + """Compatibility shim for LIBERO init-state files under torch>=2.6.""" + kwargs.setdefault("weights_only", False) + return _ORIGINAL_TORCH_LOAD(*args, **kwargs) + + +@dataclass +class _TrajectorySpec: + """One LIBERO task / init-state pair to evaluate.""" + + task_id: int + trial_id: int + + +@dataclass +class _TrajectoryResult: + """Per-trajectory rollout summary.""" + + stack: Literal["torchrl", "simplevla"] + task_id: int + trial_id: int + success: bool + outer_steps: int | None + env_steps: int + reward_sum: float | None = None + first_tokens: list[int] | None = None + first_action: list[float] | None = None + tokens: list[list[int]] | None = None + actions: list[list[float]] | None = None + + +@dataclass +class _StackSummary: + """Success-rate summary for one rollout stack.""" + + stack: Literal["torchrl", "simplevla"] + trajectories: int + successes: int + success_rate: float + env_steps: int + + +@dataclass +class _ParityResult: + """Per-trajectory parity metrics between the two rollout stacks.""" + + task_id: int + trial_id: int + success_agreement: bool + torchrl_success: bool + simplevla_success: bool + first_token_exact_match: bool | None + first_tokens_exact_match: bool | None + first_token_compare_count: int + first_token_mismatch_index: int | None + torchrl_first_token_at_mismatch: int | None + simplevla_first_token_at_mismatch: int | None + first_action_compare_count: int + first_action_max_abs_diff: float | None + first_action_mismatch_index: int | None + + +@dataclass +class _ParitySummary: + """Aggregate parity metrics for the report header.""" + + trajectories: int + success_agreements: int + success_agreement_rate: float + first_token_exact_matches: int + first_token_exact_match_rate: float + first_tokens_exact_matches: int + first_tokens_exact_match_rate: float + first_action_max_abs_diff: float | None + + +def _insert_optional_root(root: str | None) -> None: + if root and root not in sys.path: + sys.path.insert(0, root) + + +def _load_verl_rollout(simplevla_root: str | None): + """Load SimpleVLA-RL's VeRL rollout class as an optional dependency.""" + global _DataProto, _RobHFRollout, _has_verl + _insert_optional_root(simplevla_root) + if not _has_verl: + _has_verl = importlib.util.find_spec("verl") is not None + if not _has_verl: + raise ImportError( + "Could not import 'verl'. Set SIMPLEVLA_RL_ROOT or pass " + "--simplevla-root pointing to the SimpleVLA-RL checkout." + ) + if _DataProto is None or _RobHFRollout is None: + from verl import DataProto + from verl.workers.rollout import RobHFRollout + + _DataProto = DataProto + _RobHFRollout = RobHFRollout + return _DataProto, _RobHFRollout + + +def _git_commit(path: Path) -> str | None: + if not path.exists(): + return None + try: + return subprocess.check_output( + ["git", "-C", str(path), "rev-parse", "HEAD"], text=True + ).strip() + except Exception: + return None + + +def _jsonable(value: Any) -> Any: + if isinstance(value, torch.Tensor): + data = value.detach().cpu() + if data.numel() == 1: + return data.reshape(()).item() + return data.tolist() + if isinstance(value, np.ndarray): + return value.tolist() + if isinstance(value, np.generic): + return value.item() + if isinstance(value, Path): + return str(value) + if isinstance(value, dict): + return {str(key): _jsonable(val) for key, val in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(item) for item in value] + return value + + +def _parse_int_list(value: str) -> list[int]: + values: list[int] = [] + for part in value.split(","): + part = part.strip() + if not part: + continue + if "-" in part: + start, end = part.split("-", 1) + values.extend(range(int(start), int(end) + 1)) + else: + values.append(int(part)) + return values + + +def _trajectory_specs( + task_ids: Iterable[int], + trial_ids: Iterable[int], +) -> list[_TrajectorySpec]: + return [ + _TrajectorySpec(task_id=int(task_id), trial_id=int(trial_id)) + for task_id in task_ids + for trial_id in trial_ids + ] + + +def _set_seed(seed: int) -> None: + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +def _load_cfg(args: argparse.Namespace, task_id: int) -> DictConfig: + cfg = OmegaConf.load(Path(args.config)) + cfg.env.backend = "libero" + cfg.env.task_suite = args.task_suite + cfg.env.task_ids = [int(task_id)] + cfg.env.num_envs = 1 + cfg.env.eval_num_envs = 1 + cfg.env.parallel_group_repeats = False + cfg.env.render_gpu_ids = [int(args.render_gpu)] + cfg.env.eval_render_gpu_ids = [int(args.render_gpu)] + cfg.env.seed = int(args.seed) + cfg.policy.backend = "openvla" + cfg.policy.mode = "tokens" + cfg.policy.checkpoint = args.checkpoint + cfg.policy.unnorm_key = args.unnorm_key + cfg.policy.dataset_statistics = args.dataset_statistics + cfg.policy.device = str(args.policy_device) + cfg.policy.temperature = float(args.temperature) + cfg.policy.top_k = args.top_k + cfg.policy.use_wrist_image = False + cfg.policy.use_proprio = False + cfg.policy.num_images_in_input = 1 + cfg.policy.center_crop = bool(args.center_crop) + cfg.policy.image_backend = args.torchrl_image_backend + cfg.policy.gripper_binarize = True + cfg.policy.gripper_binarize_threshold = float(args.gripper_binarize_threshold) + cfg.policy.gripper_invert = bool(args.gripper_invert) + cfg.policy.lora_rank = int(args.lora_rank) + cfg.logger.backend = "none" + cfg.logger.mode = "disabled" + return cfg + + +def _simplevla_cfg(args: argparse.Namespace) -> DictConfig: + return OmegaConf.create( + { + "vla": "openvla-oft", + "pretrained_checkpoint": args.checkpoint, + "unnorm_key": args.unnorm_key, + "model_family": "openvla", + "task_suite_name": args.task_suite, + "num_steps_wait": int(args.settle_steps), + "center_crop": bool(args.center_crop), + "num_images_in_input": 1, + "use_proprio": False, + "action_chunks_len": int(args.chunk_size), + "temperature": float(args.temperature), + "do_sample": bool(args.sample_actions), + "micro_batch_size": int(args.simplevla_batch_size), + "val_micro_batch_size": int(args.simplevla_batch_size), + "max_prompt_length": int(args.max_prompt_length), + "experiment_name": args.experiment_name, + } + ) + + +def _policy_and_tokenizer(args: argparse.Namespace): + device = torch.device(args.policy_device) + cfg = _load_cfg(args, int(args.task_ids[0])) + policy = vla_utils.make_policy(cfg, device) + policy.eval() + tokenizer = vla_utils.make_action_tokenizer(cfg, policy) + return cfg, policy, tokenizer + + +def _make_torchrl_env( + cfg: DictConfig, + tokenizer, + spec: _TrajectorySpec, + args: argparse.Namespace, +) -> TransformedEnv: + env_kwargs = dict(cfg.env.env_kwargs or {}) + env_kwargs["render_gpu_device_id"] = int(args.render_gpu) + base = LiberoEnv( + args.task_suite, + task_id=int(spec.task_id), + camera_height=int(cfg.env.camera_height), + camera_width=int(cfg.env.camera_width), + env_kwargs=env_kwargs, + wrist_camera=None, + from_pixels=False, + max_episode_steps=int(args.max_env_steps), + settle_steps=int(args.settle_steps), + init_state_mode="fixed", + init_state_id=int(spec.trial_id), + ) + return TransformedEnv(base, vla_utils._chunk_transform(cfg, tokenizer)) + + +def _tensor_bool(value: torch.Tensor | None) -> bool: + if value is None: + return False + return bool(torch.as_tensor(value).reshape(-1).any().item()) + + +def _first_flat(value: torch.Tensor, limit: int) -> list[int] | list[float]: + data = value.detach().cpu() + while data.ndim > 2: + data = data[0] + return data.reshape(-1)[: min(data.numel(), limit)].tolist() + + +def _decode_env_actions(policy, tokenizer, tokens: torch.Tensor) -> torch.Tensor: + decoded = tokenizer.decode(tokens) + gripper_postprocess = getattr(policy, "gripper_postprocess", None) + if gripper_postprocess is not None: + decoded = gripper_postprocess.postprocess(decoded) + return decoded + + +def _torchrl_rollout_one( + cfg: DictConfig, + policy, + tokenizer, + spec: _TrajectorySpec, + args: argparse.Namespace, +) -> _TrajectoryResult: + env = _make_torchrl_env(cfg, tokenizer, spec, args) + reward_sum = 0.0 + first_tokens: list[int] | None = None + first_action: list[float] | None = None + tokens: list[list[int]] = [] + actions: list[list[float]] = [] + env_steps = 0 + outer_steps = 0 + success = False + try: + env.set_seed(int(args.seed)) + td = env.reset() + interaction = ( + InteractionType.RANDOM + if bool(args.sample_actions) + else InteractionType.DETERMINISTIC + ) + with torch.no_grad(), set_interaction_type(interaction): + for _ in range(int(args.max_outer_steps)): + outer_steps += 1 + policy_td = policy(td) + action_tokens = policy_td.get(ACTION_TOKENS_KEY).detach().cpu() + decoded = _decode_env_actions(policy, tokenizer, action_tokens) + tokens.append(_first_flat(action_tokens, action_tokens.numel())) + actions.append(_first_flat(decoded.float(), decoded.numel())) + if first_tokens is None: + first_tokens = tokens[-1][:32] + first_action = _first_flat(decoded.float(), 16) + next_td = env.step(policy_td) + reward = next_td.get(("next", "reward"), None) + if reward is not None: + reward_sum += float(torch.as_tensor(reward).sum().item()) + env_steps += int(args.chunk_size) + success = success or _tensor_bool( + next_td.get(("next", "success"), None) + ) + done = _tensor_bool(next_td.get(("next", "done"), None)) + if done or success: + break + td = step_mdp(next_td) + finally: + env.close() + return _TrajectoryResult( + stack="torchrl", + task_id=int(spec.task_id), + trial_id=int(spec.trial_id), + success=bool(success), + outer_steps=outer_steps, + env_steps=int(min(env_steps, int(args.max_env_steps))), + reward_sum=float(reward_sum), + first_tokens=first_tokens, + first_action=first_action, + tokens=tokens, + actions=actions, + ) + + +def _run_torchrl_stack( + args: argparse.Namespace, + specs: list[_TrajectorySpec], + policy, + tokenizer, +) -> list[_TrajectoryResult]: + results = [] + cfg_cache: dict[int, DictConfig] = {} + for spec in specs: + cfg = cfg_cache.setdefault(spec.task_id, _load_cfg(args, spec.task_id)) + result = _torchrl_rollout_one(cfg, policy, tokenizer, spec, args) + torchrl_logger.info( + "TorchRL rollout task=%s trial=%s success=%s env_steps=%s", + spec.task_id, + spec.trial_id, + result.success, + result.env_steps, + ) + results.append(result) + return results + + +def _make_simplevla_prompts(specs: list[_TrajectorySpec], meta_info: dict[str, Any]): + DataProto, _ = _load_verl_rollout(None) + batch = TensorDict( + { + "task_id": torch.tensor( + [[spec.task_id] for spec in specs], dtype=torch.int64 + ), + "trial_id": torch.tensor( + [[spec.trial_id] for spec in specs], dtype=torch.int64 + ), + "trial_seed": torch.full((len(specs), 1), -1, dtype=torch.int64), + }, + batch_size=[len(specs)], + ) + non_tensor_batch = { + "task_suite_name": np.array( + [meta_info["task_suite_name"] for _ in specs], dtype=object + ) + } + return DataProto( + batch=batch, + non_tensor_batch=non_tensor_batch, + meta_info={"n_samples": 1}, + ) + + +def _run_simplevla_batch( + rollout, + specs: list[_TrajectorySpec], + task_suite_name: str, + policy, + tokenizer, +) -> list[_TrajectoryResult]: + prompts = _make_simplevla_prompts( + specs, + {"task_suite_name": task_suite_name}, + ) + output = rollout.generate_sequences(prompts) + complete = output.batch["complete"].detach().cpu().bool() + finish_step = output.batch["finish_step"].detach().cpu().long() + responses = output.batch["responses"].detach().cpu().long() + vocab_size = getattr(rollout.module, "vocab_size", None) + num_bins = getattr(tokenizer, "num_bins", 256) + results = [] + for index, spec in enumerate(specs): + trajectory_tokens = responses[index].detach().cpu().long() + if vocab_size is not None: + trajectory_tokens = trajectory_tokens - (int(vocab_size) - int(num_bins)) + trajectory_actions = [ + _decode_env_actions( + policy, + tokenizer, + step_tokens.reshape(int(step_tokens.numel() // 7), 7), + ) + for step_tokens in trajectory_tokens + ] + results.append( + _TrajectoryResult( + stack="simplevla", + task_id=int(spec.task_id), + trial_id=int(spec.trial_id), + success=bool(complete[index].item()), + outer_steps=None, + env_steps=int(finish_step[index].item()), + first_tokens=trajectory_tokens[0].reshape(-1)[:32].tolist(), + first_action=trajectory_actions[0] + .detach() + .cpu() + .float() + .reshape(-1)[:16] + .tolist(), + tokens=[step.reshape(-1).tolist() for step in trajectory_tokens], + actions=[ + step.detach().cpu().float().reshape(-1).tolist() + for step in trajectory_actions + ], + ) + ) + return results + + +def _run_simplevla_stack( + args: argparse.Namespace, + specs: list[_TrajectorySpec], + policy, +) -> list[_TrajectoryResult]: + _load_verl_rollout(args.simplevla_root) + _, RobHFRollout = _load_verl_rollout(args.simplevla_root) + rollout = RobHFRollout(policy.model, _simplevla_cfg(args)) + results: list[_TrajectoryResult] = [] + batch_size = int(args.simplevla_batch_size) + # SimpleVLA-RL calls LIBERO's ``get_task_init_states``, which still uses + # ``torch.load(path)``. Torch>=2.6 defaults ``weights_only=True`` and + # rejects these trusted benchmark numpy arrays. Use PyTorch's process-wide + # compatibility knob so SimpleVLA's worker process also sees it. + old_force_no_weights_only = os.environ.get("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD") + os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = "1" + torch.load = _torch_load_trusted_init_state + try: + for start in range(0, len(specs), batch_size): + batch_specs = specs[start : start + batch_size] + batch_results = _run_simplevla_batch( + rollout, + batch_specs, + args.task_suite, + policy, + policy.action_tokenizer, + ) + for result in batch_results: + torchrl_logger.info( + "SimpleVLA rollout task=%s trial=%s success=%s env_steps=%s", + result.task_id, + result.trial_id, + result.success, + result.env_steps, + ) + results.extend(batch_results) + finally: + torch.load = _ORIGINAL_TORCH_LOAD + if old_force_no_weights_only is None: + os.environ.pop("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", None) + else: + os.environ["TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD"] = old_force_no_weights_only + return results + + +def _summary( + stack: Literal["torchrl", "simplevla"], + results: list[_TrajectoryResult], +) -> _StackSummary: + successes = sum(int(result.success) for result in results) + trajectories = len(results) + return _StackSummary( + stack=stack, + trajectories=trajectories, + successes=successes, + success_rate=float(successes / trajectories) if trajectories else float("nan"), + env_steps=sum(result.env_steps for result in results), + ) + + +def _first_mismatch( + torchrl_values: list[int] | list[float] | None, + simplevla_values: list[int] | list[float] | None, + *, + atol: float = 0.0, +) -> tuple[int, int | None]: + if torchrl_values is None or simplevla_values is None: + return 0, 0 if torchrl_values != simplevla_values else None + compare_count = min(len(torchrl_values), len(simplevla_values)) + for index in range(compare_count): + if abs(float(torchrl_values[index]) - float(simplevla_values[index])) > atol: + return compare_count, index + if len(torchrl_values) != len(simplevla_values): + return compare_count, compare_count + return compare_count, None + + +def _value_at(values: list[int] | None, index: int | None) -> int | None: + if values is None or index is None or index >= len(values): + return None + return int(values[index]) + + +def _max_abs_diff( + torchrl_values: list[float] | None, + simplevla_values: list[float] | None, +) -> float | None: + if torchrl_values is None or simplevla_values is None: + return None + count = min(len(torchrl_values), len(simplevla_values)) + if count == 0: + return None + torchrl_array = np.asarray(torchrl_values[:count], dtype=np.float64) + simplevla_array = np.asarray(simplevla_values[:count], dtype=np.float64) + return float(np.abs(torchrl_array - simplevla_array).max()) + + +def _parity_results( + torchrl_results: list[_TrajectoryResult], + simplevla_results: list[_TrajectoryResult], + *, + action_diff_atol: float, +) -> list[_ParityResult]: + parity = [] + for torchrl_result, simplevla_result in zip( + torchrl_results, simplevla_results, strict=True + ): + if (torchrl_result.task_id, torchrl_result.trial_id) != ( + simplevla_result.task_id, + simplevla_result.trial_id, + ): + raise RuntimeError( + "TorchRL and SimpleVLA results are not aligned: " + f"{torchrl_result.task_id}/{torchrl_result.trial_id} vs " + f"{simplevla_result.task_id}/{simplevla_result.trial_id}." + ) + token_count, token_mismatch = _first_mismatch( + torchrl_result.first_tokens, + simplevla_result.first_tokens, + ) + action_count, action_mismatch = _first_mismatch( + torchrl_result.first_action, + simplevla_result.first_action, + atol=float(action_diff_atol), + ) + first_token_exact_match = None + if ( + torchrl_result.first_tokens is not None + and simplevla_result.first_tokens is not None + ): + first_token_exact_match = ( + bool(torchrl_result.first_tokens) + and bool(simplevla_result.first_tokens) + and torchrl_result.first_tokens[0] == simplevla_result.first_tokens[0] + ) + parity.append( + _ParityResult( + task_id=int(torchrl_result.task_id), + trial_id=int(torchrl_result.trial_id), + success_agreement=( + bool(torchrl_result.success) == bool(simplevla_result.success) + ), + torchrl_success=bool(torchrl_result.success), + simplevla_success=bool(simplevla_result.success), + first_token_exact_match=first_token_exact_match, + first_tokens_exact_match=token_mismatch is None, + first_token_compare_count=int(token_count), + first_token_mismatch_index=token_mismatch, + torchrl_first_token_at_mismatch=_value_at( + torchrl_result.first_tokens, token_mismatch + ), + simplevla_first_token_at_mismatch=_value_at( + simplevla_result.first_tokens, token_mismatch + ), + first_action_compare_count=int(action_count), + first_action_max_abs_diff=_max_abs_diff( + torchrl_result.first_action, + simplevla_result.first_action, + ), + first_action_mismatch_index=action_mismatch, + ) + ) + return parity + + +def _parity_summary(parity: list[_ParityResult]) -> _ParitySummary: + trajectories = len(parity) + success_agreements = sum(int(item.success_agreement) for item in parity) + first_token_exact_matches = sum( + int(bool(item.first_token_exact_match)) for item in parity + ) + first_tokens_exact_matches = sum( + int(bool(item.first_tokens_exact_match)) for item in parity + ) + first_action_diffs = [ + item.first_action_max_abs_diff + for item in parity + if item.first_action_max_abs_diff is not None + ] + return _ParitySummary( + trajectories=trajectories, + success_agreements=success_agreements, + success_agreement_rate=( + float(success_agreements / trajectories) if trajectories else float("nan") + ), + first_token_exact_matches=first_token_exact_matches, + first_token_exact_match_rate=( + float(first_token_exact_matches / trajectories) + if trajectories + else float("nan") + ), + first_tokens_exact_matches=first_tokens_exact_matches, + first_tokens_exact_match_rate=( + float(first_tokens_exact_matches / trajectories) + if trajectories + else float("nan") + ), + first_action_max_abs_diff=( + float(max(first_action_diffs)) if first_action_diffs else None + ), + ) + + +def _first_mismatch_boundary(parity: list[_ParityResult]) -> dict[str, Any] | None: + for item in parity: + if item.first_token_mismatch_index is not None: + return { + "kind": "first_tokens", + "task_id": item.task_id, + "trial_id": item.trial_id, + "index": item.first_token_mismatch_index, + "torchrl": item.torchrl_first_token_at_mismatch, + "simplevla": item.simplevla_first_token_at_mismatch, + } + if item.first_action_mismatch_index is not None: + return { + "kind": "first_action", + "task_id": item.task_id, + "trial_id": item.trial_id, + "index": item.first_action_mismatch_index, + "max_abs_diff": item.first_action_max_abs_diff, + } + if not item.success_agreement: + return { + "kind": "success", + "task_id": item.task_id, + "trial_id": item.trial_id, + "torchrl": item.torchrl_success, + "simplevla": item.simplevla_success, + } + return None + + +def _write_report( + args: argparse.Namespace, + specs: list[_TrajectorySpec], + torchrl_results: list[_TrajectoryResult], + simplevla_results: list[_TrajectoryResult], +) -> Path: + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + stamp = datetime.now(UTC).strftime("%Y%m%d_%H%M%S") + out_path = out_dir / f"simplevla_torchrl_rollout_compare_{stamp}.json" + simplevla_root = Path(args.simplevla_root) if args.simplevla_root else None + parity = _parity_results( + torchrl_results, + simplevla_results, + action_diff_atol=float(args.action_diff_atol), + ) + report = { + "script": str(Path(__file__).resolve()), + "created_utc": datetime.now(UTC).isoformat(), + "args": vars(args), + "torch": torch.__version__, + "torchrl": getattr(__import__("torchrl"), "__version__", None), + "simplevla_root": str(simplevla_root) if simplevla_root is not None else None, + "simplevla_commit": ( + _git_commit(simplevla_root) if simplevla_root is not None else None + ), + "specs": [asdict(spec) for spec in specs], + "summaries": { + "torchrl": asdict(_summary("torchrl", torchrl_results)), + "simplevla": asdict(_summary("simplevla", simplevla_results)), + "parity": asdict(_parity_summary(parity)), + }, + "first_mismatch_boundary": _first_mismatch_boundary(parity), + "parity": [asdict(result) for result in parity], + "trajectories": { + "torchrl": [asdict(result) for result in torchrl_results], + "simplevla": [asdict(result) for result in simplevla_results], + }, + } + out_path.write_text(json.dumps(_jsonable(report), indent=2), encoding="utf-8") + torchrl_logger.info("Wrote rollout comparison report to %s", out_path) + return out_path + + +def main() -> None: + default_config = Path(__file__).parent / "config" / "vla_grpo_libero.yaml" + parser = argparse.ArgumentParser( + description="Compare TorchRL VLA and SimpleVLA-RL/VeRL LIBERO rollouts." + ) + parser.add_argument("--config", default=str(default_config)) + parser.add_argument("--simplevla-root", default=os.environ.get("SIMPLEVLA_RL_ROOT")) + parser.add_argument("--task-suite", default="libero_spatial") + parser.add_argument("--task-ids", type=_parse_int_list, default=[0]) + parser.add_argument("--trial-ids", type=_parse_int_list, default=[0]) + parser.add_argument( + "--checkpoint", default="Haozhan72/Openvla-oft-SFT-libero-spatial-traj1" + ) + parser.add_argument( + "--dataset-statistics", + default="moojink/openvla-7b-oft-finetuned-libero-spatial", + ) + parser.add_argument("--unnorm-key", default="libero_spatial_no_noops") + parser.add_argument("--policy-device", default="cuda:0") + parser.add_argument("--render-gpu", type=int, default=0) + parser.add_argument("--temperature", type=float, default=1.0) + parser.add_argument("--top-k", type=int, default=None) + parser.add_argument( + "--sample-actions", + action=argparse.BooleanOptionalAction, + default=False, + help="Use sampled token actions in both TorchRL and SimpleVLA rollouts.", + ) + parser.add_argument( + "--action-diff-atol", + type=float, + default=1e-6, + help="Tolerance for reporting the first decoded-action mismatch.", + ) + parser.add_argument("--gripper-binarize-threshold", type=float, default=0.0) + parser.add_argument( + "--gripper-invert", + action=argparse.BooleanOptionalAction, + default=True, + ) + parser.add_argument("--lora-rank", type=int, default=0) + parser.add_argument( + "--center-crop", + action=argparse.BooleanOptionalAction, + default=True, + ) + parser.add_argument( + "--torchrl-image-backend", + choices=["torchvision", "pil", "tensorflow"], + default="torchvision", + ) + parser.add_argument("--settle-steps", type=int, default=10) + parser.add_argument("--chunk-size", type=int, default=8) + parser.add_argument("--max-env-steps", type=int, default=512) + parser.add_argument("--max-outer-steps", type=int, default=64) + parser.add_argument("--max-prompt-length", type=int, default=512) + parser.add_argument("--simplevla-batch-size", type=int, default=1) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--out-dir", default="/tmp") + parser.add_argument("--experiment-name", default="simplevla_torchrl_compare") + args = parser.parse_args() + + if args.simplevla_root is None: + raise ValueError( + "Pass --simplevla-root or set SIMPLEVLA_RL_ROOT to the " + "SimpleVLA-RL checkout." + ) + _set_seed(int(args.seed)) + specs = _trajectory_specs(args.task_ids, args.trial_ids) + _, policy, tokenizer = _policy_and_tokenizer(args) + try: + _set_seed(int(args.seed)) + torchrl_results = _run_torchrl_stack(args, specs, policy, tokenizer) + _set_seed(int(args.seed)) + simplevla_results = _run_simplevla_stack(args, specs, policy) + finally: + del policy + torch.cuda.empty_cache() + report = _write_report(args, specs, torchrl_results, simplevla_results) + parity = _parity_results( + torchrl_results, + simplevla_results, + action_diff_atol=float(args.action_diff_atol), + ) + summaries = { + "torchrl": asdict(_summary("torchrl", torchrl_results)), + "simplevla": asdict(_summary("simplevla", simplevla_results)), + "parity": asdict(_parity_summary(parity)), + "first_mismatch_boundary": _first_mismatch_boundary(parity), + "report": str(report), + } + sys.stdout.write(json.dumps(_jsonable(summaries), indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/sota-implementations/vla_grpo/config/vla_grpo_libero.yaml b/sota-implementations/vla_grpo/config/vla_grpo_libero.yaml index 35ab8586f0b..319c27b50bf 100644 --- a/sota-implementations/vla_grpo/config/vla_grpo_libero.yaml +++ b/sota-implementations/vla_grpo/config/vla_grpo_libero.yaml @@ -1,19 +1,16 @@ # SimpleVLA-RL (arXiv:2509.09674) on LIBERO: OpenVLA-OFT token variant (7B), -# the paper's hyper-parameters. Parallel MuJoCo workers (one process each) -# feed a single training device; see the README for the multi-GPU notes and +# the paper's hyper-parameters. H100 fast mode uses shared policy-server +# MultiCollector rollout workers; see the README for the multi-GPU notes and # the LoRA fallback. Requires LIBERO (and its deps) plus transformers/timm. # -# Per-iteration accounting: groups_per_iter x group_size = 512 trajectories, +# Per-target-wave accounting: groups_per_iter x group_size = 320 trajectories, # each up to max_outer_steps = 64 chunk decisions (512 env steps / chunk 8). env: backend: libero task_suite: libero_spatial # libero_spatial | libero_object | libero_goal | libero_10 task_ids: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] # one worker per task (num_envs must cover them) - # Parallel MuJoCo processes (training); >= len(task_ids). For GRPO, each - # worker owns serial group ids, so keep num_envs <= groups_per_iter (ideally - # a divisor, often equal) or no worker may finish a full group per iteration. - num_envs: 10 + num_envs: 80 # false: each worker replays group_size rollouts serially for one group. # true: group_size workers share the same task/init state/group id and # collect a GRPO group in parallel. This improves group-completion throughput @@ -26,17 +23,20 @@ env: # batches can overcollect under one policy, but without a group barrier they # may also drift; the LIBERO wrapper groups parallel repeats by cycled # init-state id to reduce this waste. - parallel_group_repeats: false + parallel_group_repeats: true eval_num_envs: 10 # parallel MuJoCo processes (evaluation); >= len(task_ids) camera_height: 256 camera_width: 256 render_backend: egl # egl (GPU/headless) | osmesa (CPU fallback) - render_gpu_ids: [0] # EGL-visible ids; override to [0,1,...] to spread workers - eval_render_gpu_ids: null # null = reuse render_gpu_ids for eval workers + render_gpu_ids: [2, 3, 4, 5] # H100 fast mode render GPUs for rollout workers + eval_render_gpu_ids: [7] # reserve GPU 7 for eval/video rendering + render_gpu_device_zero_fallback: true + env_kwargs: null chunk_size: 8 # NUM_ACTIONS_CHUNK of the checkpoint max_env_steps: 512 # base env steps per episode (the paper's cap) max_outer_steps: 64 # = max_env_steps / chunk_size, in chunk decisions - train_init_state_mode: random # random | cycle | fixed; eval always cycles + train_init_state_mode: cycle # random | cycle | fixed; eval always cycles + train_init_state_id: 0 # selected LIBERO init state when train_init_state_mode=fixed seed: 0 tokenizer: @@ -44,6 +44,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,58 +57,54 @@ policy: temperature: 1.6 # rollout sampling temperature (greedy eval) top_k: null # null = full categorical; small values keep token exploration local use_wrist_image: false + # L1 reference path: set mode=l1, checkpoint=moojink/openvla-7b-oft-finetuned-libero-spatial, + # use_wrist_image=true, use_proprio=true, num_images_in_input=2. + action_head_file: action_head--150000_checkpoint.pt + proprio_projector_file: proprio_projector--150000_checkpoint.pt + use_proprio: false + num_images_in_input: 1 center_crop: true # the SimpleVLA-RL checkpoints are trained with augmentations - image_backend: torchvision # OpenVLA image preprocessing backend: torchvision | pil + image_backend: torchvision # OpenVLA preprocessing: torchvision | pil | tensorflow (reference) # the model emits a continuous gripper; LIBERO/robosuite needs a firm +/-1 # open/close (demo gripper is binary, -1 open / +1 close). SimpleVLA-RL - # applies sign(2*g - 1) then inverts, i.e. a 0.5 binarization threshold - # before the open/close sign flip. + # applies sign(2 * g - 1) then inverts, so threshold 0.0 here is equivalent + # to threshold 0.5 on the raw decoded gripper scalar. gripper_binarize: true - gripper_binarize_threshold: 0.5 + gripper_binarize_threshold: 0.0 gripper_invert: true lora_rank: 32 # de-risk fallback; set to 0/null for full fine-tuning (needs FSDP) lora_target_modules: [q_proj, k_proj, v_proj, o_proj] - device: null # null = auto (cuda if available) + device: cuda:0 # training device collector: group_size: 8 # n: rollouts per initial state (GRPO group) candidate_group_size: null # null = group_size; e.g. 16 samples 16 and selects 8 - # Initial states per iteration (512 trajectories). Must be >= env.num_envs; - # for best throughput, use a multiple of env.num_envs so no partial groups - # are discarded at the synchronous update boundary. If - # env.parallel_group_repeats=true, use env.num_envs / group_size instead. - groups_per_iter: 64 - # Safety cap in target group waves. The collector writes complete - # trajectories directly to the replay buffer every one-outer-step internal - # poll; one wave is approximately max_outer_steps * - # ceil(groups_per_iter * group_size / num_envs) polls. Keeping MCAdvantage - # queues across these same-policy polls lets partial GRPO groups finish - # without crossing a policy update. - max_collect_batches_per_iter: 1 - # If a capped wave produces too few useful replay decisions, keep the same - # rollout policy and queued partial groups for another capped wave instead - # of updating on a tiny replay batch and clearing those queues immediately. - max_same_policy_collect_attempts: 2 - min_replay_decisions: null # null/0 = collect one target group wave + # Initial states per iteration (320 trajectories). With shared fast mode and + # env.parallel_group_repeats=true, use + # num_collectors * envs_per_collector / group_size. + groups_per_iter: 40 + # null/0 = one full selected rollout set: + # groups_per_iter * group_size * max_outer_steps decisions. + min_replay_decisions: null total_iters: 100 # the paper's total_epochs - policy_device: null # null = policy.device; set e.g. cuda:1 for rollout inference - # Execution-mode switches for throughput ablations: - # - false/false: regular synchronous TorchRL Collector. - # - true/false: async env slots with one request per policy forward. - # - true/true: async env slots plus auto-batched policy inference. - # - false/true: sync env stepping through the policy server path. - async_env: false - async_policy: false - env_backend: threading # AsyncBatchedCollector env backend: threading | multiprocessing - policy_backend: threading # inference transport: threading | multiprocessing | ray | monarch - server_backend: thread # process server needs a policy_factory and is not used here - server_max_batch_size: null # null = env.num_envs when async_policy=true + policy_device: cuda:1 # H100 default: rollout/eval inference server device + # null = one model forward for the complete request batch. Set to 1 for + # single-environment numerical parity while keeping batched env collection. + policy_micro_batch_size: null + replay_wait_s: 0.5 + replay_log_s: 30.0 + num_collectors: 4 + envs_per_collector: 80 + num_threads: 1 + server_max_batch_size: 1 # one batched ParallelEnv request; avoids mixed eval/rollout modes server_min_batch_size: 1 - server_timeout: 0.01 + server_timeout: 0.001 server_collect_stats: true server_stats_window_size: 1024 max_inflight_per_env: 1 - storing_device: null + env_sub_threads: 1 + storing_device: cpu + use_buffers: null advantage: trajectory_return: sum # binary success return per trajectory @@ -120,6 +117,9 @@ advantage: buffer: device: cpu # decisions hold raw uint8 images; keep the buffer off-GPU + shared_init: true + capacity_group_waves: 7 # replay capacity as multiples of one target group wave + consume_after_n_samples: 1 # sampled decisions leave the buffer after one update pass loss: clip_epsilon: [0.2, 0.28] # asymmetric DAPO Clip-Higher (per-token, see ratio_level) @@ -142,12 +142,11 @@ logger: mode: online eval_iter: 4 # evaluate every N iterations (the paper's cadence) eval_episodes: 50 # greedy episodes per evaluation (cycled init states) - eval_process: false # true = run eval in a dedicated process - eval_async: false # with eval_process, submit eval and log result later - eval_device: null # null = policy.device/auto; set e.g. cuda:7 with eval_process - eval_state_dir: null # null = tempfile dir for trainable-weight snapshots - record_video: true # render eval rollouts via VideoRecorder (needs a logger) - video_episodes: 2 # greedy episodes per recorded video (one camera stream) + eval_backend: process # TorchRL Evaluator backend: thread | process | ray + eval_async: false # submit eval and log result later + eval_busy_policy: skip # skip | error | queue + eval_timeout_s: 1800.0 + eval_device: null # null = rollout device; set e.g. cuda:7 for eval video_fps: 8 # frames per second of the logged video (one frame per decision) checkpoint: diff --git a/sota-implementations/vla_grpo/config/vla_grpo_toy.yaml b/sota-implementations/vla_grpo/config/vla_grpo_toy.yaml index 4f84fe14779..f563bf8d26c 100644 --- a/sota-implementations/vla_grpo/config/vla_grpo_toy.yaml +++ b/sota-implementations/vla_grpo/config/vla_grpo_toy.yaml @@ -11,7 +11,15 @@ env: success_tol: 0.35 # sized so a random policy succeeds sometimes (cold-start signal) max_outer_steps: 6 # episode truncation, in chunk decisions render_size: 64 # side length of the from_pixels eval-video frame - num_envs: 1 # async-env workers; ToyVLAEnv grouped rollouts are per worker + num_envs: 1 + eval_num_envs: 1 + parallel_group_repeats: false + train_init_state_mode: random + render_backend: null + render_gpu_ids: null + eval_render_gpu_ids: null + render_gpu_device_zero_fallback: true + env_kwargs: null seed: 0 tokenizer: @@ -26,33 +34,24 @@ collector: group_size: 4 # n: rollouts per initial state (GRPO group) candidate_group_size: null # null = group_size; larger values oversample then select group_size groups_per_iter: 8 # initial states per iteration - # Safety cap in target group waves. The collector writes complete - # trajectories directly to the replay buffer every one-outer-step internal - # poll; one wave is approximately max_outer_steps * - # ceil(groups_per_iter * group_size / num_envs) polls. - max_collect_batches_per_iter: 1 - # Retry a capped wave once under the same policy if too few useful replay - # decisions have reached the replay buffer. - max_same_policy_collect_attempts: 2 min_replay_decisions: null # null/0 = collect one target group wave total_iters: 200 - # Execution-mode switches for throughput ablations: - # - false/false: regular synchronous TorchRL Collector. - # - true/false: async env slots with one request per policy forward. - # - true/true: async env slots plus auto-batched policy inference. - # - false/true: sync env stepping through the policy server path. - async_env: false - async_policy: false - env_backend: threading # AsyncBatchedCollector env backend: threading | multiprocessing - policy_backend: threading # inference transport: threading | multiprocessing | ray | monarch - server_backend: thread # process server needs a policy_factory and is not used here - server_max_batch_size: null # null = number of async envs when async_policy=true + replay_wait_s: 0.05 + replay_log_s: 30.0 + policy_device: null + policy_micro_batch_size: null + num_collectors: 1 + envs_per_collector: 1 + num_threads: 1 + env_sub_threads: 1 + server_max_batch_size: 1 server_min_batch_size: 1 server_timeout: 0.01 server_collect_stats: true server_stats_window_size: 1024 max_inflight_per_env: 1 - storing_device: null + storing_device: cpu + use_buffers: null advantage: trajectory_return: sum # binary success return per trajectory @@ -63,6 +62,9 @@ advantage: buffer: device: null # null = training device + shared_init: false + capacity_group_waves: 1 # replay capacity as multiples of one target group wave + consume_after_n_samples: 1 # sampled decisions leave the buffer after one update pass loss: clip_epsilon: [0.2, 0.28] # asymmetric DAPO clip; a float gives symmetric clipping @@ -85,14 +87,13 @@ logger: mode: online eval_iter: 10 # evaluate every N iterations eval_episodes: 50 # greedy episodes per evaluation - eval_process: false # true = run eval in a dedicated process - eval_async: false # with eval_process, submit eval and log result later - eval_device: null # null = policy.device/auto; set e.g. cuda:7 with eval_process - eval_state_dir: null # null = tempfile dir for trainable-weight snapshots - record_video: true # render eval rollouts via VideoRecorder (needs a logger) - video_episodes: 2 # greedy episodes per recorded video + eval_backend: thread # TorchRL Evaluator backend: thread | process | ray + eval_async: false # submit eval and log result later + eval_busy_policy: skip # skip | error | queue + eval_timeout_s: 300.0 + eval_device: null # null = rollout device; set e.g. cuda:7 for eval video_fps: 4 # frames per second of the logged video (one frame per decision) checkpoint: save_iter: 25 # save every N iterations (0 disables) - resume: null # path to a checkpoint_latest.pt to resume from + resume: null # path to a checkpoint_latest TensorDict directory to resume from diff --git a/sota-implementations/vla_grpo/openvla.py b/sota-implementations/vla_grpo/openvla.py index 01747b450e8..b9e27961629 100644 --- a/sota-implementations/vla_grpo/openvla.py +++ b/sota-implementations/vla_grpo/openvla.py @@ -2,17 +2,21 @@ # # 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 module also provides a separate continuous L1-head wrapper for the official +OpenVLA-OFT checkpoint family. That path is deterministic and intended for +reference/evaluation rollouts; the token GRPO training semantics stay +unchanged. + +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 +25,25 @@ 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 import nn 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 +51,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 +75,87 @@ # 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") + + +class _MLPResNetBlock(nn.Module): + """Residual MLP block used by the OpenVLA-OFT L1 action head.""" + + def __init__(self, dim: int) -> None: + super().__init__() + self.ffn = nn.Sequential(nn.LayerNorm(dim), nn.Linear(dim, dim), nn.ReLU()) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.ffn(x) + x + + +class _MLPResNet(nn.Module): + """Small residual MLP used by the OpenVLA-OFT L1 action head.""" + + def __init__( + self, + *, + num_blocks: int, + input_dim: int, + hidden_dim: int, + output_dim: int, + ) -> None: + super().__init__() + self.layer_norm1 = nn.LayerNorm(input_dim) + self.fc1 = nn.Linear(input_dim, hidden_dim) + self.relu = nn.ReLU() + self.mlp_resnet_blocks = nn.ModuleList( + [_MLPResNetBlock(hidden_dim) for _ in range(num_blocks)] + ) + self.layer_norm2 = nn.LayerNorm(hidden_dim) + self.fc2 = nn.Linear(hidden_dim, output_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.relu(self.fc1(self.layer_norm1(x))) + for block in self.mlp_resnet_blocks: + x = block(x) + return self.fc2(self.layer_norm2(x)) + + +class _L1RegressionActionHead(nn.Module): + """Continuous OpenVLA-OFT L1 action head.""" + + def __init__( + self, + *, + input_dim: int = 4096, + hidden_dim: int = 4096, + action_dim: int = 7, + ) -> None: + super().__init__() + from openvla_oft.constants import ACTION_DIM, NUM_ACTIONS_CHUNK + + self.action_dim = int(action_dim) + self.chunk_size = int(NUM_ACTIONS_CHUNK) + self.model = _MLPResNet( + num_blocks=2, + input_dim=input_dim * ACTION_DIM, + hidden_dim=hidden_dim, + output_dim=action_dim, + ) + + def predict_action(self, actions_hidden_states: torch.Tensor) -> torch.Tensor: + batch_size = actions_hidden_states.shape[0] + rearranged = actions_hidden_states.reshape(batch_size, self.chunk_size, -1) + return self.model(rearranged) + + +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,274 +164,91 @@ 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) -def register_openvla_oft() -> None: - """Register the vendored token-OFT classes with the transformers Auto* APIs.""" - from openvla_oft.configuration_prismatic import OpenVLAConfig - from openvla_oft.modeling_prismatic import OpenVLAForActionPrediction - from openvla_oft.processing_prismatic import ( - PrismaticImageProcessor, - PrismaticProcessor, - ) - from transformers import ( - AutoConfig, - AutoImageProcessor, - AutoModelForVision2Seq, - AutoProcessor, - ) - - # The vendored modeling code predates the attention-implementation - # dispatch added in recent transformers (>=4.5x), which probes - # ``_supports_sdpa`` / ``_supports_flash_attn_*`` class attributes during - # ``PreTrainedModel.__init__``. Declaring no support routes the model to - # eager attention -- exactly the path OpenVLA was trained and evaluated - # with -- so the numerics (and the SFT success rate) are preserved while - # the vendored code stays verbatim. - for attr in ( - "_supports_sdpa", - "_supports_flash_attn_2", - "_supports_flash_attn_3", - "_supports_flex_attn", - "_supports_attention_backend", - ): - if not hasattr(OpenVLAForActionPrediction, attr): - setattr(OpenVLAForActionPrediction, attr, False) - - try: - AutoConfig.register("openvla", OpenVLAConfig) - AutoImageProcessor.register(OpenVLAConfig, PrismaticImageProcessor) - AutoProcessor.register(OpenVLAConfig, PrismaticProcessor) - AutoModelForVision2Seq.register(OpenVLAConfig, OpenVLAForActionPrediction) - except ValueError: - # already registered - pass - - -class OpenVLAOFTWrapper(VLAWrapperBase): - """Token-head OpenVLA-OFT policy speaking the canonical VLA TensorDict schema. - - Reads ``("observation", "image")`` (uint8 ``[*B, 3, H, W]``), optionally - ``("observation", "wrist_image")``, and ``language_instruction``; writes - ``("vla_action", "tokens")`` (``[*B, chunk_size, action_dim]`` ids in the - 256-way action-token *window*) and their ``("vla_action", "log_probs")``. - Decode the tokens to environment actions with :attr:`action_tokenizer` - (e.g. through :class:`~torchrl.envs.transforms.ActionTokenizerTransform`). - - Args: - model: a vendored ``OpenVLAForActionPrediction`` (token variant). - processor: the matching ``PrismaticProcessor``. +class OpenVLAInputTransform(Transform): + """Build OpenVLA prompt tokens and image tensors from canonical VLA inputs. - Keyword Args: - unnorm_key (str, optional): dataset key of the normalization - statistics (e.g. ``"libero_spatial_no_noops"``). Defaults to the - checkpoint's single key when unambiguous. - temperature (float, optional): sampling temperature, applied - identically when sampling rollout actions and when recomputing - log-probabilities at loss time. Defaults to ``1.0``. - top_k (int, optional): if provided, restrict each action-token - categorical to its top-k logits before sampling and log-prob - recomputation. This keeps exploration local while preserving the - rollout/loss distribution contract. Defaults to ``None`` (full - categorical). - default_interaction_type (InteractionType, optional): token readout - when no exploration context is active (``RANDOM`` samples, else - argmax); the forward otherwise follows the ambient - :func:`~torchrl.envs.utils.exploration_type`. Defaults to - ``InteractionType.DETERMINISTIC``. See - :class:`~torchrl.modules.vla.VLAWrapperBase`. - log_probs_mode (str, optional): ``"sequence"`` or ``"token"``. See - :class:`~torchrl.modules.vla.VLAWrapperBase`. Defaults to - ``"sequence"``. - use_wrist_image (bool, optional): read - ``("observation", "wrist_image")`` as the second camera (the - checkpoint must have been trained with two input images). - Defaults to ``False``. - center_crop (bool, optional): center-crop the images to 90% area - before the processor resize, as in SimpleVLA-RL evaluation with - augmented training. Defaults to ``False``. - 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"``. - 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_invert (bool, optional): flip the gripper open/close sign after - (optional) binarization, for checkpoints whose convention is - opposite the environment's. Defaults to ``False``. + 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, - model, processor, *, - unnorm_key: str | None = None, - temperature: float = 1.0, - top_k: int | None = None, - default_interaction_type: InteractionType = InteractionType.DETERMINISTIC, - log_probs_mode: LogProbsMode = "sequence", use_wrist_image: bool = False, center_crop: bool = False, - image_backend: Literal["torchvision", "pil"] = "torchvision", - gripper_binarize: bool = False, - gripper_binarize_threshold: float = 0.0, - gripper_invert: 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: - from openvla_oft.constants import ACTION_DIM, NUM_ACTIONS_CHUNK - - 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}.") - num_bins = len(model.bin_centers) + 1 - if unnorm_key is None and model.norm_stats is not None: - if len(model.norm_stats) != 1: - raise ValueError( - "the checkpoint carries statistics for several datasets; " - f"pass unnorm_key explicitly (options: {sorted(model.norm_stats)})." - ) - unnorm_key = next(iter(model.norm_stats)) - action_tokenizer = None - if model.norm_stats is not None and unnorm_key is not None: - action_tokenizer = VocabTailActionTokenizer.from_norm_stats( - model.norm_stats, - unnorm_key, - num_bins=num_bins, - gripper_binarize=gripper_binarize, - gripper_binarize_threshold=gripper_binarize_threshold, - gripper_invert=gripper_invert, - ) + 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__( - action_dim=ACTION_DIM, - chunk_size=NUM_ACTIONS_CHUNK, - action_head="tokens", - vocab_size=num_bins, - action_tokenizer=action_tokenizer, - use_state=False, - default_interaction_type=default_interaction_type, - log_probs_mode=log_probs_mode, + in_keys=in_keys, + out_keys=[ + input_ids_key, + attention_mask_key, + pixel_values_key, + ], ) - 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.use_wrist_image = bool(use_wrist_image) self.center_crop = bool(center_crop) self.image_backend = image_backend - self.gripper_binarize = bool(gripper_binarize) - self.gripper_binarize_threshold = float(gripper_binarize_threshold) - self.gripper_invert = bool(gripper_invert) - self.num_bins = num_bins + 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.unnorm_key = unnorm_key self._image_preprocessor = self._make_image_preprocessor( processor, image_backend ) if self._image_preprocessor is not None: self.image_backend = self._image_preprocessor.backend - if self.use_wrist_image: - self.in_keys = [*self.in_keys, ("observation", "wrist_image")] - - # -- loading ---------------------------------------------------------- - @classmethod - def from_pretrained( - cls, - pretrained_model_name_or_path: str, - *, - torch_dtype: torch.dtype = torch.bfloat16, - device: torch.device | str | None = None, - dataset_statistics: str | None = None, - **kwargs, - ) -> OpenVLAOFTWrapper: - """Load a SimpleVLA-RL token-OFT checkpoint (e.g. ``Haozhan72/...``). - - Args: - dataset_statistics (str, optional): a path to an OpenVLA - ``dataset_statistics.json`` or a HF repo id shipping one, whose - action-normalization stats are merged into ``model.norm_stats``. - The SimpleVLA-RL LIBERO checkpoints omit their fine-tuning - dataset's stats (only the base pretraining datasets remain), so - point this at the matching official OFT release (e.g. - ``moojink/openvla-7b-oft-finetuned-libero-spatial``) -- it is the - same LIBERO data, hence the same normalization. - """ - from openvla_oft.modeling_prismatic import OpenVLAForActionPrediction - from openvla_oft.processing_prismatic import PrismaticProcessor - - from transformers import AutoConfig - - register_openvla_oft() - # The SimpleVLA-RL token-OFT checkpoints ship a base-Prismatic - # config.json that omits the OFT architecture flags the modeling code - # reads. ``use_proprio`` is the only one read unconditionally; this - # single-image token variant has no proprio projector, so default it - # to False (which skips the projector and the dependent proprio_dim). - config = AutoConfig.from_pretrained( - pretrained_model_name_or_path, trust_remote_code=False - ) - if not hasattr(config, "use_proprio"): - config.use_proprio = False - model = OpenVLAForActionPrediction.from_pretrained( - pretrained_model_name_or_path, - config=config, - torch_dtype=torch_dtype, - low_cpu_mem_usage=True, - trust_remote_code=False, - # the vendored model implements only eager attention - attn_implementation="eager", - ) - if dataset_statistics is not None: - extra = _load_dataset_statistics(dataset_statistics) - if getattr(model, "norm_stats", None) is None: - model.norm_stats = {} - model.norm_stats.update(extra) - if device is not None: - model = model.to(device) - model.eval() - processor = PrismaticProcessor.from_pretrained( - pretrained_model_name_or_path, trust_remote_code=False - ) - return cls(model, processor, **kwargs) - def make_action_tokenizer(self) -> VocabTailActionTokenizer: - """Return the matching window-id tokenizer (tokens -> env actions).""" - if self.action_tokenizer is None: - raise RuntimeError( - "OpenVLA-OFT action statistics are missing; cannot build an " - "action tokenizer." - ) - return self.action_tokenizer + 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) - # -- 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) @@ -348,17 +262,31 @@ def _to_pil(self, image: torch.Tensor): 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 def _instructions(self, tensordict: TensorDictBase, batch: int) -> list[str]: - instruction = tensordict.get(self.tensor_keys.instruction) + 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: @@ -388,19 +316,30 @@ def _image_processor_config( if ( input_sizes is None or normalize_params is None - or len(input_sizes) != 1 - or len(normalize_params) != 1 - or tuple(input_sizes[0][-2:]) != (_OPENVLA_IMAGE_SIZE, _OPENVLA_IMAGE_SIZE) + 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(normalize_params[0]["mean"], dtype=torch.float32), - torch.as_tensor(normalize_params[0]["std"], dtype=torch.float32), + 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"] + self, + processor, + image_backend: Literal["torchvision", "pil", "tensorflow"], ) -> OpenVLAImagePreprocessor | None: image_processor = getattr(processor, "image_processor", None) if image_processor is None: @@ -452,12 +391,10 @@ def _preprocess(self, images, wrist_images, instructions): 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), + 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) @@ -485,15 +422,103 @@ def _preprocess_slow(self, images, wrist_images, instructions): ) 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), + self._to_device(input_ids), + self._to_device(attention_mask), + self._to_device(pixel_values, dtype=self.dtype), ) - # -- the parallel-decoding forward ------------------------------------- + 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 @@ -506,11 +531,6 @@ def _window_logits(self, input_ids, attention_mask, pixel_values): 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 @@ -554,12 +574,739 @@ def _window_logits(self, input_ids, attention_mask, pixel_values): 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 + 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 + from openvla_oft.modeling_prismatic import OpenVLAForActionPrediction + from openvla_oft.processing_prismatic import ( + PrismaticImageProcessor, + PrismaticProcessor, + ) + from transformers import ( + AutoConfig, + AutoImageProcessor, + AutoModelForVision2Seq, + AutoProcessor, + ) + + # The vendored modeling code predates the attention-implementation + # dispatch added in recent transformers (>=4.5x), which probes + # ``_supports_sdpa`` / ``_supports_flash_attn_*`` class attributes during + # ``PreTrainedModel.__init__``. Declaring no support routes the model to + # eager attention -- exactly the path OpenVLA was trained and evaluated + # with -- so the numerics (and the SFT success rate) are preserved while + # the vendored code stays verbatim. + for attr in ( + "_supports_sdpa", + "_supports_flash_attn_2", + "_supports_flash_attn_3", + "_supports_flex_attn", + "_supports_attention_backend", + ): + if not hasattr(OpenVLAForActionPrediction, attr): + setattr(OpenVLAForActionPrediction, attr, False) + + try: + AutoConfig.register("openvla", OpenVLAConfig) + AutoImageProcessor.register(OpenVLAConfig, PrismaticImageProcessor) + AutoProcessor.register(OpenVLAConfig, PrismaticProcessor) + AutoModelForVision2Seq.register(OpenVLAConfig, OpenVLAForActionPrediction) + except ValueError: + # already registered + pass + + +class OpenVLAOFTWrapper(VLAWrapperBase): + """Token-head OpenVLA-OFT policy speaking the canonical VLA TensorDict schema. + + Reads ``("observation", "image")`` (uint8 ``[*B, 3, H, W]``), optionally + ``("observation", "wrist_image")``, and ``language_instruction``; writes + ``("vla_action", "tokens")`` (``[*B, chunk_size, action_dim]`` ids in the + 256-way action-token *window*) and their ``("vla_action", "log_probs")``. + Decode the tokens to environment actions with :attr:`action_tokenizer` + (e.g. through :class:`~torchrl.envs.transforms.ActionTokenizerTransform`). + + Args: + model: a vendored ``OpenVLAForActionPrediction`` (token variant). + processor: the matching ``PrismaticProcessor``. + + Keyword Args: + unnorm_key (str, optional): dataset key of the normalization + statistics (e.g. ``"libero_spatial_no_noops"``). Defaults to the + checkpoint's single key when unambiguous. + temperature (float, optional): sampling temperature, applied + identically when sampling rollout actions and when recomputing + log-probabilities at loss time. Defaults to ``1.0``. + top_k (int, optional): if provided, restrict each action-token + categorical to its top-k logits before sampling and log-prob + 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 + :func:`~torchrl.envs.utils.exploration_type`. Defaults to + ``InteractionType.DETERMINISTIC``. See + :class:`~torchrl.modules.vla.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). + Defaults to ``False``. + center_crop (bool, optional): center-crop the images to 90% area + before the processor resize, as in SimpleVLA-RL evaluation with + augmented training. Defaults to ``False``. + 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; ``"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`` 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``. + """ + + def __init__( + self, + model, + processor, + *, + 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", "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 + + 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}.") + num_bins = len(model.bin_centers) + 1 + if unnorm_key is None and model.norm_stats is not None: + if len(model.norm_stats) != 1: + raise ValueError( + "the checkpoint carries statistics for several datasets; " + f"pass unnorm_key explicitly (options: {sorted(model.norm_stats)})." + ) + unnorm_key = next(iter(model.norm_stats)) + action_tokenizer = None + if model.norm_stats is not None and unnorm_key is not None: + action_tokenizer = VocabTailActionTokenizer.from_norm_stats( + model.norm_stats, + unnorm_key, + num_bins=num_bins, + ) + super().__init__( + action_dim=ACTION_DIM, + chunk_size=NUM_ACTIONS_CHUNK, + action_head="tokens", + vocab_size=num_bins, + action_tokenizer=action_tokenizer, + 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 + self.gripper_binarize = bool(gripper_binarize) + self.gripper_binarize_threshold = float(gripper_binarize_threshold) + self.gripper_invert = bool(gripper_invert) + self.num_bins = num_bins + self.unnorm_key = unnorm_key + 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, + ) + 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")] + + # -- loading ---------------------------------------------------------- + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str, + *, + torch_dtype: torch.dtype = torch.bfloat16, + device: torch.device | str | None = None, + dataset_statistics: str | None = None, + **kwargs, + ) -> OpenVLAOFTWrapper: + """Load a SimpleVLA-RL token-OFT checkpoint (e.g. ``Haozhan72/...``). + + Args: + dataset_statistics (str, optional): a path to an OpenVLA + ``dataset_statistics.json`` or a HF repo id shipping one, whose + action-normalization stats are merged into ``model.norm_stats``. + The SimpleVLA-RL LIBERO checkpoints omit their fine-tuning + dataset's stats (only the base pretraining datasets remain), so + point this at the matching official OFT release (e.g. + ``moojink/openvla-7b-oft-finetuned-libero-spatial``) -- it is the + same LIBERO data, hence the same normalization. + """ + from openvla_oft.modeling_prismatic import OpenVLAForActionPrediction + from openvla_oft.processing_prismatic import PrismaticProcessor + + from transformers import AutoConfig + + register_openvla_oft() + # The SimpleVLA-RL token-OFT checkpoints ship a base-Prismatic + # config.json that omits the OFT architecture flags the modeling code + # reads. ``use_proprio`` is the only one read unconditionally; this + # single-image token variant has no proprio projector, so default it + # to False (which skips the projector and the dependent proprio_dim). + config = AutoConfig.from_pretrained( + pretrained_model_name_or_path, trust_remote_code=False + ) + if not hasattr(config, "use_proprio"): + config.use_proprio = False + model = OpenVLAForActionPrediction.from_pretrained( + pretrained_model_name_or_path, + config=config, + torch_dtype=torch_dtype, + low_cpu_mem_usage=True, + trust_remote_code=False, + # the vendored model implements only eager attention + attn_implementation="eager", + ) + if dataset_statistics is not None: + extra = _load_dataset_statistics(dataset_statistics) + if getattr(model, "norm_stats", None) is None: + model.norm_stats = {} + model.norm_stats.update(extra) + if device is not None: + model = model.to(device) + model.eval() + processor = PrismaticProcessor.from_pretrained( + pretrained_model_name_or_path, trust_remote_code=False + ) + return cls(model, processor, **kwargs) + + def make_action_tokenizer(self) -> VocabTailActionTokenizer: + """Return the matching window-id tokenizer (tokens -> env actions).""" + if self.action_tokenizer is None: + raise RuntimeError( + "OpenVLA-OFT action statistics are missing; cannot build an " + "action tokenizer." + ) + return self.action_tokenizer + + # -- preprocessing (shared by rollout and loss-time recompute) --------- + def _to_pil(self, image: torch.Tensor): + return self.input_transform._to_pil(image) + + def _instructions(self, tensordict: TensorDictBase, batch: int) -> list[str]: + return self.input_transform._instructions(tensordict, batch) + + def _prompt_input_ids(self, instruction: str) -> torch.Tensor | None: + return self.input_transform._prompt_input_ids(instruction) + + @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 len(input_sizes) != 1 + or len(normalize_params) != 1 + or tuple(input_sizes[0][-2:]) != (_OPENVLA_IMAGE_SIZE, _OPENVLA_IMAGE_SIZE) + ): + return None + return ( + int(input_sizes[0][-1]), + torch.as_tensor(normalize_params[0]["mean"], dtype=torch.float32), + torch.as_tensor(normalize_params[0]["std"], dtype=torch.float32), + ) + + def _make_image_preprocessor( + self, + processor, + image_backend: Literal["torchvision", "pil", "tensorflow"], + ) -> OpenVLAImagePreprocessor | None: + return self.input_transform._make_image_preprocessor(processor, image_backend) + + def _pixel_values_from_images(self, images: torch.Tensor) -> torch.Tensor | None: + return self.input_transform._pixel_values_from_images(images) + + def _preprocess(self, images, wrist_images, instructions): + return self.input_transform._preprocess(images, wrist_images, instructions) + + def _preprocess_slow(self, images, wrist_images, instructions): + return self.input_transform._preprocess_slow(images, wrist_images, instructions) + + # -- the parallel-decoding forward ------------------------------------- + def _window_logits(self, input_ids, attention_mask, pixel_values): + return self.model_transform._window_logits( + input_ids, attention_mask, pixel_values + ) def _action_logits(self, tensordict: TensorDictBase) -> torch.Tensor: + 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 + + +class OpenVLAOFTL1Wrapper(OpenVLAOFTWrapper): + """Continuous L1-head OpenVLA-OFT policy for reference/evaluation rollouts. + + This wrapper targets the official OpenVLA-OFT LIBERO checkpoints that use + two camera images and an 8-D proprio vector. It reads the canonical TorchRL + LIBERO observation, converts the default 9-D quaternion state to the + OpenVLA 8-D axis-angle state when needed, normalizes proprio using the + checkpoint statistics, and writes a continuous ``("vla_action", "chunk")``. + """ + + def __init__( + self, + model, + processor, + action_head: nn.Module, + proprio_projector: nn.Module | None, + *, + unnorm_key: str | None = None, + default_interaction_type: InteractionType = InteractionType.DETERMINISTIC, + use_proprio: bool = True, + use_wrist_image: bool = True, + center_crop: bool = True, + image_backend: Literal["torchvision", "pil", "tensorflow"] = "torchvision", + gripper_binarize: bool = True, + gripper_binarize_threshold: float = 0.0, + gripper_invert: bool = True, + return_vla_action_container: bool = True, + ) -> None: + from openvla_oft.constants import ACTION_DIM, NUM_ACTIONS_CHUNK + + if unnorm_key is None and model.norm_stats is not None: + if len(model.norm_stats) != 1: + raise ValueError( + "the checkpoint carries statistics for several datasets; " + f"pass unnorm_key explicitly (options: {sorted(model.norm_stats)})." + ) + unnorm_key = next(iter(model.norm_stats)) + if model.norm_stats is not None and unnorm_key not in model.norm_stats: + no_noops_key = f"{unnorm_key}_no_noops" + if no_noops_key in model.norm_stats: + unnorm_key = no_noops_key + VLAWrapperBase.__init__( + self, + action_dim=ACTION_DIM, + chunk_size=NUM_ACTIONS_CHUNK, + action_head="continuous", + use_state=use_proprio, + default_interaction_type=default_interaction_type, + return_vla_action_container=return_vla_action_container, + ) + self.model = model + self.processor = processor + self.action_head_module = action_head + self.proprio_projector = proprio_projector + self.unnorm_key = unnorm_key + self.use_proprio = bool(use_proprio) + self.use_wrist_image = bool(use_wrist_image) + self.center_crop = bool(center_crop) + self.image_backend = image_backend + self.gripper_binarize = bool(gripper_binarize) + self.gripper_binarize_threshold = float(gripper_binarize_threshold) + self.gripper_invert = bool(gripper_invert) + 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, + ) + self.gripper_postprocess = GripperPostProcessTransform( + action_key=self.tensor_keys.action_chunk, + rescale=True, + binarize=gripper_binarize, + threshold=gripper_binarize_threshold, + invert=gripper_invert, + ) + 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")] + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str, + *, + torch_dtype: torch.dtype = torch.bfloat16, + device: torch.device | str | None = None, + dataset_statistics: str | None = None, + action_head_file: str = "action_head--150000_checkpoint.pt", + proprio_projector_file: str = "proprio_projector--150000_checkpoint.pt", + use_proprio: bool = True, + num_images_in_input: int = 2, + **kwargs, + ) -> OpenVLAOFTL1Wrapper: + """Load an official continuous-head OpenVLA-OFT checkpoint.""" + from openvla_oft.constants import ACTION_DIM + from openvla_oft.modeling_prismatic import ( + OpenVLAForActionPrediction, + ProprioProjector, + ) + from openvla_oft.processing_prismatic import PrismaticProcessor + from openvla_oft.train_utils import load_component_state_dict + from transformers import AutoConfig + + register_openvla_oft() + config = AutoConfig.from_pretrained( + pretrained_model_name_or_path, trust_remote_code=False + ) + config.use_proprio = False + config.proprio_dim = 8 + model = OpenVLAForActionPrediction.from_pretrained( + pretrained_model_name_or_path, + config=config, + torch_dtype=torch_dtype, + low_cpu_mem_usage=True, + trust_remote_code=False, + attn_implementation="eager", + ) + model.vision_backbone.set_num_images_in_input(int(num_images_in_input)) + if dataset_statistics is not None: + model.norm_stats = _load_dataset_statistics(dataset_statistics) + if device is not None: + model = model.to(device) + model.eval() + processor = PrismaticProcessor.from_pretrained( + pretrained_model_name_or_path, trust_remote_code=False + ) + device = next(model.parameters()).device + dtype = next(model.parameters()).dtype + action_head = _L1RegressionActionHead( + input_dim=model.llm_dim, + hidden_dim=model.llm_dim, + action_dim=ACTION_DIM, + ).to(device=device, dtype=dtype) + action_head.load_state_dict( + load_component_state_dict( + _resolve_component_checkpoint( + pretrained_model_name_or_path, + action_head_file, + "action_head", + ) + ) + ) + action_head.eval() + proprio_projector = None + if use_proprio: + proprio_projector = ProprioProjector( + llm_dim=model.llm_dim, proprio_dim=8 + ).to(device=device, dtype=dtype) + proprio_projector.load_state_dict( + load_component_state_dict( + _resolve_component_checkpoint( + pretrained_model_name_or_path, + proprio_projector_file, + "proprio_projector", + ) + ) + ) + proprio_projector.eval() + return cls( + model, + processor, + action_head, + proprio_projector, + use_proprio=use_proprio, + **kwargs, + ) + + def _openvla_state(self, state: torch.Tensor) -> torch.Tensor: + if state.shape[-1] == 8: + return state + if state.shape[-1] != 9: + raise ValueError( + "OpenVLA L1 proprio expects an 8-D axis-angle state or the " + f"default 9-D LIBERO quaternion state, got shape {state.shape}." + ) + position = state[..., :3] + quat = state[..., 3:7] + gripper = state[..., 7:] + xyz = quat[..., :3] + w = quat[..., 3].clamp(-1.0, 1.0) + den = (1.0 - w.square()).clamp_min(0.0).sqrt() + angle = 2.0 * torch.acos(w) + axis_angle = torch.where( + den.unsqueeze(-1) > 1e-6, + xyz * angle.unsqueeze(-1) / den.clamp_min(1e-6).unsqueeze(-1), + torch.zeros_like(xyz), + ) + return torch.cat((position, axis_angle, gripper), dim=-1) + + def _normalize_proprio(self, state: torch.Tensor) -> np.ndarray: + if not self.use_proprio: + raise RuntimeError("OpenVLA L1 proprio normalization requires state input.") + stats = self.model.norm_stats[self.unnorm_key]["proprio"] + if "q01" in stats: + low = torch.as_tensor(stats["q01"], dtype=state.dtype, device=state.device) + high = torch.as_tensor(stats["q99"], dtype=state.dtype, device=state.device) + else: + low = torch.as_tensor(stats["min"], dtype=state.dtype, device=state.device) + high = torch.as_tensor(stats["max"], dtype=state.dtype, device=state.device) + mask = torch.as_tensor( + stats.get("mask", np.ones(low.shape, dtype=bool)), + dtype=torch.bool, + device=state.device, + ) + normalized = 2.0 * (state - low) / (high - low).clamp_min(1e-8) - 1.0 + normalized = torch.where(mask, normalized, state) + return normalized.clamp(-1.0, 1.0).detach().cpu().numpy().astype(np.float32) + + def _sort_padding_for_predict_action(self, input_ids, attention_mask): + pad_token_id = self.processor.tokenizer.pad_token_id + padding_mask = (~input_ids.ne(pad_token_id)).int() + sorted_indices = torch.argsort( + padding_mask, dim=1, descending=True, stable=True + ) + return ( + torch.gather(input_ids, 1, sorted_indices), + torch.gather(attention_mask, 1, sorted_indices), + ) + + def _postprocess_actions(self, actions: torch.Tensor) -> torch.Tensor: + return self.gripper_postprocess.postprocess(actions) + + def _predict_chunk(self, tensordict: TensorDictBase) -> torch.Tensor: images = tensordict.get(self.tensor_keys.image) batch_dims = images.shape[:-3] images = images.reshape(-1, *images.shape[-3:]) @@ -571,18 +1318,59 @@ def _action_logits(self, tensordict: TensorDictBase) -> torch.Tensor: 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 + input_ids, attention_mask = self._sort_padding_for_predict_action( + input_ids, attention_mask ) + proprio = None + if self.use_proprio: + state = self._get_state(tensordict).reshape(images.shape[0], -1).float() + proprio = self._normalize_proprio(self._openvla_state(state)) + chunks = [] + for index in range(images.shape[0]): + row_proprio = None if proprio is None else proprio[index : index + 1] + with torch.no_grad(): + actions, _ = self.model.predict_action( + input_ids=input_ids[index : index + 1], + pixel_values=pixel_values[index : index + 1], + attention_mask=attention_mask[index : index + 1], + unnorm_key=self.unnorm_key, + proprio=row_proprio, + proprio_projector=self.proprio_projector, + action_head=self.action_head_module, + noisy_action_projector=None, + use_film=False, + ) + actions = np.asarray(actions, dtype=np.float32) + if actions.ndim == 3 and actions.shape[0] == 1: + actions = actions[0] + chunks.append(torch.as_tensor(actions, device=images.device)) + chunk = torch.stack(chunks, dim=0).reshape( + *batch_dims, self.chunk_size, self.action_dim + ) + return self._postprocess_actions(chunk) + + def _predict(self, tensordict: TensorDictBase) -> torch.Tensor: + return self._predict_chunk(tensordict).flatten(-2) + + +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 cb89393da37..fec8241b899 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, + OpenVLAOFTL1Wrapper, + 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], + ) + + +def test_local_advantage_worker_metrics_and_reset(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", +) +def test_real_checkpoint_microbatch_tokens_and_actions_match(): + """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): @@ -141,6 +279,44 @@ def __init__(self): def get_input_embeddings(self): return self.embed + class _TinyL1OFT(nn.Module): + """Tiny deterministic stand-in for the continuous OpenVLA-OFT head.""" + + def __init__(self): + super().__init__() + self.param = nn.Parameter(torch.zeros(())) + self.norm_stats = { + "libero_spatial_no_noops": { + "proprio": { + "q01": [-1.0] * 8, + "q99": [1.0] * 8, + "mask": [True] * 8, + } + } + } + self.proprios = [] + self.pixel_shapes = [] + + def predict_action( + self, + *, + input_ids, + pixel_values, + attention_mask, + unnorm_key, + proprio, + proprio_projector, + action_head, + noisy_action_projector, + use_film, + ): + self.pixel_shapes.append(tuple(pixel_values.shape)) + self.proprios.append(None if proprio is None else proprio.copy()) + action = np.zeros((CHUNK, ACT_DIM), dtype=np.float32) + action[:, :-1] = 0.25 + action[:, -1] = 1.0 + return action, None + class _FakeProcessor: """Deterministic stand-in for PrismaticProcessor.""" @@ -240,302 +416,241 @@ def policy(): ) -def test_make_collector_casts_deviceless_env_to_cpu(monkeypatch): - captured = {} - - class _FakeCollector: - def __init__(self, *args, **kwargs): - captured["kwargs"] = kwargs - - class _FakeEnv: - batch_size = torch.Size([2]) - device = None - - cfg = SimpleNamespace( - collector=SimpleNamespace(groups_per_iter=2, group_size=1), - env=SimpleNamespace(max_outer_steps=3), - ) - monkeypatch.setattr(utils, "Collector", _FakeCollector) - - with warnings.catch_warnings(): - warnings.simplefilter("error") - collector = utils.make_collector( - cfg, _FakeEnv(), object(), torch.device("cuda:0") - ) - - assert isinstance(collector, _FakeCollector) - assert captured["kwargs"]["policy_device"] == torch.device("cuda:0") - assert captured["kwargs"]["env_device"] == torch.device("cpu") - - -def test_make_collector_parallel_groups_use_logical_worker_count(monkeypatch): +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", + } + logger.update(logger_overrides) + return SimpleNamespace(**logger) + + +def test_make_collector_uses_multicollector_policy_server(monkeypatch): captured = {} - class _FakeCollector: + class _FakeMultiCollector: def __init__(self, *args, **kwargs): - captured["kwargs"] = kwargs - - class _FakeEnv: - batch_size = torch.Size([32]) - device = torch.device("cpu") - - 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) - - with pytest.warns(UserWarning, match="parallel_group_repeats=true"): - collector = utils.make_collector( - cfg, _FakeEnv(), object(), torch.device("cuda:0") - ) - - assert isinstance(collector, _FakeCollector) - assert captured["kwargs"]["trajs_per_batch"] == 16 * 8 - assert captured["kwargs"]["frames_per_batch"] == 32 * 3 - - -def test_make_collector_parallel_group_wave_has_no_warning(monkeypatch): - captured = {} - - class _FakeCollector: - def __init__(self, *args, **kwargs): - captured["kwargs"] = kwargs - - class _FakeEnv: - batch_size = torch.Size([32]) - device = torch.device("cpu") - - 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")) - - assert isinstance(collector, _FakeCollector) - assert captured["kwargs"]["trajs_per_batch"] == 4 * 8 - assert captured["kwargs"]["frames_per_batch"] == 32 * 3 + captured["multi_args"] = args + captured["multi_kwargs"] = kwargs + class _FakeServer: + def __init__(self, **kwargs): + captured["server_kwargs"] = kwargs + self.transport = kwargs["transport"] -def test_make_collector_can_write_to_replay_buffer(monkeypatch): - captured = {} + def start(self): + captured["server_started"] = True + return self - class _FakeCollector: - def __init__(self, *args, **kwargs): - captured["kwargs"] = kwargs + class _FakeTransport: + def __init__(self, **kwargs): + captured["transport_kwargs"] = kwargs + self.clients = [] - class _FakeEnv: - batch_size = torch.Size([4]) - device = torch.device("cpu") + def client(self): + client = object() + self.clients.append(client) + return client + policy = _fake_token_policy() cfg = SimpleNamespace( - collector=SimpleNamespace(groups_per_iter=4, group_size=2), - env=SimpleNamespace(max_outer_steps=3), + collector=_complete_collector_cfg(policy_micro_batch_size=1), + env=_complete_toy_env_cfg(), ) replay_buffer = object() - def hook(_): - return None - - monkeypatch.setattr(utils, "Collector", _FakeCollector) + monkeypatch.setattr(utils, "MultiCollector", _FakeMultiCollector) + monkeypatch.setattr(utils, "ProcessInferenceServer", _FakeServer) + monkeypatch.setattr(utils, "MPTransport", _FakeTransport) - collector = utils.make_collector( + collector, server, eval_policy = utils.make_collector( cfg, - _FakeEnv(), - object(), + policy, torch.device("cpu"), + tokenizer=utils.UniformActionTokenizer(16, low=-1.0, high=1.0), replay_buffer=replay_buffer, - post_collect_hook=hook, ) - 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_collector_async_env_uses_async_batched_collector(monkeypatch): - captured = {} - - class _FakeAsyncCollector: - def __init__(self, *args, **kwargs): - captured["args"] = args - captured["kwargs"] = kwargs - - def __iter__(self): - return self - - def __next__(self): - raise StopIteration - - def server_stats(self, *, reset=False): - return {"requests": 0} - - def shutdown(self): - captured["shutdown"] = True - - class _FakeEnv: - batch_size = torch.Size([1]) - device = torch.device("cpu") - + 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(): + policy = _fake_token_policy() cfg = SimpleNamespace( - collector=SimpleNamespace( - groups_per_iter=4, - group_size=2, - async_env=True, - async_policy=True, - server_min_batch_size=2, - ), - env=SimpleNamespace( - 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, - seed=0, + 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), ) - monkeypatch.setattr(utils, "AsyncBatchedCollector", _FakeAsyncCollector) - - collector = utils.make_collector( - cfg, - _FakeEnv(), - object(), - torch.device("cpu"), - tokenizer=object(), - replay_buffer=object(), - ) - collector._ensure_collector() - assert len(captured["kwargs"]["create_env_fn"]) == 4 - assert captured["kwargs"]["yield_completed_trajectories"] - server_config = captured["kwargs"]["server_config"] - assert server_config.max_batch_size == 4 - assert server_config.min_batch_size == 2 + 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), + ) -def test_make_collector_async_env_without_policy_batching(monkeypatch): +def test_make_collector_allows_cross_subcollector_groups_with_shared_rb(monkeypatch): captured = {} - class _FakeAsyncCollector: + class _FakeMultiCollector: def __init__(self, *args, **kwargs): - captured["kwargs"] = kwargs - - def __iter__(self): - return self + captured["multi_args"] = args + captured["multi_kwargs"] = kwargs - def __next__(self): - raise StopIteration + class _FakeServer: + def __init__(self, **kwargs): + self.transport = kwargs["transport"] - def server_stats(self, *, reset=False): - return {} + def start(self): + return self - def shutdown(self): + class _FakeTransport: + def __init__(self, **kwargs): pass - class _FakeEnv: - batch_size = torch.Size([1]) - device = torch.device("cpu") + def client(self): + return object() + policy = _fake_token_policy() cfg = SimpleNamespace( - collector=SimpleNamespace( - groups_per_iter=2, - group_size=2, - async_env=True, - async_policy=False, - ), - env=SimpleNamespace( - 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=2, - seed=0, + 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), ) - monkeypatch.setattr(utils, "AsyncBatchedCollector", _FakeAsyncCollector) + replay_buffer = SimpleNamespace(shared=True) - collector = utils.make_collector( + monkeypatch.setattr(utils, "MultiCollector", _FakeMultiCollector) + monkeypatch.setattr(utils, "ProcessInferenceServer", _FakeServer) + monkeypatch.setattr(utils, "MPTransport", _FakeTransport) + + collector, _, _ = utils.make_collector( cfg, - _FakeEnv(), - object(), + policy, torch.device("cpu"), - tokenizer=object(), + tokenizer=utils.UniformActionTokenizer(16, low=-1.0, high=1.0), + replay_buffer=replay_buffer, ) - collector._ensure_collector() - server_config = captured["kwargs"]["server_config"] - assert server_config.max_batch_size == 1 - assert server_config.timeout == 0.0 + assert isinstance(collector, _FakeMultiCollector) + assert len(captured["multi_args"][0]) == 4 -def test_make_collector_sync_env_can_use_policy_server(monkeypatch): +def test_make_evaluator_process_backend_uses_factories(monkeypatch): captured = {} - class _FakeCollector: - def __init__(self, *args, **kwargs): - captured["collector_args"] = args - captured["collector_kwargs"] = kwargs - self.requested_frames_per_batch = kwargs["frames_per_batch"] - - def shutdown(self, *args, **kwargs): - captured["collector_shutdown"] = True - - def reset(self, *args, **kwargs): - captured["collector_reset"] = True - - class _FakeServer: - def __init__(self, *args, **kwargs): - captured["server_args"] = args - captured["server_kwargs"] = kwargs - - def start(self): - return self - - def shutdown(self): - captured["server_shutdown"] = True - - def stats(self, *, reset=False): - return {"requests": 0} + 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 - class _FakeEnv: - batch_size = torch.Size([2]) - device = None + monkeypatch.setattr(utils, "Evaluator", _FakeEvaluator) - policy = SimpleNamespace( - in_keys=["observation"], out_keys=[("vla_action", "tokens")] - ) + policy = _fake_token_policy() cfg = SimpleNamespace( - collector=SimpleNamespace( - groups_per_iter=2, - group_size=1, - async_policy=True, - ), - env=SimpleNamespace(max_outer_steps=3), + 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"), ) - monkeypatch.setattr(utils, "Collector", _FakeCollector) - monkeypatch.setattr(utils, "InferenceServer", _FakeServer) - - collector = utils.make_collector(cfg, _FakeEnv(), policy, torch.device("cpu")) - assert isinstance(collector, utils._ServerBackedCollector) - assert isinstance(captured["collector_args"][1], utils.PolicyClientModule) - assert captured["server_kwargs"]["server_config"].max_batch_size == 2 - assert captured["collector_kwargs"]["policy_device"] == torch.device("cpu") - assert captured["collector_kwargs"]["trust_policy"] is True - collector.shutdown() - assert captured["server_shutdown"] + 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 def test_make_replay_buffer_scales_capacity_with_overcollection(): @@ -543,19 +658,29 @@ def test_make_replay_buffer_scales_capacity_with_overcollection(): collector=SimpleNamespace( groups_per_iter=2, group_size=3, - max_collect_batches_per_iter=4, + 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, + ), ) 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(): @@ -564,7 +689,6 @@ def test_make_replay_buffer_scales_capacity_with_candidate_group_size(): 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( @@ -572,8 +696,14 @@ def test_make_replay_buffer_scales_capacity_with_candidate_group_size(): 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, + ), ) replay_buffer, advantage = utils.make_replay_buffer(cfg, torch.device("cpu")) @@ -584,13 +714,55 @@ def test_make_replay_buffer_scales_capacity_with_candidate_group_size(): assert advantage.candidate_selection_min_size == 4 +def test_make_action_tokenizer_returns_none_for_l1_openvla(): + cfg = SimpleNamespace( + policy=SimpleNamespace(backend="openvla", mode="l1"), + tokenizer=SimpleNamespace(vocab_size=16), + ) + assert ( + utils.make_action_tokenizer(cfg, SimpleNamespace(action_tokenizer=None)) is None + ) + + +def test_chunk_transform_without_tokenizer_consumes_vla_chunk(): + cfg = SimpleNamespace(env=SimpleNamespace(max_outer_steps=1)) + + transform = utils._chunk_transform(cfg, None) + + assert transform[0].out_keys_inv == [ACTION_CHUNK_KEY] + + +def test_chunk_transform_openvla_tokens_decodes_then_postprocesses(): + 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, + ), + ) + + transform = utils._chunk_transform( + cfg, + utils.UniformActionTokenizer(16, low=-1.0, high=1.0), + ) + + 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"] + + def test_libero_worker_assignment_serial_and_parallel_groups(): class _EnvCfg(dict): def __getattr__(self, name): return self[name] cfg = SimpleNamespace( - collector=SimpleNamespace(group_size=8), + collector=SimpleNamespace(group_size=8, candidate_group_size=None), env=_EnvCfg( { "task_ids": [10, 11, 12], @@ -636,13 +808,17 @@ def _fake_libero_env(*args, **kwargs): 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), + collector=SimpleNamespace(group_size=8, candidate_group_size=None), policy=SimpleNamespace(use_wrist_image=False), ) monkeypatch.setattr(utils, "LiberoEnv", _fake_libero_env) @@ -654,6 +830,7 @@ def _fake_libero_env(*args, **kwargs): 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 @@ -666,6 +843,74 @@ def test_forward_shapes(self, policy): assert out[ACTION_TOKENS_KEY].max() < N_BINS assert out[LOG_PROBS_KEY].shape == (2,) + 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() + + with torch.no_grad(): + stacked = policy.policy_stack(obs.clone()) + logits = policy._action_logits(obs.clone()) + + 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 + + out = transform(TensorDict({"action": actions}, batch_size=[2]))["action"] + + torch.testing.assert_close(out[0, :, -1], torch.ones(CHUNK)) + torch.testing.assert_close(out[1, :, -1], -torch.ones(CHUNK)) + + 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() + + with torch.no_grad(): + out = policy(obs) + + 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) + + 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 # at rollout time and the loss-time recompute agree exactly, so the @@ -729,6 +974,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") @@ -788,6 +1047,39 @@ def test_action_tokenizer_decode(self, policy): # decode -> encode is the identity on emitted tokens torch.testing.assert_close(tokenizer.encode(actions), out[ACTION_TOKENS_KEY]) + def test_l1_forward_uses_wrist_image_and_quaternion_proprio(self): + model = _TinyL1OFT() + policy = OpenVLAOFTL1Wrapper( + model, + _BatchProcessor(), + nn.Identity(), + nn.Identity(), + use_proprio=True, + use_wrist_image=True, + center_crop=False, + image_backend="pil", + ) + obs = _make_obs(batch=2) + obs["observation", "wrist_image"] = torch.randint( + 0, 256, (2, 3, 32, 32), dtype=torch.uint8 + ) + state = torch.zeros(2, 9) + state[:, 6] = 1.0 # identity quaternion in xyzw convention + obs["observation", "state"] = state + + out = policy(obs) + + assert out[ACTION_CHUNK_KEY].shape == (2, CHUNK, ACT_DIM) + assert model.pixel_shapes == [(1, 6, 224, 224), (1, 6, 224, 224)] + assert len(model.proprios) == 2 + for proprio in model.proprios: + assert proprio.shape == (1, 8) + np.testing.assert_allclose(proprio[0, 3:6], np.zeros(3)) + torch.testing.assert_close( + out[ACTION_CHUNK_KEY][..., -1], + -torch.ones(2, CHUNK), + ) + @pytest.mark.parametrize("ratio_level", ["sequence", "token"]) def test_clip_ppo_loss_integration(self, ratio_level): torch.manual_seed(0) diff --git a/sota-implementations/vla_grpo/utils.py b/sota-implementations/vla_grpo/utils.py index a3523118dd2..0ea1660fe9f 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 lazy_stack, TensorDictBase -from torchrl.collectors import AsyncBatchedCollector, 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,43 +52,118 @@ ToyVLAEnv, TransformedEnv, ) -from torchrl.envs.utils import ExplorationType, set_exploration_type +from torchrl.envs.utils import ExplorationType from torchrl.modules.inference_server import ( - InferenceServer, InferenceServerConfig, + MPTransport, PolicyClientModule, - ThreadingTransport, + 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 _cfg_get(section, key: str, default=None): - get = getattr(section, "get", None) - if get is not None: - return get(key, default) - return getattr(section, key, default) + +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 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) @@ -95,14 +173,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] @@ -111,15 +195,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, @@ -142,7 +285,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. @@ -155,34 +303,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 - - 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), - ) + OpenVLAOFTWrapper = _openvla_wrapper_cls(cfg.policy.mode) + + 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, @@ -196,8 +354,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( @@ -209,64 +369,98 @@ 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 - ) - 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") - ), - group_repeats=worker_group_repeats, - group_id_offset=group_id_offset, - group_id_mode="init_state" if parallel_group_repeats else "episode", + cfg, + worker_idx, + eval_mode=eval_mode, + render_gpu_device_id=render_gpu_device_id, ) + 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(getattr(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": - default = 1 + return 1 else: - default = cfg.env.eval_num_envs if eval_mode else cfg.env.num_envs - key = "eval_num_envs" if eval_mode else "num_envs" - return int(_cfg_get(cfg.env, key, default)) + return int(cfg.env.eval_num_envs if eval_mode else cfg.env.num_envs) def _validate_libero_env_count( @@ -276,7 +470,7 @@ def _validate_libero_env_count( parallel_group_repeats = ( not eval_mode and group_repeats is not None - and bool(_cfg_get(cfg.env, "parallel_group_repeats", False)) + and bool(cfg.env.parallel_group_repeats) ) task_coverage_envs = num_envs if parallel_group_repeats: @@ -294,7 +488,7 @@ def _validate_libero_env_count( "when env.parallel_group_repeats=true so every parallel " f"group has exactly {group_repeats} workers ({num_envs=})." ) - if _cfg_get(cfg.env, "train_init_state_mode", "random") == "random": + 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 " @@ -321,7 +515,7 @@ def _validate_libero_env_count( def _make_env_worker( cfg, - tokenizer: ActionTokenizerBase, + tokenizer: ActionTokenizerBase | None, worker_idx: int, *, group_repeats: int | None = None, @@ -329,19 +523,23 @@ def _make_env_worker( 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_seed = None if seed is None else int(seed) + int(worker_idx) + 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_get(cfg.env, "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, - group_id_offset=worker_idx * GROUP_ID_OFFSET, + group_id_offset=worker_idx_with_offset * GROUP_ID_OFFSET, batch_size=[], seed=worker_seed, device=device, @@ -353,56 +551,26 @@ def _make_env_worker( 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)) - - -def make_async_env_factories( - cfg, - tokenizer: ActionTokenizerBase, - *, - group_repeats: int | None = None, - seed: int | None = None, - device: torch.device | None = None, - eval_mode: bool = False, - from_pixels: bool = False, - num_envs: int | None = None, -) -> list[Callable[[], TransformedEnv]]: - """Build one transformed VLA env factory per async collection slot.""" - override = num_envs is not None - if num_envs is None: - num_envs = _num_envs_from_cfg(cfg, eval_mode=eval_mode) - if cfg.env.backend == "libero": - _validate_libero_env_count( - cfg, - num_envs, - group_repeats=group_repeats, - eval_mode=eval_mode, - override=override, - ) - return [ - partial( - _make_env_worker, + return TransformedEnv( + base, + _chunk_transform( cfg, tokenizer, - worker_idx, - group_repeats=group_repeats, - seed=seed, - device=device, - eval_mode=eval_mode, - from_pixels=from_pixels, - ) - for worker_idx in range(num_envs) - ] + decode_actions_in_env=decode_actions_in_env, + ), + ) def make_env( cfg, - tokenizer: ActionTokenizerBase, + tokenizer: ActionTokenizerBase | None, *, group_repeats: int | None = None, seed: int | None = None, @@ -410,6 +578,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( @@ -417,7 +588,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, @@ -431,55 +602,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, [ @@ -490,6 +620,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) ], @@ -503,7 +635,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( @@ -514,38 +685,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, @@ -564,6 +727,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) @@ -592,387 +762,198 @@ def make_optimizer(cfg, loss_module: ClipPPOLoss): return optim, scheduler -class _ServerBackedCollector: - """Collector wrapper that owns a thread-backed inference server.""" +_WEIGHT_STRATEGY = WeightStrategy(extract_as="tensordict") - def __init__(self, collector: Collector, server: InferenceServer) -> None: - self.collector = collector - self.server = server - self.requested_frames_per_batch = collector.requested_frames_per_batch - def __iter__(self): - return iter(self.collector) +def policy_weights(policy: torch.nn.Module) -> TensorDictBase: + """Detached CPU TensorDict snapshot of a policy's parameters and buffers.""" + return _WEIGHT_STRATEGY.extract_weights(policy).data.detach().clone().cpu() - def reset(self, *args, **kwargs) -> None: - self.collector.reset(*args, **kwargs) - def shutdown(self, *args, **kwargs) -> None: - try: - self.collector.shutdown(*args, **kwargs) - finally: - self.server.shutdown() - - def server_stats(self, *, reset: bool = False) -> dict[str, float | int]: - return self.server.stats(reset=reset) - - def __getattr__(self, name): - return getattr(self.collector, name) - - -class _BatchedPolicyClientModule(PolicyClientModule): - """Split a synchronous batched env observation into server requests.""" - - def forward(self, tensordict: TensorDictBase) -> TensorDictBase: - if tensordict.ndim == 0: - return super().forward(tensordict) - batch_size = tensordict.batch_size - flat_tensordict = tensordict.reshape(-1) - futures = [self.submit(td) for td in flat_tensordict.unbind(0)] - result = lazy_stack([future.result() for future in futures], 0) - result = result.reshape(batch_size) - self._check_policy_lag(result) - return result - - -class _AsyncReplayCollector: - """AsyncBatchedCollector wrapper that writes complete trajectories to replay.""" - - yields_complete_trajectories = True - requested_frames_per_batch = 1 - - def __init__( - self, - *, - create_env_fn: list[Callable[[], EnvBase]], - policy: VLAWrapperBase, - replay_buffer: TensorDictReplayBuffer | None, - collector_kwargs: dict, - ) -> None: - self._create_env_fn = create_env_fn - self._policy = policy - self._replay_buffer = replay_buffer - self._collector_kwargs = collector_kwargs - self._collector = None - self._iterator = None - self._last_server_stats: dict[str, float | int] = {} - self.num_envs = len(create_env_fn) - - def _ensure_collector(self): - if self._collector is None: - self._collector = AsyncBatchedCollector( - create_env_fn=self._create_env_fn, - policy=self._policy, - **self._collector_kwargs, - ) - self._iterator = iter(self._collector) - return self._collector - - def __iter__(self): - return self - - def __next__(self): - self._ensure_collector() - traj = next(self._iterator) - if self._replay_buffer is not None: - self._replay_buffer.extend(traj) - return traj - - def pause_collection(self) -> None: - if self._collector is None: - return - self._last_server_stats = self._collector.server_stats(reset=True) - self._collector.shutdown() - self._collector = None - self._iterator = None - - def reset(self, *args, **kwargs) -> None: - self.pause_collection() - - def shutdown(self, *args, **kwargs) -> None: - self.pause_collection() - - def server_stats(self, *, reset: bool = False) -> dict[str, float | int]: - if self._collector is not None: - return self._collector.server_stats(reset=reset) - result = dict(self._last_server_stats) - if reset: - self._last_server_stats = {} - return result +def sync_policy_server(policy_server: ProcessInferenceServer, policy: torch.nn.Module): + """Push trainer policy weights to the shared process policy server.""" + return policy_server.update_model_weights(policy_weights(policy)) + + +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: - env_get = getattr( - cfg.env, - "get", - lambda key, default=None: getattr(cfg.env, key, default), - ) return ( cfg.collector.group_size - if env_get("parallel_group_repeats", False) + if cfg.env.parallel_group_repeats else candidate_group_size(cfg) ) -def _server_config_from_collector(cfg, *, num_envs: int) -> InferenceServerConfig: - collector_get = getattr( - cfg.collector, - "get", - lambda key, default=None: getattr(cfg.collector, key, default), - ) - async_policy = bool(collector_get("async_policy", False)) - if async_policy: - max_batch_size = int(collector_get("server_max_batch_size", None) or num_envs) - min_batch_size = int(collector_get("server_min_batch_size", 1)) - timeout = float(collector_get("server_timeout", 0.01)) - else: - max_batch_size = 1 - min_batch_size = 1 - timeout = 0.0 +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=max_batch_size, - min_batch_size=min_batch_size, - timeout=timeout, - collect_stats=bool(collector_get("server_collect_stats", True)), - stats_window_size=int(collector_get("server_stats_window_size", 1024)), + 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, *, - tokenizer: ActionTokenizerBase | None = None, - replay_buffer: TensorDictReplayBuffer | None = None, + tokenizer: ActionTokenizerBase | None, + replay_buffer: TensorDictReplayBuffer, post_collect_hook: Callable[[TensorDictBase], None] | None = None, -) -> Collector: - """Build a VLA rollout collector. - - The default path is the synchronous TorchRL ``Collector`` used by the - original recipe. Set ``collector.async_env=true`` to use - ``AsyncBatchedCollector`` env slots; set ``collector.async_policy=true`` to - route policy calls through an inference server with configurable - auto-batching. The async-env path yields complete trajectories and writes - them to the replay buffer directly, so the training loop can use the same - replay/advantage machinery across execution modes. - - With no replay buffer on the synchronous path, each yielded batch holds - complete, done-terminated trajectories, concatenated along time - (``trajs_per_batch`` with ``traj_format="cat"``). With a replay buffer, - TorchRL's collector writer path pushes complete trajectories to storage as - each internal rollout batch finishes. - - 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. """ - collector_get = getattr( - cfg.collector, - "get", - lambda key, default=None: getattr(cfg.collector, key, default), - ) - async_env = bool(collector_get("async_env", False)) - async_policy = bool(collector_get("async_policy", False)) - if async_env: - num_envs = _num_envs_from_cfg(cfg) - else: - 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." - ) - 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." + "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." ) - # 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) - ) - server_config = _server_config_from_collector(cfg, num_envs=num_envs) - if async_env: - if tokenizer is None: - raise ValueError( - "tokenizer is required when collector.async_env=true so async " - "environment factories can decode action tokens." - ) - create_env_fn = make_async_env_factories( + + 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 = getattr(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 if cfg.env.backend == "toy" else None, - ) - return _AsyncReplayCollector( - create_env_fn=create_env_fn, - policy=policy, - replay_buffer=replay_buffer, - collector_kwargs={ - "frames_per_batch": 1, - "total_frames": -1, - "yield_completed_trajectories": True, - "env_backend": collector_get("env_backend", "threading"), - "policy_backend": collector_get("policy_backend", "threading"), - "server_backend": collector_get("server_backend", "thread"), - "server_config": server_config, - "policy_device": device, - "output_device": env_device, - "env_device": env_device, - "storing_device": collector_get("storing_device", None), - "max_inflight_per_env": collector_get("max_inflight_per_env", 1), - "verbose": bool(collector_get("verbose", False)), - }, + device=env_device, + worker_idx_offset=collector_idx * envs_per_collector, + render_gpu_device_id=_render_gpu_for_subcollector(cfg, collector_idx), ) - - collector_policy = policy - collector_policy_device = device - server = None - if async_policy: - transport = ThreadingTransport() - server = InferenceServer( - policy, - transport, - server_config=server_config, - policy_device=device, - output_device=env_device, - ).start() - collector_policy = _BatchedPolicyClientModule( - transport, - in_keys=getattr(policy, "in_keys", None), - out_keys=getattr(policy, "out_keys", None), - max_inflight=None, + 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, + propagate_interaction_type=True, ) - collector_policy_device = env_device - - collector = Collector( - env, - collector_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, + propagate_interaction_type=True, + ) + 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=collector_policy_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, - trust_policy=True if async_policy else None, ) - if server is not None: - return _ServerBackedCollector(collector, server) - return collector - + 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, @@ -996,29 +977,384 @@ 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) + ] + ) + 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" + 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 ) - recorder.dump(step=step) + ) + 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() + if micro_batches % accumulate: + 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 335862263ef..fd307e51d85 100644 --- a/sota-implementations/vla_grpo/vla-grpo.py +++ b/sota-implementations/vla_grpo/vla-grpo.py @@ -6,885 +6,207 @@ 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") - ) - - # 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_device = auto_device(cfg.policy.device) + rollout_device = ( + torch.device(cfg.collector.policy_device) + if cfg.collector.policy_device + else train_device ) - collector_get = getattr( - cfg.collector, - "get", - lambda key, default=None: getattr(cfg.collector, key, default), + eval_device = ( + torch.device(cfg.logger.eval_device) + if cfg.logger.eval_device + else rollout_device ) - async_rollout_env = bool(collector_get("async_env", False)) - async_rollout_policy = bool(collector_get("async_policy", False)) - train_group_repeats = ( - cfg.collector.group_size - if env_get("parallel_group_repeats", False) - else collection_group_size - ) - train_env = make_env( - cfg, - tokenizer, - group_repeats=train_group_repeats, - seed=cfg.env.seed, - device=device if cfg.env.backend == "toy" else None, - num_envs=1 if async_rollout_env 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 + ) + sync_policy_server(policy_server, policy) 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 = int(getattr(collector, "num_envs", 0)) or ( - train_env.batch_size[0] if train_env.batch_size else 1 - ) - if getattr(collector, "yields_complete_trajectories", False): - collect_polls_per_group_wave = max(episodes_per_iter, 1) - else: - 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 + sync_policy_server(policy_server, policy) + 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: - collector_server_stats = {} - 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 - - pause_collection = getattr(collector, "pause_collection", None) - if pause_collection is not None: - pause_collection() - server_stats = getattr(collector, "server_stats", None) - if server_stats is not None: - collector_server_stats = server_stats(reset=True) - - if min_replay_decisions > 0 and len(replay_buffer) < min_replay_decisions: - safety_cap_hit = True - 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 - ), - "collector/async_env": float(async_rollout_env), - "collector/async_policy": float(async_rollout_policy), - } - for key, value in collector_server_stats.items(): - if isinstance(value, (float, int)): - group_metrics[f"policy_server/{key}"] = value - # 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 - ): + 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) + + 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: + 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/test/collectors/test_evaluator.py b/test/collectors/test_evaluator.py index 2376c2fb7bd..08c3c2c5c84 100644 --- a/test/collectors/test_evaluator.py +++ b/test/collectors/test_evaluator.py @@ -718,6 +718,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) diff --git a/test/data/test_vla.py b/test/data/test_vla.py index aa5200817b4..26cd9fd965e 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,28 @@ 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) + assert decoded.dtype == torch.float64 + np.testing.assert_array_equal(decoded.numpy(), expected) + def test_from_norm_stats(self): norm_stats = { "libero_spatial_no_noops": { @@ -440,9 +487,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 +502,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/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_inference_server.py b/test/test_inference_server.py index 92db060ae26..8dbd2667311 100644 --- a/test/test_inference_server.py +++ b/test/test_inference_server.py @@ -17,6 +17,11 @@ from tensordict import lazy_stack, TensorDict from tensordict.base import TensorDictBase from tensordict.nn import TensorDictModule +from tensordict.nn.probabilistic import ( + interaction_type, + InteractionType, + set_interaction_type, +) from torchrl.modules.inference_server import ( InferenceClient, @@ -168,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 = [] @@ -294,6 +323,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): @@ -506,6 +558,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.""" @@ -820,6 +888,27 @@ def test_no_weight_sync(self): result = client(td) assert "action" in result.keys() + def test_policy_client_can_propagate_interaction_type(self): + class _InteractionPolicy(nn.Module): + def forward(self, td: TensorDictBase) -> TensorDictBase: + value = 1 if interaction_type() is InteractionType.RANDOM else 0 + return TensorDict( + {"action": torch.full(td.batch_size, value)}, + batch_size=td.batch_size, + ) + + transport = ThreadingTransport() + policy = _InteractionPolicy() + with InferenceServer(policy, transport, max_batch_size=4): + client = PolicyClientModule( + transport, + out_keys=["action"], + propagate_interaction_type=True, + ) + with set_interaction_type(InteractionType.RANDOM): + result = client(TensorDict({}, batch_size=[1])) + assert result["action"].item() == 1 + # --------------------------------------------------------------------------- # AsyncBatchedCollector tests 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..a664b26086f 100644 --- a/torchrl/collectors/_evaluator.py +++ b/torchrl/collectors/_evaluator.py @@ -706,25 +706,36 @@ def _extract_metrics_from_trajectories( ) -> 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")``. """ - 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: + 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 +743,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()) 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/_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/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..b23b6f876c8 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,11 @@ 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. See + :meth:`from_norm_stats`. Args: num_bins (int): number of bin edges per action dimension (the OpenVLA @@ -258,8 +262,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,8 +293,8 @@ 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) normalized = 2.0 * (actions - norm_low) / scale - 1.0 @@ -298,6 +305,38 @@ def encode(self, actions: torch.Tensor) -> torch.Tensor: return self.num_bins - digitized def decode(self, tokens: torch.Tensor) -> torch.Tensor: + 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().numpy() + norm_high = self.norm_high.detach().numpy() + norm_mask = self.norm_mask.detach().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) + return torch.from_numpy(np.ascontiguousarray(actions)) + # operate on the tokens' device (policy and env may differ): move the # small lookup/normalization buffers there so the output rides the # input device @@ -309,8 +348,8 @@ def decode(self, tokens: torch.Tensor) -> torch.Tensor: 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_low = self.norm_low.to(device=tokens.device, dtype=actions.dtype) + norm_high = self.norm_high.to(device=tokens.device, dtype=actions.dtype) 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) @@ -371,8 +410,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/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/_client.py b/torchrl/modules/inference_server/_client.py index 7af3a5329a5..2d335940a06 100644 --- a/torchrl/modules/inference_server/_client.py +++ b/torchrl/modules/inference_server/_client.py @@ -11,10 +11,20 @@ import torch from tensordict.base import TensorDictBase from tensordict.nn import TensorDictModuleBase +from tensordict.nn.probabilistic import interaction_type from tensordict.utils import NestedKey from torchrl.modules.inference_server._transport import InferenceTransport +_REMOTE_INTERACTION_TYPE_KEY = "_torchrl_inference_interaction_type" +_INTERACTION_TYPE_TO_CODE = { + "mode": 0, + "median": 1, + "mean": 2, + "random": 3, + "deterministic": 4, +} + class _ImmediateFuture: def __init__(self, result: TensorDictBase | BaseException): @@ -82,6 +92,10 @@ class PolicyClientModule(TensorDictModuleBase): policy_version_key (NestedKey, optional): key that contains the behavior policy version returned by the server. Defaults to ``"policy_version"``. + propagate_interaction_type (bool, optional): if ``True``, the active + :func:`tensordict.nn.interaction_type` is attached to each request + so the server can execute the remote policy under the caller's + exploration context. Defaults to ``False``. Examples: >>> import torch @@ -117,6 +131,7 @@ def __init__( target_policy_version: int | None = None, max_policy_lag: int | None = None, policy_version_key: NestedKey = "policy_version", + propagate_interaction_type: bool = False, ) -> None: super().__init__() if isinstance(client, InferenceTransport): @@ -128,6 +143,7 @@ def __init__( self.target_policy_version = target_policy_version self.max_policy_lag = max_policy_lag self.policy_version_key = policy_version_key + self.propagate_interaction_type = bool(propagate_interaction_type) self._inflight_sem = ( threading.BoundedSemaphore(max_inflight) if max_inflight is not None @@ -168,6 +184,19 @@ def submit(self, tensordict: TensorDictBase): Returns: Future-like object whose ``result()`` method returns a TensorDict. """ + if self.propagate_interaction_type: + current_interaction_type = interaction_type() + if current_interaction_type is not None: + tensordict = tensordict.clone(recurse=False) + tensordict.set( + _REMOTE_INTERACTION_TYPE_KEY, + torch.full( + tensordict.batch_size, + _INTERACTION_TYPE_TO_CODE[current_interaction_type.value], + dtype=torch.int8, + device=tensordict.device or torch.device("cpu"), + ), + ) release = self._acquire_inflight() submit = getattr(self.client, "submit", None) if submit is None: diff --git a/torchrl/modules/inference_server/_mp.py b/torchrl/modules/inference_server/_mp.py index ef512e0b674..c127f6ae323 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,14 +36,32 @@ 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._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: - return self._ctx.Queue() + if self._manager is None: + return self._ctx.Queue() + return self._manager.Queue() + + def __getstate__(self) -> dict: + state = super().__getstate__() + state["_manager"] = None + return state + + def close(self) -> None: + if self._manager is not None: + self._manager.shutdown() def client(self) -> _QueueInferenceClient: """Create an actor-side client with a dedicated response queue. 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 8d2fc4e8463..7f971ef59bf 100644 --- a/torchrl/modules/inference_server/_server.py +++ b/torchrl/modules/inference_server/_server.py @@ -4,6 +4,7 @@ # LICENSE file in the root directory of this source tree. from __future__ import annotations +import contextlib import multiprocessing as mp import queue import threading @@ -12,18 +13,72 @@ from concurrent.futures import Future from multiprocessing.synchronize import Event as MPEvent from statistics import mean +from typing import Any import torch from tensordict import lazy_stack from tensordict.base import TensorDictBase +from tensordict.nn.probabilistic import InteractionType, set_interaction_type from tensordict.utils import NestedKey from torch import nn +from torchrl.modules.inference_server._client import _REMOTE_INTERACTION_TYPE_KEY from torchrl.modules.inference_server._config import ( InferenceDeviceConfig, InferenceServerConfig, ) from torchrl.modules.inference_server._transport import InferenceTransport +from torchrl.weight_update import WeightStrategy + +_CODE_TO_INTERACTION_TYPE = { + 0: InteractionType.MODE, + 1: InteractionType.MEDIAN, + 2: InteractionType.MEAN, + 3: InteractionType.RANDOM, + 4: InteractionType.DETERMINISTIC, +} + + +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: @@ -159,7 +214,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 @@ -273,6 +328,30 @@ def _mark_weight_update(self) -> None: self._policy_version += 1 self._num_weight_updates += 1 + 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: @@ -330,8 +409,8 @@ def _poll_weight_update(self) -> None: return with self._model_lock: weights = ws.receive(timeout=0.0) - if weights is not None: - self._mark_weight_update() + if weights is not None: + self._mark_weight_update() def _set_policy_version(self, result_batch: TensorDictBase) -> TensorDictBase: """Annotate inference outputs with the behavior policy version.""" @@ -348,6 +427,30 @@ def _set_policy_version(self, result_batch: TensorDictBase) -> TensorDictBase: ) return result_batch.set(self.policy_version_key, version) + def _interaction_type_context(self, batch: TensorDictBase): + code = batch.get(_REMOTE_INTERACTION_TYPE_KEY, default=None) + if code is None: + return contextlib.nullcontext(), batch + if not isinstance(code, torch.Tensor): + interaction_type_value = _CODE_TO_INTERACTION_TYPE[int(code)] + else: + flat_code = code.reshape(-1) + if flat_code.numel() == 0: + return contextlib.nullcontext(), batch.exclude( + _REMOTE_INTERACTION_TYPE_KEY, inplace=False + ) + interaction_code = int(flat_code[0].item()) + if not flat_code.eq(interaction_code).all(): + raise RuntimeError( + "InferenceServer received a mixed interaction-type batch. " + "Use homogeneous server requests or a smaller max_batch_size." + ) + interaction_type_value = _CODE_TO_INTERACTION_TYPE[interaction_code] + return ( + set_interaction_type(interaction_type_value), + batch.exclude(_REMOTE_INTERACTION_TYPE_KEY, inplace=False), + ) + @torch.no_grad() def _run(self) -> None: self._init_weight_sync() @@ -403,10 +506,14 @@ def _run(self) -> None: ] forward_start = time.monotonic() with self._model_lock: - result_batch = self.model(batch) - if self.output_device is not None: - result_batch = result_batch.to(self.output_device) - result_batch = self._set_policy_version(result_batch) + interaction_context, batch = self._interaction_type_context( + batch + ) + with interaction_context: + result_batch = self.model(batch) + if self.output_device is not None: + result_batch = result_batch.to(self.output_device) + result_batch = self._set_policy_version(result_batch) forward_ms = (time.monotonic() - forward_start) * 1000.0 self._record_batch_stats( batch_size=len(callbacks), @@ -484,6 +591,22 @@ def _process_server_entry( "alive": server.is_alive, "policy_version": server.policy_version, } + elif command == "update_model_weights": + weights = kwargs["weights"] + mark_weight_update = kwargs.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 command == "shutdown": shutdown_event.set() payload = {"accepted": True} @@ -690,9 +813,14 @@ def _request_control( raise TimeoutError( f"Timed out waiting for ProcessInferenceServer {command!r}." ) - response_id, ok, payload = self._control_response_queue.get( - timeout=remaining - ) + try: + response_id, ok, payload = self._control_response_queue.get( + timeout=remaining + ) + except queue.Empty as exc: + raise TimeoutError( + f"Timed out waiting for ProcessInferenceServer {command!r}." + ) from exc if response_id != request_id: continue if not ok: @@ -732,6 +860,16 @@ def stats(self, *, reset: bool = False) -> dict[str, float | int]: """ return self._request_control("stats", {"reset": reset}) + def update_model_weights( + self, weights: TensorDictBase, *, mark_weight_update: bool = True + ) -> dict[str, bool]: + """Apply TensorDict weights to the model hosted by the child process.""" + return self._request_control( + "update_model_weights", + {"weights": weights, "mark_weight_update": mark_weight_update}, + timeout=300.0, + ) + def health(self) -> dict[str, int | bool | None]: """Return a lightweight child-process health snapshot.""" process = self._process 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..3805b91c5d0 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 @@ -1238,6 +1239,8 @@ 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": + return self manager = mp.Manager() queues = {group: list(queue) for group, queue in self.queues.items()} self.queues = _MCAdvantageSharedQueues( @@ -1274,6 +1277,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 +1327,33 @@ 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.""" + try: + return tensordict.contiguous() + except RuntimeError as err: + if "Failed to stack tensors within a tensordict" not in str(err): + raise + 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 +1369,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 +1385,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 +1401,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 +1459,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 +1541,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: