Spike/vectorized env

roboeval/batched/__init__.py

@@ -0,0 +1,32 @@
+"""Vectorized-environment namespace for roboeval.
+
+Single-env users should keep importing from ``roboeval`` directly. This
+subpackage adds parallel types/Protocols for N-env vectorized rollouts
+(Isaac Lab, gym.vector, Brax, MJX), strictly isolated from the single-env
+public API so existing callers see no behavior change.
+"""
+
+from .environment import BatchedEnvironmentAdapter
+from .policy import (
+    BatchedPolicy,
+    BatchedPolicyAdapter,
+    from_single,
+    normalize_batched_policy,
+)
+from .runner import BatchedEvalRunner
+from .scheduler import SlotScheduler, SlotTask
+from .types import BatchedState, BatchedStepOutcome
+
+
+__all__ = [
+    "BatchedEnvironmentAdapter",
+    "BatchedEvalRunner",
+    "BatchedPolicy",
+    "BatchedPolicyAdapter",
+    "BatchedState",
+    "BatchedStepOutcome",
+    "SlotScheduler",
+    "SlotTask",
+    "from_single",
+    "normalize_batched_policy",
+]

roboeval/batched/environment.py

@@ -0,0 +1,64 @@
+"""BatchedEnvironmentAdapter Protocol — the vectorized counterpart of
+EnvironmentAdapter.
+
+A batched adapter wraps a vectorized simulator (Isaac Lab ManagerBasedRLEnv,
+gym.vector.VectorEnv, Brax, MJX, etc.) and exposes a uniform N-env interface
+to the runner.
+
+Method semantics:
+  * reset(scenarios)        — bulk reset all num_envs slots. ``len(scenarios)``
+                              must equal ``num_envs``. Returns one State per
+                              slot.
+  * step(actions)           — apply one action per slot, return a
+                              BatchedStepOutcome with num_envs entries. The
+                              runner detects terminal slots from
+                              ``outcome.terminals[i]``.
+  * reset_slots(slots,
+                scenarios)  — selective reset for slots that just terminated.
+                              Slots are 0-indexed; ``len(scenarios)`` must
+                              match ``len(slots)``. Returns the new initial
+                              state for each reset slot. Implementations whose
+                              underlying sim auto-resets (Gymnasium vector)
+                              should still implement this — at minimum to
+                              re-seed and to read off the post-reset obs.
+
+The Protocol intentionally does NOT take a scenario on step(): scenarios are
+fixed for the lifetime of the slot's current episode and the runner tracks
+the mapping. Reset is where scenario-dependent state (seeds, options) flows
+into the env.
+
+A vectorized adapter should expose ``num_envs`` as either an attribute or a
+read-only property.
+"""
+
+from __future__ import annotations
+
+from typing import Protocol
+
+from roboeval.core import Action, Scenario, State
+
+from .types import BatchedState, BatchedStepOutcome
+
+
+class BatchedEnvironmentAdapter(Protocol):
+    """Protocol every vectorized environment adapter should implement."""
+
+    num_envs: int
+
+    def reset(self, scenarios: list[Scenario]) -> BatchedState:
+        """Bulk-reset all slots. ``len(scenarios) == num_envs`` required."""
+        ...
+
+    def step(self, actions: list[Action]) -> BatchedStepOutcome:
+        """Apply one action per slot and return the batched transition."""
+        ...
+
+    def reset_slots(
+        self, slots: list[int], scenarios: list[Scenario]
+    ) -> list[State]:
+        """Reset a subset of slots, returning the new initial state for each.
+
+        Called by the runner when slots terminate and get new scenario
+        assignments. ``len(slots) == len(scenarios)`` required.
+        """
+        ...

roboeval/batched/policy.py

@@ -0,0 +1,102 @@
+"""BatchedPolicy Protocol + adapter + single-state shim.
+
+A batched policy receives a list of per-slot states and returns one Decision
+per slot. This is the right interface for GPU-batched inference (VLAs,
+diffusion policies, transformer-based controllers) where calling the model
+N times sequentially defeats the point of vectorization.
+
+For legacy single-state policies (plain functions taking one State and
+returning a dict/tuple/Decision), use ``from_single(fn)`` to get a
+BatchedPolicy that iterates internally. It is a strict throughput regression
+relative to a natively-batched policy, but it preserves backward compat for
+the existing single-env policy zoo.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable, Protocol
+
+from roboeval.adapters import _normalize_decision
+from roboeval.core import Decision, State
+
+from .types import BatchedState
+
+
+class BatchedPolicy(Protocol):
+    """Protocol every batched policy should implement."""
+
+    name: str
+
+    def decide(self, states: BatchedState) -> list[Decision]:
+        """Return one Decision per slot. ``len(states) == len(returned)``."""
+        ...
+
+
+@dataclass(frozen=True)
+class BatchedPolicyAdapter:
+    """Normalize a user-provided batched policy to the BatchedPolicy Protocol.
+
+    Accepts either:
+      * an object with ``.decide(states)`` returning list[Decision | dict |
+        tuple | raw_action]
+      * a callable ``fn(states) -> list[...]`` (same item shapes as above)
+
+    Each per-slot return value is normalized through ``_normalize_decision``
+    (the same helper used in single-env adapters), so dict/tuple/raw returns
+    are accepted uniformly.
+    """
+
+    name: str
+    policy: Any
+
+    def decide(self, states: BatchedState) -> list[Decision]:
+        if hasattr(self.policy, "decide"):
+            raw = self.policy.decide(states)
+        elif callable(self.policy):
+            raw = self.policy(states)
+        else:
+            raise TypeError(
+                f"Batched policy {self.name!r} is not callable and has no decide() method."
+            )
+        if not isinstance(raw, list):
+            raise TypeError(
+                f"Batched policy {self.name!r} must return a list, got {type(raw).__name__}."
+            )
+        if len(raw) != len(states):
+            raise ValueError(
+                f"Batched policy {self.name!r} returned {len(raw)} decisions "
+                f"for {len(states)} input states."
+            )
+        return [_normalize_decision(item) for item in raw]
+
+
+def normalize_batched_policy(
+    policy: Any, name: str | None = None
+) -> BatchedPolicyAdapter:
+    if isinstance(policy, BatchedPolicyAdapter):
+        return policy
+    policy_name = name or getattr(policy, "version", None) or getattr(policy, "__name__", None)
+    if not policy_name:
+        policy_name = policy.__class__.__name__
+    return BatchedPolicyAdapter(name=str(policy_name), policy=policy)
+
+
+def from_single(
+    policy: Callable[[State], Any], name: str | None = None
+) -> BatchedPolicyAdapter:
+    """Wrap a single-state policy as a BatchedPolicy by looping over slots.
+
+    Throughput note: each call sequentially evaluates the inner policy N
+    times. Use for backward compat with single-env policy zoos; for
+    production GPU-batched models, write a natively-batched policy instead.
+    """
+    inner_name = name or getattr(policy, "version", None) or getattr(policy, "__name__", None)
+    if not inner_name:
+        inner_name = policy.__class__.__name__
+
+    def _batched(states: BatchedState) -> list[Any]:
+        return [policy(state) for state in states]
+
+    _batched.__name__ = f"batched_{inner_name}"
+    return BatchedPolicyAdapter(name=str(inner_name), policy=_batched)