[Feature] Add async VLA GRPO collection modes

ghstack-source-id: af32c10
Pull-Request: #3909

Author: vmoens

sota-implementations/vla_grpo/README.md

@@ -204,6 +204,32 @@ 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:
+
+```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
+```
+
+`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.
+
 Eval rollouts can also be rendered to video (`logger.record_video=true`, on by
 default). A dedicated single-environment recorder is built with
 `from_pixels=True`: `ToyVLAEnv` renders the tracking scene, while `LiberoEnv`

sota-implementations/vla_grpo/config/vla_grpo_libero.yaml

@@ -91,6 +91,23 @@ collector:
   min_replay_decisions: null # null/0 = collect one target group wave
   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
+  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
 
 advantage:
   trajectory_return: sum # binary success return per trajectory

sota-implementations/vla_grpo/config/vla_grpo_toy.yaml

@@ -11,6 +11,7 @@ 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
   seed: 0
 
 tokenizer:
@@ -35,6 +36,23 @@ collector:
   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
+  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
 
 advantage:
   trajectory_return: sum # binary success return per trajectory

sota-implementations/vla_grpo/test_openvla.py

@@ -357,6 +357,187 @@ def hook(_):
     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")
+
+    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,
+        ),
+    )
+    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
+
+
+def test_make_collector_async_env_without_policy_batching(monkeypatch):
+    captured = {}
+
+    class _FakeAsyncCollector:
+        def __init__(self, *args, **kwargs):
+            captured["kwargs"] = kwargs
+
+        def __iter__(self):
+            return self
+
+        def __next__(self):
+            raise StopIteration
+
+        def server_stats(self, *, reset=False):
+            return {}
+
+        def shutdown(self):
+            pass
+
+    class _FakeEnv:
+        batch_size = torch.Size([1])
+        device = torch.device("cpu")
+
+    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,
+        ),
+    )
+    monkeypatch.setattr(utils, "AsyncBatchedCollector", _FakeAsyncCollector)
+
+    collector = utils.make_collector(
+        cfg,
+        _FakeEnv(),
+        object(),
+        torch.device("cpu"),
+        tokenizer=object(),
+    )
+    collector._ensure_collector()
+
+    server_config = captured["kwargs"]["server_config"]
+    assert server_config.max_batch_size == 1
+    assert server_config.timeout == 0.0
+
+
+def test_make_collector_sync_env_can_use_policy_server(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 _FakeEnv:
+        batch_size = torch.Size([2])
+        device = None
+
+    policy = SimpleNamespace(
+        in_keys=["observation"], out_keys=[("vla_action", "tokens")]
+    )
+    cfg = SimpleNamespace(
+        collector=SimpleNamespace(
+            groups_per_iter=2,
+            group_size=1,
+            async_policy=True,
+        ),
+        env=SimpleNamespace(max_outer_steps=3),
+    )
+    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"]
+
+
 def test_make_replay_buffer_scales_capacity_with_overcollection():
     cfg = SimpleNamespace(
         collector=SimpleNamespace(