From b43dc4c1ed5a2550cc028d39158928bdb904b0d4 Mon Sep 17 00:00:00 2001 From: Aryan Pareek Date: Wed, 3 Jun 2026 18:55:28 +0530 Subject: [PATCH 1/6] Add Isaac Lab integration spike under roboeval/integrations/isaac/ Wraps any single-env Isaac Lab environment (gym.make("Isaac-...-v0", num_envs=1)) into the EnvironmentAdapter Protocol so policies can be evaluated against Isaac Lab tasks through EvalRunner. Zero changes to existing SDK files. The adapter handles three Isaac-specific concerns that distinguish it from the existing Gymnasium adapter: 1. Always-batched envs. Isaac envs subclass gym.vector.VectorEnv and return tensors with a leading batch dimension even at num_envs=1. The adapter slices batch_index=0 (configurable) to expose single-episode semantics to the runner. 2. GPU tensors. Observations come back as torch.Tensor on cuda:0. The tensor_to_numpy helper moves any tensor to CPU and converts to numpy before it reaches the runner's JSON writer. 3. Tensor actions on the env device. Policies typically return numpy arrays or Python scalars; the adapter wraps them in a batched torch tensor on the env's device with the expected (num_envs, action_dim) shape. Vectorized eval (num_envs > 1 producing N independent rollouts) is explicitly out of scope for this spike. The adapter accepts num_envs > 1 but only reads batch_index=0; a clear warning fires explaining the parallelism cost. True vector eval requires SDK runner changes documented as Phase 3+ work. Files added: - roboeval/integrations/isaac/__init__.py (public re-exports) - roboeval/integrations/isaac/adapter.py (adapter + 6 default hooks + tensor_to_numpy helper + action shape normalization) - roboeval/integrations/isaac/demo_rollout.py (Isaac-Cartpole-Direct-v0 manual rollout demo) - roboeval/integrations/isaac/README.md (usage, mapping table, cloud GPU setup workflow) - roboeval/integrations/isaac/notes.md (design rationale, 12 known edge cases, real-Isaac validation plan) - roboeval/integrations/isaac/requirements.txt (Isaac Sim + Isaac Lab install pointers; not a single pip line) - roboeval/integrations/isaac/tests/__init__.py - roboeval/integrations/isaac/tests/test_adapter.py (46 unit + integration tests using mock Isaac envs; runs on Mac, no Isaac install needed) What's verified by the test suite: - All 46 tests pass on Mac in 23ms via mocked Isaac envs (real torch tensors, no Isaac dependency) - Adapter integrates cleanly with EvalRunner + Ruleset - info["isaac"] namespace preserved through to StepRecord - Rich env info flows through to_serializable + JSON write - Hook overrides work (outcome, action, info_keys, etc.) - Multi-episode reset behavior - Action shape normalization across 4 input shapes (scalar, 1-D array matching action_dim, 1-D with batch shape, pre-batched 2-D tensor) What's NOT verified (requires cloud GPU pass): - Real Isaac-Cartpole-Direct-v0 smoke test - GPU->CPU tensor sync against actual cuda:0 tensors - env.device / env.sim_device attribute resolution on real Isaac - env.reset(options=...) acceptance across Isaac task variants See notes.md section 12 for the cloud GPU validation plan. Validate: python -m unittest discover -s roboeval/integrations/isaac/tests --- roboeval/integrations/isaac/README.md | 151 ++++ roboeval/integrations/isaac/__init__.py | 49 ++ roboeval/integrations/isaac/adapter.py | 409 +++++++++++ roboeval/integrations/isaac/demo_rollout.py | 100 +++ roboeval/integrations/isaac/notes.md | 162 +++++ roboeval/integrations/isaac/requirements.txt | 33 + roboeval/integrations/isaac/tests/__init__.py | 0 .../integrations/isaac/tests/test_adapter.py | 675 ++++++++++++++++++ 8 files changed, 1579 insertions(+) create mode 100644 roboeval/integrations/isaac/README.md create mode 100644 roboeval/integrations/isaac/__init__.py create mode 100644 roboeval/integrations/isaac/adapter.py create mode 100644 roboeval/integrations/isaac/demo_rollout.py create mode 100644 roboeval/integrations/isaac/notes.md create mode 100644 roboeval/integrations/isaac/requirements.txt create mode 100644 roboeval/integrations/isaac/tests/__init__.py create mode 100644 roboeval/integrations/isaac/tests/test_adapter.py diff --git a/roboeval/integrations/isaac/README.md b/roboeval/integrations/isaac/README.md new file mode 100644 index 0000000..8983a97 --- /dev/null +++ b/roboeval/integrations/isaac/README.md @@ -0,0 +1,151 @@ +# Isaac Lab Integration + +Wrap any single-env [Isaac Lab](https://github.com/isaac-sim/IsaacLab) environment as a roboeval `EnvironmentAdapter` so existing policies can be evaluated against Isaac Lab tasks through `EvalRunner` — no changes to the SDK core. + +## Important — what runs where + +Isaac Sim + Isaac Lab require **NVIDIA GPU + Linux or Windows**. They will not run on macOS. + +The recommended workflow is: + +1. **Write adapter / policy code locally** (Mac is fine; the adapter file imports torch but doesn't import Isaac Lab at the module top) +2. **Run mocked unit tests locally** (this directory's tests use mock Isaac envs, no Isaac install needed) +3. **Validate against real Isaac on a cloud GPU** (RunPod, Lambda Labs, AWS p3, GCP T4, internal Linux workstation) — see "Cloud GPU workflow" below + +## Quick start (once Isaac Lab is installed) + +```python +import gymnasium as gym + +from roboeval import EvalRunner, Ruleset, Scenario, require_metric, require_outcome +from roboeval.integrations.isaac import IsaacEnvironmentAdapter + + +def naive_policy(state): + obs = state.get("policy") or state.get("observation") + pole_angle = float(obs[2]) + return {"action": 0 if pole_angle < 0 else 1} + + +env = gym.make("Isaac-Cartpole-Direct-v0", num_envs=1) +adapter = IsaacEnvironmentAdapter(env=env, name="isaac_cartpole") + +report = EvalRunner( + policies=[naive_policy], + scenarios=[Scenario("smoke", {"seed": 0}, max_steps=200)], + ruleset=Ruleset([ + require_outcome("terminated_success"), + require_metric("episode_return", ">=", 50.0), + ]), + baseline_policy="naive_policy", + environment=adapter, +).run() + +report.save("runs/isaac_smoke") +``` + +For a manual rollout (prints step-by-step output): + +```bash +python -m roboeval.integrations.isaac.demo_rollout +``` + +## How Isaac Lab maps to the SDK + +Isaac Lab envs are gymnasium-compatible but always vectorized. The adapter handles three differences from vanilla gym envs: + +| Concern | What Isaac Lab does | What the adapter does | +|---------|---------------------|------------------------| +| **Batch dimension** | Even `num_envs=1` envs return tensors with shape `(1, ...)` | Slices `batch_index=0` (configurable) to expose single-episode semantics | +| **GPU tensors** | Observations are `torch.Tensor` on `cuda:0` | Coerces to CPU numpy via `tensor_to_numpy()` before they hit the runner's JSON writer | +| **Tensor actions** | `env.step()` expects `torch.Tensor` on the env's device, shape `(num_envs, action_dim)` | Wraps user actions (numpy arrays, scalars, lists) in a batched torch tensor on the right device | + +Plus the standard six translation hooks (matching `GymnasiumEnvironmentAdapter`): + +| Hook | Default behavior | Override when... | +|------|------------------|-------------------| +| `observation_to_state` | If dict → pass through; else wrap as `{"observation": obs}` | User wants to rename/restructure keys | +| `action_from_decision` | Identity | User wants a custom action vocabulary (`"left"` → `0`) | +| `outcome_from_step` | `terminated && reward > 0` → success; else failure / truncated / progress | Env uses `info["is_success"]` or task-specific signal | +| `events_from_step` | Emit `episode_terminated`, `episode_truncated`, `reward_negative` | Add domain-specific tags | +| `seed_from_scenario` | Read `scenario.initial_state["seed"]`, then `metadata["seed"]` | Custom seed routing | +| `options_from_scenario` | Read `scenario.metadata["reset_options"]` | Env supports task-specific reset options | + +## The `StepOutcome` shape produced + +```python +StepOutcome( + next_state={"policy": ndarray([cart_pos, cart_vel, pole_angle, pole_vel])}, + outcome="progress", # or terminated_success / failure / truncated + failure_label="", # populated on failure + terminal=False, # True when terminated or truncated + metrics={ + "reward": 1.0, # this step's reward + "episode_return": 17.0, # running sum since reset + }, + events=["episode_terminated"], # tags for rule filtering + info={ + "isaac": { + "terminated": False, + "truncated": False, + "raw_info": {...}, # whatever info dict the env emits + "batch_index": 0, # which env slice we're reading + } + }, +) +``` + +## What this integration does NOT support (yet) + +- **`num_envs > 1` parallelism.** The SDK runner is single-episode. The adapter accepts `num_envs > 1` but warns and reads only `batch_index=0`. Other envs run but are ignored. For true vector eval, pass `num_envs=1` to `gym.make()`; vector-eval throughput will come when the SDK runner gains batched execution. +- **Render-mode frame capture.** `StepOutcome.artifacts` is empty. A follow-up can add an `artifacts_from_step` hook for `render_mode="rgb_array"` frames. +- **Task-specific success defaults.** Each Isaac task has different success conventions (`info["is_success"]`, terminal reward shaping, etc.). The default `outcome_from_step` is a starting point; override per task. + +## Cloud GPU workflow + +When you don't have a Linux+NVIDIA workstation handy: + +### Option A — RunPod (recommended for spikes) + +1. Create a RunPod account, billing set up. +2. Spin up a pod with an A40 / RTX 4090 / A100 (whatever's cheap; A40 around $0.40/hr). +3. Choose the "Isaac Sim" template if available, or start from `runpod/pytorch:2.x-py3.10-cuda12.1`. +4. SSH in or open the in-browser terminal. +5. Install Isaac Sim: + ```bash + pip install isaacsim==4.5.* --extra-index-url https://pypi.nvidia.com + ``` +6. Install Isaac Lab: + ```bash + git clone https://github.com/isaac-sim/IsaacLab.git + cd IsaacLab && ./isaaclab.sh --install + ``` +7. Install roboeval from your branch: + ```bash + git clone https://github.com/quarqlabs/roboeval.git + cd roboeval && git checkout spike/isaac-integration + pip install -e . + ``` +8. Run the smoke test: + ```bash + python -m roboeval.integrations.isaac.demo_rollout + ``` + +Expect ~5–20 min for the install, then sub-second per-step in sim. Total spike validation cost: a few dollars. + +### Option B — Lambda Labs / AWS / GCP + +Same shape; pick the lowest-cost GPU instance with Linux + CUDA. Persistence of the disk matters if you'll iterate over multiple sessions. + +### Option C — Internal Linux workstation + +If the team has a Linux box with an NVIDIA GPU, that's the most cost-effective dev experience. Just install Isaac Sim + Isaac Lab + roboeval and iterate. + +## Known sharp edges + +- **CartPole `outcome_from_step` defaults are imperfect.** Isaac-Cartpole-Direct gives `reward = +1` every step including the terminal step where the pole falls, so the default classifier reports `terminated_success` even when the pole fell. Score with `require_metric("episode_return", ">=", threshold)` rather than `require_outcome("terminated_success")` for these envs, OR override `outcome_from_step` with the env-specific success detector. (Same sharp edge applies to vanilla Gymnasium CartPole; documented in the Gymnasium integration too.) +- **`env.reset(options=...)` not universally accepted.** Some Isaac envs don't accept the `options` kwarg. The adapter catches the `TypeError` and falls back to `env.reset(seed=...)` only. +- **`env.spec.max_episode_steps` may conflict with `scenario.max_steps`.** Whichever is smaller wins. Treat `scenario.max_steps` as authoritative; if the env's TimeLimit fires first, episodes are reported as `truncated`. +- **GPU memory.** Isaac Sim is heavy. A40 / RTX 4090 / A100 all have enough memory for single-env CartPole; manipulation tasks may need bigger GPUs or `num_envs=1` discipline. + +See `notes.md` for the full design rationale and follow-up items. diff --git a/roboeval/integrations/isaac/__init__.py b/roboeval/integrations/isaac/__init__.py new file mode 100644 index 0000000..0a52f58 --- /dev/null +++ b/roboeval/integrations/isaac/__init__.py @@ -0,0 +1,49 @@ +"""Isaac Lab integration for roboeval. + +Wraps any single-env Isaac Lab environment (``gym.make("Isaac-...-v0", num_envs=1)``) +into roboeval's ``EnvironmentAdapter`` so policies can be evaluated against Isaac +through ``EvalRunner``. + +Quick start:: + + import gymnasium as gym + from roboeval import EvalRunner, Ruleset, Scenario, require_metric + from roboeval.integrations.isaac import IsaacEnvironmentAdapter + + env = gym.make("Isaac-Cartpole-Direct-v0", num_envs=1) + adapter = IsaacEnvironmentAdapter(env=env, name="isaac_cartpole") + + report = EvalRunner( + policies=[my_policy], + scenarios=[Scenario("smoke", {"seed": 0}, max_steps=200)], + ruleset=Ruleset([require_metric("episode_return", ">=", 50.0)]), + baseline_policy="my_policy", + environment=adapter, + ).run() + report.save("runs/isaac_smoke") + +Isaac Lab installation is non-trivial (NVIDIA GPU + Linux/Windows + Isaac Sim ++ Isaac Lab build). See ``README.md`` for the workflow. +""" + +from .adapter import ( + IsaacEnvironmentAdapter, + default_action_from_decision, + default_events_from_step, + default_observation_to_state, + default_options_from_scenario, + default_outcome_from_step, + default_seed_from_scenario, + tensor_to_numpy, +) + +__all__ = [ + "IsaacEnvironmentAdapter", + "default_action_from_decision", + "default_events_from_step", + "default_observation_to_state", + "default_options_from_scenario", + "default_outcome_from_step", + "default_seed_from_scenario", + "tensor_to_numpy", +] diff --git a/roboeval/integrations/isaac/adapter.py b/roboeval/integrations/isaac/adapter.py new file mode 100644 index 0000000..99a0453 --- /dev/null +++ b/roboeval/integrations/isaac/adapter.py @@ -0,0 +1,409 @@ +"""Isaac Lab ↔ roboeval environment adapter. + +Wraps an Isaac Lab single-env ``gymnasium.VectorEnv`` (i.e. constructed with +``num_envs=1``) into roboeval's ``EnvironmentAdapter`` Protocol. + +Isaac Lab vs Gymnasium +---------------------- +Isaac Lab envs are gymnasium-compatible at the API level but differ in three +practical ways the adapter handles: + +1. **They're always vectorized.** Even at ``num_envs=1``, observations, + rewards, and termination flags come back with a batch dimension. This + adapter takes the ``batch_index=0`` slice and exposes a single-episode + view to the runner. + +2. **Tensors live on GPU.** Observations come back as ``torch.Tensor`` with + ``device='cuda:0'``. Reports require JSON-safe types, so the adapter + coerces tensors to CPU numpy before they leave the boundary. + +3. **Actions must be tensors with a batch dim.** Policies typically return + numpy arrays or Python scalars; the adapter converts to torch tensors on + the right device with the batch dimension prepended. + +The six translation hooks match the Gymnasium adapter's surface so users +already familiar with that pattern see the same shape here. +""" + +from __future__ import annotations + +import warnings +from dataclasses import dataclass +from typing import Any, Callable + +from roboeval.core import Action, Scenario, State, to_serializable +from roboeval.environment import StepOutcome + + +# ── Public default hooks ──────────────────────────────────────────────────── + + +def default_observation_to_state(obs: Any) -> dict[str, Any]: + """Wrap a (batch-sliced, CPU-numpy) observation into a State dict. + + Mirrors the Gymnasium adapter: dict observations pass through; non-dict + observations get wrapped as ``{"observation": obs}``. Tensor coercion + happens upstream in ``_slice_and_coerce_obs``, so by the time this hook + runs, values are already JSON-safe. + """ + if isinstance(obs, dict): + return dict(obs) + return {"observation": obs} + + +def default_action_from_decision(action: Action) -> Any: + """Pass through the decision action unchanged. + + The adapter handles the torch.Tensor conversion + batch-dim prepending + in ``_to_batched_torch_action`` — this hook lets users translate from + a custom action vocabulary if they want. + """ + return action + + +def default_outcome_from_step( + reward: float, terminated: bool, truncated: bool, info: dict +) -> tuple[str, str]: + """Default ``(outcome, failure_label)`` mapping. + + Override for env-specific success detection — e.g. many Isaac Lab tasks + expose ``info["success"]`` or ``info["is_success"]`` as a boolean signal, + which is more reliable than reward sign. + """ + if terminated: + if reward > 0: + return ("terminated_success", "") + return ("terminated_failure", "terminated_failure") + if truncated: + return ("truncated", "timeout") + return ("progress", "") + + +def default_events_from_step( + reward: float, terminated: bool, truncated: bool, info: dict +) -> list[str]: + """Emit short event tags useful for rule filters and report highlights.""" + events: list[str] = [] + if terminated: + events.append("episode_terminated") + if truncated: + events.append("episode_truncated") + if reward < 0: + events.append("reward_negative") + return events + + +def default_seed_from_scenario(scenario: Scenario) -> int | None: + """Pull a seed from the scenario, preferring ``initial_state``.""" + seed = scenario.initial_state.get("seed") + if seed is None: + seed = scenario.metadata.get("seed") + return int(seed) if seed is not None else None + + +def default_options_from_scenario(scenario: Scenario) -> dict | None: + """Pull reset options from ``scenario.metadata['reset_options']``.""" + options = scenario.metadata.get("reset_options") + if isinstance(options, dict): + return dict(options) + return None + + +def tensor_to_numpy(value: Any) -> Any: + """Coerce a torch.Tensor (CPU or GPU) to a numpy array. + + Public so user overrides of ``observation_to_state`` can call it. + Returns ``value`` unchanged if it isn't tensor-like. + """ + if hasattr(value, "is_cuda") and getattr(value, "is_cuda", False): + value = value.detach().cpu() + elif hasattr(value, "detach"): + value = value.detach() + if hasattr(value, "numpy"): + return value.numpy() + return value + + +# ── The adapter ───────────────────────────────────────────────────────────── + + +@dataclass +class IsaacEnvironmentAdapter: + """Wraps a single-env Isaac Lab env into roboeval's ``EnvironmentAdapter``. + + Style mirrors ``GymnasiumEnvironmentAdapter`` and the existing + ``CallableEnvironmentAdapter`` pattern from ``roboeval.environment``. + + Parameters + ---------- + env : gymnasium.vector.VectorEnv + An Isaac Lab env constructed with ``num_envs=1``. Anything else is + refused at construction with a clear error message. + name : str + Display name for reports (the runner reads this). + batch_index : int + Which env in the batch to read/write. Defaults to ``0``. The adapter + warns if ``num_envs > 1`` and uses this index. + observation_to_state, action_from_decision, outcome_from_step, + events_from_step, seed_from_scenario, options_from_scenario : + Override hooks. ``None`` selects the corresponding ``default_*`` + function in this module. + info_keys : list[str] | None + Optional allowlist for keys passed through from Isaac's ``info`` + dict into ``StepOutcome.info["isaac"]["raw_info"]``. + """ + + env: Any # gym.vector.VectorEnv + name: str = "isaac_env" + batch_index: int = 0 + + observation_to_state: Callable[[Any], dict[str, Any]] | None = None + action_from_decision: Callable[[Action], Any] | None = None + outcome_from_step: ( + Callable[[float, bool, bool, dict], tuple[str, str]] | None + ) = None + events_from_step: ( + Callable[[float, bool, bool, dict], list[str]] | None + ) = None + seed_from_scenario: Callable[[Scenario], int | None] | None = None + options_from_scenario: Callable[[Scenario], dict | None] | None = None + + info_keys: list[str] | None = None + + def __post_init__(self) -> None: + num_envs = getattr(self.env, "num_envs", None) + if num_envs is not None and num_envs != 1 and self.batch_index >= num_envs: + raise ValueError( + f"batch_index={self.batch_index} out of range for env with " + f"num_envs={num_envs}. Use 0 <= batch_index < num_envs." + ) + if num_envs is not None and num_envs > 1 and self.batch_index == 0: + warnings.warn( + f"Isaac env has num_envs={num_envs} but the SDK runner is " + f"single-episode. Reading batch_index={self.batch_index}; the " + f"other {num_envs - 1} envs run but are ignored. For best " + f"throughput pass num_envs=1.", + stacklevel=2, + ) + + # Wire defaults + if self.observation_to_state is None: + self.observation_to_state = default_observation_to_state + if self.action_from_decision is None: + self.action_from_decision = default_action_from_decision + if self.outcome_from_step is None: + self.outcome_from_step = default_outcome_from_step + if self.events_from_step is None: + self.events_from_step = default_events_from_step + if self.seed_from_scenario is None: + self.seed_from_scenario = default_seed_from_scenario + if self.options_from_scenario is None: + self.options_from_scenario = default_options_from_scenario + + self._episode_return: float = 0.0 + self._device: str | None = None + + # ── EnvironmentAdapter Protocol ──────────────────────────────────────── + + def reset(self, scenario: Scenario) -> State: + seed = self.seed_from_scenario(scenario) + options = self.options_from_scenario(scenario) + try: + obs, _info = self.env.reset(seed=seed, options=options) + except TypeError: + # Some Isaac envs don't accept options kwarg + obs, _info = self.env.reset(seed=seed) + self._episode_return = 0.0 + return self._state_from_obs(obs) + + def step(self, action: Action, scenario: Scenario) -> StepOutcome: + decision_action = self.action_from_decision(action) + gym_action = self._to_batched_torch_action(decision_action) + obs, reward, terminated, truncated, info = self.env.step(gym_action) + + reward_scalar = self._scalar(reward, default=0.0) + terminated_scalar = self._scalar(terminated, default=False, dtype=bool) + truncated_scalar = self._scalar(truncated, default=False, dtype=bool) + self._episode_return += float(reward_scalar) + + return self._build_outcome( + obs, + float(reward_scalar), + bool(terminated_scalar), + bool(truncated_scalar), + info, + ) + + def close(self) -> None: + """Forward close to the wrapped env.""" + close = getattr(self.env, "close", None) + if callable(close): + close() + + # ── Internals ────────────────────────────────────────────────────────── + + def _state_from_obs(self, obs: Any) -> State: + """Slice the batch dim and coerce tensors to numpy, then apply hook.""" + sliced = self._slice_and_coerce_obs(obs) + return self.observation_to_state(sliced) + + def _slice_and_coerce_obs(self, obs: Any) -> Any: + """Take batch_index slice and convert tensors to numpy. + + Isaac obs is typically either: + - a dict like {"policy": tensor of shape (N, obs_dim)}, OR + - a tensor of shape (N, obs_dim) directly. + Both shapes are handled. + """ + if isinstance(obs, dict): + return {k: tensor_to_numpy(self._index(v)) for k, v in obs.items()} + return tensor_to_numpy(self._index(obs)) + + def _index(self, value: Any) -> Any: + """Return ``value[batch_index]`` if it supports indexing, else value.""" + try: + return value[self.batch_index] + except (TypeError, IndexError, KeyError): + return value + + def _scalar(self, value: Any, default: Any, dtype: type | None = None) -> Any: + """Reduce a batched tensor/array to a Python scalar at batch_index.""" + sliced = self._index(value) + if hasattr(sliced, "item"): + try: + return sliced.item() + except (ValueError, RuntimeError): + pass + if hasattr(sliced, "numpy"): + arr = tensor_to_numpy(sliced) + if hasattr(arr, "item"): + try: + return arr.item() + except (ValueError, RuntimeError): + pass + if dtype is bool: + try: + return bool(sliced) + except (TypeError, ValueError): + return default + try: + return float(sliced) + except (TypeError, ValueError): + return default + + def _filter_info(self, info: dict) -> dict: + if self.info_keys is None: + return dict(info) if isinstance(info, dict) else {} + if not isinstance(info, dict): + return {} + return {key: info[key] for key in self.info_keys if key in info} + + def _to_batched_torch_action(self, action: Any) -> Any: + """Convert action to torch.Tensor on the env's device with a batch dim. + + Accepts: torch.Tensor (any shape), numpy array, Python scalar or list. + Returns a torch tensor with shape ``(num_envs, *action_dim)``. Device + is inferred from the env if possible; falls back to ``cuda`` if + available, else ``cpu``. + """ + try: + import torch + except ImportError as exc: + raise ImportError( + "PyTorch is required for the Isaac adapter. Install via " + "`pip install torch` or as part of your Isaac Lab install." + ) from exc + + if not isinstance(action, torch.Tensor): + action = torch.as_tensor(action) + + # Bring to the right device + device = self._resolve_device(action) + if action.device != torch.device(device): + action = action.to(device) + + # Ensure batch dim matches env.num_envs + num_envs = getattr(self.env, "num_envs", 1) + if action.ndim == 0: + # scalar -> (num_envs, 1) + action = action.view(1, 1).expand(num_envs, 1) + elif action.ndim == 1: + # (action_dim,) -> (num_envs, action_dim) if batch_index == 0 and + # action_dim matches; otherwise treat first dim as batch + action_space = getattr(self.env, "single_action_space", None) + expected_dim = ( + action_space.shape[0] + if action_space is not None and getattr(action_space, "shape", None) + else None + ) + if expected_dim is not None and action.shape[0] == expected_dim: + action = action.unsqueeze(0).expand(num_envs, -1) + elif action.shape[0] == num_envs: + # Already batched as 1-D — add singleton action dim + action = action.unsqueeze(-1) + else: + action = action.unsqueeze(0).expand(num_envs, -1) + # ndim >= 2 we assume already batched correctly + + return action + + def _resolve_device(self, fallback_tensor: Any) -> str: + if self._device is not None: + return self._device + + # Try to read the env's device attribute + env_device = getattr(self.env, "device", None) or getattr( + self.env, "sim_device", None + ) + if env_device is not None: + self._device = str(env_device) + return self._device + + # Fall back to detecting from existing tensor or cuda availability + try: + import torch + except ImportError: + self._device = "cpu" + return self._device + + if hasattr(fallback_tensor, "device"): + self._device = str(fallback_tensor.device) + return self._device + + self._device = "cuda" if torch.cuda.is_available() else "cpu" + return self._device + + def _build_outcome( + self, + obs: Any, + reward: float, + terminated: bool, + truncated: bool, + info: dict, + ) -> StepOutcome: + next_state = self._state_from_obs(obs) + outcome, failure_label = self.outcome_from_step( + reward, terminated, truncated, info + ) + events = self.events_from_step(reward, terminated, truncated, info) + filtered_info = self._filter_info(info) + + return StepOutcome( + next_state=next_state, + outcome=outcome, + failure_label=failure_label, + terminal=bool(terminated or truncated), + metrics={ + "reward": reward, + "episode_return": float(self._episode_return), + }, + events=events, + info={ + "isaac": { + "terminated": terminated, + "truncated": truncated, + "raw_info": to_serializable(filtered_info), + "batch_index": self.batch_index, + } + }, + ) diff --git a/roboeval/integrations/isaac/demo_rollout.py b/roboeval/integrations/isaac/demo_rollout.py new file mode 100644 index 0000000..6931d96 --- /dev/null +++ b/roboeval/integrations/isaac/demo_rollout.py @@ -0,0 +1,100 @@ +"""Manual rollout demo for the Isaac Lab integration. + +Runs Isaac-Cartpole-Direct-v0 through ``IsaacEnvironmentAdapter`` and prints +the resulting ``StepOutcome`` for each step. Demonstrates: + +- Wrapping a single-env Isaac Lab environment (``num_envs=1``) +- GPU tensor → numpy state coercion happening automatically +- Action passthrough from a numpy/scalar policy to a batched torch tensor +- ``info["isaac"]`` namespace with terminated / truncated / raw_info + +Requires Isaac Sim + Isaac Lab installed (NVIDIA GPU + Linux/Windows). Will +not run on Mac. See ``README.md`` for the recommended cloud-GPU workflow. + +Run:: + + python -m roboeval.integrations.isaac.demo_rollout +""" + +from __future__ import annotations + +import sys + + +def naive_policy(state: dict) -> int: + """Push the cart in the direction the pole is leaning. + + Isaac Cartpole observation typically exposes + ``state["policy"]`` or ``state["observation"]`` as a 4-vector + ``[cart_position, cart_velocity, pole_angle, pole_velocity]``. + """ + obs = state.get("policy") + if obs is None: + obs = state.get("observation") + if obs is None: + # Unknown observation shape; default to 0 + return 0 + pole_angle = float(obs[2]) if len(obs) > 2 else 0.0 + return 0 if pole_angle < 0 else 1 + + +def main() -> None: + try: + import gymnasium as gym + except ImportError: + print( + "ERROR: gymnasium is required. Install via `pip install gymnasium>=0.29`.", + file=sys.stderr, + ) + sys.exit(1) + + try: + import isaaclab # noqa: F401 - just to verify Isaac Lab is installed + # NOTE: Some Isaac Lab installs require additional setup (running + # ``./isaaclab.sh -p ...``); see the README for details. + except ImportError: + print( + "ERROR: Isaac Lab is not installed. This demo requires a working " + "Isaac Sim + Isaac Lab install. See README.md for setup.", + file=sys.stderr, + ) + sys.exit(1) + + from roboeval.core import Scenario + from roboeval.integrations.isaac import IsaacEnvironmentAdapter + + env_id = "Isaac-Cartpole-Direct-v0" + print(f"Creating env: {env_id} (num_envs=1)") + env = gym.make(env_id, num_envs=1) + + adapter = IsaacEnvironmentAdapter(env=env, name="isaac_cartpole") + + scenario = Scenario( + name="isaac_cartpole_smoke", + initial_state={"seed": 0}, + max_steps=200, + ) + + state = adapter.reset(scenario) + print(f"[reset] state keys={list(state.keys())}") + + for step in range(scenario.max_steps): + action = naive_policy(state) + outcome = adapter.step(action, scenario) + print( + f"[step {step:3d}] action={action} " + f"outcome={outcome.outcome:<22} terminal={outcome.terminal} " + f"reward={outcome.metrics['reward']:+.2f} " + f"return={outcome.metrics['episode_return']:.2f} " + f"events={outcome.events}" + ) + state = outcome.next_state + if outcome.terminal: + break + + print(f"\n✓ Demo complete. Final return: {adapter._episode_return:.2f}") + adapter.close() + + +if __name__ == "__main__": + main() diff --git a/roboeval/integrations/isaac/notes.md b/roboeval/integrations/isaac/notes.md new file mode 100644 index 0000000..df50c21 --- /dev/null +++ b/roboeval/integrations/isaac/notes.md @@ -0,0 +1,162 @@ +# Design Notes — Isaac Lab Integration Spike + +Engineering notes captured while implementing this spike. Audience: anyone extending this adapter or templating the next integration (MuJoCo direct, PyBullet, ROS 2, real robots). + +--- + +## 1. Why a separate adapter from `GymnasiumEnvironmentAdapter`? + +Isaac Lab envs ARE gymnasium-compatible at the API level (they subclass `gym.vector.VectorEnv`). It would be possible to extend `GymnasiumEnvironmentAdapter` and override hooks. We chose a separate class for three reasons: + +1. **Different fundamental assumptions.** Gymnasium envs are usually single-env; Isaac envs are always vectorized. Embedding the batch-slicing logic in `GymnasiumEnvironmentAdapter` would either pollute it (`if isinstance(env, VectorEnv): ...` branches) or force a complicated inheritance hierarchy. + +2. **Different tensor lifecycle.** The Gymnasium adapter assumes observations are `numpy.ndarray` or JSON-safe types. The Isaac adapter handles PyTorch tensors on GPU explicitly, including device routing for actions. + +3. **Different failure modes for the user.** The Gymnasium adapter refuses `VectorEnv` outright. The Isaac adapter must accept it (because all Isaac envs are vector envs). A user error like "wrong tensor dtype" produces different debugging guidance. + +A future refactor could share helpers between the two via a base class once a third integration (MuJoCo via dm_control, say) reveals which patterns actually generalize. Premature today. + +## 2. The batch-index trick + +Isaac Lab envs return tensors with a batch dimension even at `num_envs=1`: + +```python +obs.shape # (1, obs_dim) +reward.shape # (1,) +terminated.shape # (1,) +``` + +The adapter takes `batch_index=0` slices throughout. The runner sees scalar reward, single-vector observations, single boolean done — matching the single-episode runner's expectations. + +Why expose `batch_index` as a constructor parameter rather than always using 0? Two reasons: + +- **Defensive programming.** If a user passes `num_envs > 1` by accident, they can at least pick *which* env they care about rather than getting a silent crash. +- **Forward compatibility.** When the SDK runner gains batched execution, we can ship a `BatchIsaacEnvironmentAdapter` that emits N rollouts in parallel. The `batch_index` parameter becomes a list of indices to track. + +For v1, the standard pattern is `num_envs=1` + `batch_index=0` + a warning if `num_envs > 1`. + +## 3. Tensor-to-numpy coercion happens AT THE ADAPTER BOUNDARY + +A design principle: the runner should never see tensors. The adapter does the coercion so: + +- Reports (`decision_logs.jsonl`, `episode_results.json`) are JSON-safe without depending on `to_serializable` knowing about torch +- The policy's `decide(state)` sees plain numpy / Python types, not GPU tensors +- The runner's `dict(state)` copy works without surprise + +The cost is one GPU→CPU sync per step. For Isaac Sim at single-env throughput (~30–60 Hz for Cartpole), this is negligible. For real-time control loops at 1000+ Hz this would matter — but real-time control is out of scope for the runner. + +There's a public `tensor_to_numpy()` function in `adapter.py` so users overriding `observation_to_state` can call it consistently. + +## 4. Action conversion handles three input shapes + +Policies might return: +- Python scalars (`0`, `1.5`) +- 1-D numpy arrays / lists (`[0.5, -0.3]`) +- 2-D arrays already batched (`[[0.5, -0.3]]`) + +The adapter's `_to_batched_torch_action` normalizes all three to the expected `(num_envs, action_dim)` torch tensor on the env's device. Device routing is conservative: read `env.device`, fall back to `env.sim_device`, fall back to the action's existing device, fall back to `cuda` if available, else `cpu`. + +We expose a hook (`action_from_decision`) for users to customize the user-action → policy-action translation (e.g. discrete vocabulary `"left"` → `0`). The torch conversion happens after the hook. + +## 5. The `info["isaac"]` namespace mirrors the Gymnasium adapter's `info["gymnasium"]` pattern + +Same shape: + +```python +"info": { + "isaac": { + "terminated": bool, + "truncated": bool, + "raw_info": {...}, # filtered via info_keys allowlist + "batch_index": int, + } +} +``` + +This consistency means downstream consumers (the dashboard, future agents) can find sim-specific raw data at a predictable location. The `batch_index` field is Isaac-specific; it's useful when debugging multi-env scenarios. + +## 6. What we deliberately did NOT do + +- **Vector eval support.** Out of scope for v1. The SDK runner is single-episode. When that changes, a `BatchIsaacEnvironmentAdapter` ships as a separate class. +- **`info["is_success"]` auto-detection.** Many Isaac tasks expose success this way (Franka Cabinet, Lift Cube, etc.). Tempting to default to it, but every task has different keys (`success`, `is_success`, `task_success`). We'd ship false positives for tasks that use those keys with different semantics. Instead, document the override pattern. +- **Per-task hook presets.** A `presets/franka_cabinet.py` module with the right `outcome_from_step` for that task would be useful. Out of scope until we know which tasks users actually want. +- **Render-mode frame capture.** `StepOutcome.artifacts` left empty. Should be an `artifacts_from_step` hook in v2. +- **Determinism guarantees.** Isaac Sim has its own determinism story (depends on physics version, seed, GPU model). We pass through the seed; we don't promise byte-exact replay. + +## 7. Why `num_envs > 1` warns instead of refusing + +Earlier draft refused `num_envs > 1` outright. We softened to a warning because: + +- A user debugging a vector-env-only Isaac task ("Why is my parallel eval broken?") wants the adapter to do *something*, not crash +- The Gymnasium adapter already refuses `gym.vector.VectorEnv` outright; the Isaac case needs different behavior (since Isaac envs are always VectorEnvs) +- Reading `batch_index=0` of an N-env batch is correct behavior, just throughput-wasteful + +The warning is loud enough that users notice ("Other 99 envs run but are ignored") and the docs explain the right fix (`num_envs=1`). + +## 8. Why `env.reset(options=...)` is wrapped in try/except + +Some Isaac envs accept the `options` kwarg per the standard gymnasium API; some pre-date that API change. To support both: + +```python +try: + obs, info = self.env.reset(seed=seed, options=options) +except TypeError: + obs, info = self.env.reset(seed=seed) +``` + +This is mildly defensive. The alternative (sniffing the signature) is brittle. The cost (one extra try/except per reset) is irrelevant since reset is once per scenario. + +## 9. Why mocked tests instead of real-Isaac tests + +Two reasons: + +1. **No Isaac on Mac.** Running real Isaac requires GPU + Linux. Tests should be runnable on the dev machine (which is a Mac). +2. **CI cost.** Running real Isaac in CI is impractical for an open-source SDK — requires GPU runners, which are paid. + +The mock-based tests validate: +- Tensor coercion (using real torch tensors, not Isaac) +- Batch slicing +- Action shape normalization +- Hook wiring +- `StepOutcome` field population + +They DON'T validate: +- Whether the action gets sent to Isaac's physics engine correctly +- Whether observations from Isaac actually arrive in the expected shape +- Sim-to-real behavior + +For those, the demo_rollout.py on real Isaac is the authoritative test. + +## 10. CUDA version pinning is the user's problem + +Isaac Sim 4.5 wants CUDA 11.8 or 12.x. Our adapter doesn't care which CUDA — we just use whatever torch is installed. If the user mismatches torch's CUDA with Isaac's expected CUDA, the env construction will fail before our adapter even runs. + +If you hit "RuntimeError: CUDA error: invalid device function" — that's not us; it's a CUDA/torch version mismatch. Use the install commands from `requirements.txt` exactly. + +## 11. Follow-up items for the team + +1. **Add `[isaac]` extras to `pyproject.toml`.** Won't be a one-liner (Isaac Sim isn't pip-installable in a single line), but `gymnasium>=0.29` and `torch>=2.0` can be added so `pip install roboeval[isaac]` at least gets the Python dependencies right. + +2. **Ship a `IsaacBatchEnvironmentAdapter` when the runner gains batched execution.** Reuse most of this code; emit a list of `StepOutcome` per step instead of one. + +3. **Per-task success preset modules.** `roboeval/integrations/isaac/presets/franka_cabinet.py` with pre-configured `outcome_from_step` for common Isaac Lab tasks. Document the convention so users contribute their own. + +4. **Render-mode frame capture as `artifacts`.** An `artifacts_from_step` hook that calls `env.render()` if `render_mode="rgb_array"` and stuffs the frame under `artifacts["frame"]`. Useful for failure debugging. + +5. **`env.spec.id` propagation.** Capture `env.spec.id` in the StepOutcome metadata so reports show "ran against Isaac-Cartpole-Direct-v0" without the user doing it manually. + +6. **Determinism docs.** Write a short note about Isaac's determinism story (where it's reliable, where it's not) so users know when replay-from-seed is trustworthy. + +## 12. Validation steps before merging + +This spike was written on Mac without real Isaac. Before this lands in main: + +1. **Real-Isaac smoke test on a cloud GPU box.** Run `python -m roboeval.integrations.isaac.demo_rollout`. Expect 30–200 step logs ending in `terminal=True`. + +2. **Real-Isaac EvalRunner integration test.** Wrap the adapter in `EvalRunner` with a `Ruleset` and confirm a report is produced. + +3. **Real-Isaac edge cases.** Verify `env.reset(options=...)` works on at least Isaac-Cartpole-Direct-v0 and Isaac-Lift-Cube-Franka-v0. Try `num_envs=4` to confirm the warning fires and batch_index=0 is used. + +4. **Existing SDK tests still pass on the dev machine.** This adapter must not touch core SDK files. `cd roboeval && python -m unittest discover -s tests` after `touch tests/__init__.py`. + +Document each step's results in this notes file so future contributors know what's been validated and what hasn't. diff --git a/roboeval/integrations/isaac/requirements.txt b/roboeval/integrations/isaac/requirements.txt new file mode 100644 index 0000000..50f79b2 --- /dev/null +++ b/roboeval/integrations/isaac/requirements.txt @@ -0,0 +1,33 @@ +# Isaac Lab integration dependencies. +# +# This file is informational. Isaac Sim and Isaac Lab are NOT pip-installable +# in a single line — they require NVIDIA GPU + Linux/Windows + Isaac Sim +# install + Isaac Lab build. See README.md for the workflow. +# +# What our adapter directly requires from Python land: + +gymnasium>=0.29 +torch>=2.0 + +# Isaac Lab itself follows NVIDIA's install path: +# https://isaac-sim.github.io/IsaacLab/main/source/setup/installation/index.html +# +# Typical sequence on a cloud GPU box (RunPod / Lambda Labs / etc.): +# +# 1. Install Isaac Sim (NVIDIA's pip-installer path, fastest): +# pip install isaacsim==4.5.* --extra-index-url https://pypi.nvidia.com +# +# 2. Clone + install Isaac Lab: +# git clone https://github.com/isaac-sim/IsaacLab.git +# cd IsaacLab +# ./isaaclab.sh --install +# +# 3. Verify: +# ./isaaclab.sh -p scripts/environments/list_envs.py +# +# 4. Install roboeval and the Isaac integration extras: +# pip install -e +# # (no extra pip step for the Isaac integration; the adapter is in-tree) +# +# CUDA: Isaac Sim 4.5 currently expects CUDA 11.8 or 12.x. Pin per the +# Isaac Lab release notes for the version you install. diff --git a/roboeval/integrations/isaac/tests/__init__.py b/roboeval/integrations/isaac/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/roboeval/integrations/isaac/tests/test_adapter.py b/roboeval/integrations/isaac/tests/test_adapter.py new file mode 100644 index 0000000..80ba5ea --- /dev/null +++ b/roboeval/integrations/isaac/tests/test_adapter.py @@ -0,0 +1,675 @@ +"""Unit tests for the Isaac Lab integration spike. + +These tests run on Mac without Isaac installed. They use a mock Isaac env +(`MockIsaacEnv`) that fakes the Isaac shape: gymnasium-VectorEnv-style API +returning batched torch tensors on whatever device torch defaults to. + +Run with:: + + python -m unittest discover -s roboeval/integrations/isaac/tests +""" + +from __future__ import annotations + +import json +import unittest +import warnings +from typing import Any + +import numpy as np + +try: + import torch +except ImportError: # pragma: no cover + torch = None # type: ignore[assignment] + +from roboeval import EvalRunner, Ruleset, Scenario, require_metric +from roboeval.core import to_serializable +from roboeval.environment import StepOutcome +from roboeval.integrations.isaac import ( + IsaacEnvironmentAdapter, + default_action_from_decision, + default_events_from_step, + default_observation_to_state, + default_options_from_scenario, + default_outcome_from_step, + default_seed_from_scenario, + tensor_to_numpy, +) + + +HAS_TORCH = torch is not None + + +# ───────────────────────────────────────────────────────────────────────── +# Mock Isaac environment — fakes the Isaac shape so we can test without +# a real Isaac Lab install. +# ───────────────────────────────────────────────────────────────────────── + + +class MockIsaacEnv: + """Minimal mock that mimics an Isaac Lab single-env env. + + Returns batched torch tensors from reset() / step(). Step count drives a + simple termination rule. Use the `obs_as_dict` flag to switch between + dict-of-tensors and bare-tensor observation styles. + """ + + def __init__( + self, + num_envs: int = 1, + obs_dim: int = 4, + action_dim: int = 1, + terminate_at_step: int = 10, + obs_as_dict: bool = True, + device: str | None = None, + ): + if not HAS_TORCH: + raise RuntimeError("MockIsaacEnv requires torch") + self.num_envs = num_envs + self._obs_dim = obs_dim + self._action_dim = action_dim + self._terminate_at_step = terminate_at_step + self._obs_as_dict = obs_as_dict + self.device = device or "cpu" + self._step = 0 + self._last_action: Any = None + self.spec = None + # gym-style spaces (single_action_space matters for action shaping) + import gymnasium as gym + self.single_action_space = gym.spaces.Box( + low=-1.0, high=1.0, shape=(action_dim,) + ) + + def _make_obs(self) -> Any: + # Float32 batched tensor with deterministic values + flat = torch.arange( + self.num_envs * self._obs_dim, dtype=torch.float32 + ).view(self.num_envs, self._obs_dim) + if self._obs_as_dict: + return {"policy": flat.to(self.device)} + return flat.to(self.device) + + def reset(self, seed: int | None = None, options: dict | None = None): + # Seed is honored deterministically just by the obs being a function of step + self._step = 0 + return self._make_obs(), {} + + def step(self, action): + self._step += 1 + self._last_action = action + obs = self._make_obs() + # Each step reward = +1 to mimic CartPole survival reward + reward = torch.ones(self.num_envs, dtype=torch.float32, device=self.device) + terminated = torch.tensor( + [self._step >= self._terminate_at_step] * self.num_envs, + dtype=torch.bool, + device=self.device, + ) + truncated = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + info: dict[str, Any] = {} + return obs, reward, terminated, truncated, info + + def close(self): + self._step = 0 + + +class MockIsaacEnvWithOptionsRefusal(MockIsaacEnv): + """Mock that refuses the `options=` kwarg, like older Isaac envs.""" + + def reset(self, seed: int | None = None): + return super().reset(seed=seed) + + +# ───────────────────────────────────────────────────────────────────────── +# Default hook unit tests (no env needed) +# ───────────────────────────────────────────────────────────────────────── + + +class DefaultObservationToStateTest(unittest.TestCase): + def test_dict_passes_through(self): + obs = {"policy": np.array([1.0, 2.0])} + state = default_observation_to_state(obs) + self.assertEqual(set(state.keys()), {"policy"}) + np.testing.assert_array_equal(state["policy"], obs["policy"]) + + def test_non_dict_wraps_under_observation(self): + obs = np.array([0.1, -0.5, 0.0, 0.3]) + state = default_observation_to_state(obs) + self.assertEqual(list(state.keys()), ["observation"]) + np.testing.assert_array_equal(state["observation"], obs) + + +class DefaultActionFromDecisionTest(unittest.TestCase): + def test_passthrough_int(self): + self.assertEqual(default_action_from_decision(1), 1) + + def test_passthrough_array(self): + a = np.array([0.5, -0.5]) + np.testing.assert_array_equal(default_action_from_decision(a), a) + + +class DefaultOutcomeFromStepTest(unittest.TestCase): + def test_terminated_positive_reward_is_success(self): + outcome, label = default_outcome_from_step(1.0, True, False, {}) + self.assertEqual((outcome, label), ("terminated_success", "")) + + def test_terminated_zero_reward_is_failure(self): + outcome, label = default_outcome_from_step(0.0, True, False, {}) + self.assertEqual((outcome, label), ("terminated_failure", "terminated_failure")) + + def test_truncated_is_timeout(self): + outcome, label = default_outcome_from_step(0.5, False, True, {}) + self.assertEqual((outcome, label), ("truncated", "timeout")) + + def test_non_terminal_is_progress(self): + outcome, label = default_outcome_from_step(0.1, False, False, {}) + self.assertEqual((outcome, label), ("progress", "")) + + +class DefaultEventsFromStepTest(unittest.TestCase): + def test_progress_emits_no_events(self): + self.assertEqual(default_events_from_step(1.0, False, False, {}), []) + + def test_terminated_event(self): + self.assertEqual( + default_events_from_step(1.0, True, False, {}), ["episode_terminated"] + ) + + def test_truncated_event(self): + self.assertEqual( + default_events_from_step(1.0, False, True, {}), ["episode_truncated"] + ) + + def test_negative_reward_event(self): + self.assertEqual( + default_events_from_step(-1.0, False, False, {}), ["reward_negative"] + ) + + +class DefaultSeedFromScenarioTest(unittest.TestCase): + def test_from_initial_state(self): + scenario = Scenario("s", {"seed": 42}, max_steps=10) + self.assertEqual(default_seed_from_scenario(scenario), 42) + + def test_from_metadata(self): + scenario = Scenario("s", {"foo": 1}, max_steps=10, metadata={"seed": 7}) + self.assertEqual(default_seed_from_scenario(scenario), 7) + + def test_none_when_missing(self): + scenario = Scenario("s", {"foo": 1}, max_steps=10) + self.assertIsNone(default_seed_from_scenario(scenario)) + + +class DefaultOptionsFromScenarioTest(unittest.TestCase): + def test_options_dict(self): + scenario = Scenario( + "s", {"foo": 1}, max_steps=10, metadata={"reset_options": {"a": 1}} + ) + self.assertEqual(default_options_from_scenario(scenario), {"a": 1}) + + def test_none_when_missing(self): + scenario = Scenario("s", {"foo": 1}, max_steps=10) + self.assertIsNone(default_options_from_scenario(scenario)) + + +# ───────────────────────────────────────────────────────────────────────── +# tensor_to_numpy +# ───────────────────────────────────────────────────────────────────────── + + +@unittest.skipUnless(HAS_TORCH, "torch required") +class TensorToNumpyTest(unittest.TestCase): + def test_cpu_tensor_to_numpy(self): + t = torch.tensor([1.0, 2.0, 3.0]) + arr = tensor_to_numpy(t) + self.assertIsInstance(arr, np.ndarray) + np.testing.assert_array_almost_equal(arr, [1.0, 2.0, 3.0]) + + def test_non_tensor_passes_through(self): + self.assertEqual(tensor_to_numpy(42), 42) + self.assertEqual(tensor_to_numpy("hello"), "hello") + np.testing.assert_array_equal( + tensor_to_numpy(np.array([1, 2, 3])), np.array([1, 2, 3]) + ) + + +# ───────────────────────────────────────────────────────────────────────── +# Adapter construction +# ───────────────────────────────────────────────────────────────────────── + + +@unittest.skipUnless(HAS_TORCH, "torch required") +class IsaacEnvironmentAdapterConstructionTest(unittest.TestCase): + def test_satisfies_protocol_duck_typing(self): + env = MockIsaacEnv(num_envs=1) + adapter = IsaacEnvironmentAdapter(env=env) + self.assertTrue(hasattr(adapter, "reset")) + self.assertTrue(hasattr(adapter, "step")) + self.assertTrue(callable(adapter.reset)) + self.assertTrue(callable(adapter.step)) + + def test_defaults_wired_when_hooks_none(self): + env = MockIsaacEnv(num_envs=1) + adapter = IsaacEnvironmentAdapter(env=env) + self.assertIs(adapter.observation_to_state, default_observation_to_state) + self.assertIs(adapter.action_from_decision, default_action_from_decision) + self.assertIs(adapter.outcome_from_step, default_outcome_from_step) + self.assertIs(adapter.events_from_step, default_events_from_step) + self.assertIs(adapter.seed_from_scenario, default_seed_from_scenario) + self.assertIs(adapter.options_from_scenario, default_options_from_scenario) + + def test_num_envs_gt_1_warns(self): + env = MockIsaacEnv(num_envs=4) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + IsaacEnvironmentAdapter(env=env) + self.assertTrue( + any("num_envs=4" in str(w.message) for w in caught), + "Expected num_envs > 1 warning to fire", + ) + + def test_invalid_batch_index_raises(self): + env = MockIsaacEnv(num_envs=2) + with self.assertRaises(ValueError): + IsaacEnvironmentAdapter(env=env, batch_index=5) + + def test_name_propagates(self): + env = MockIsaacEnv(num_envs=1) + adapter = IsaacEnvironmentAdapter(env=env, name="my_isaac_cartpole") + self.assertEqual(adapter.name, "my_isaac_cartpole") + + +# ───────────────────────────────────────────────────────────────────────── +# Adapter behavior — reset/step with the mock +# ───────────────────────────────────────────────────────────────────────── + + +@unittest.skipUnless(HAS_TORCH, "torch required") +class IsaacEnvironmentAdapterBehaviorTest(unittest.TestCase): + def setUp(self) -> None: + self.env = MockIsaacEnv(num_envs=1, obs_dim=4, terminate_at_step=10) + self.adapter = IsaacEnvironmentAdapter(env=self.env, name="isaac_mock") + self.scenario = Scenario("test", {"seed": 0}, max_steps=50) + + def test_reset_returns_dict_with_policy_key(self): + state = self.adapter.reset(self.scenario) + self.assertIsInstance(state, dict) + self.assertIn("policy", state) + # batch dim was sliced, so it's a 1-D array of obs_dim + self.assertEqual(state["policy"].shape, (4,)) + + def test_reset_resets_episode_return(self): + self.adapter.reset(self.scenario) + self.adapter.step(0, self.scenario) + self.adapter.step(0, self.scenario) + self.adapter.reset(self.scenario) + outcome = self.adapter.step(0, self.scenario) + self.assertEqual(outcome.metrics["episode_return"], 1.0) + + def test_step_returns_step_outcome(self): + self.adapter.reset(self.scenario) + outcome = self.adapter.step(0, self.scenario) + self.assertIsInstance(outcome, StepOutcome) + # Required fields + self.assertIsInstance(outcome.next_state, dict) + self.assertIsInstance(outcome.outcome, str) + self.assertIsInstance(outcome.failure_label, str) + self.assertIsInstance(outcome.terminal, bool) + # Optional fields populated + self.assertIsInstance(outcome.metrics, dict) + self.assertIn("reward", outcome.metrics) + self.assertIn("episode_return", outcome.metrics) + self.assertIsInstance(outcome.events, list) + self.assertIsInstance(outcome.info, dict) + self.assertIn("isaac", outcome.info) + + def test_metrics_per_spec(self): + self.adapter.reset(self.scenario) + out1 = self.adapter.step(0, self.scenario) + out2 = self.adapter.step(0, self.scenario) + out3 = self.adapter.step(0, self.scenario) + self.assertEqual(out1.metrics["reward"], 1.0) + self.assertEqual(out2.metrics["reward"], 1.0) + self.assertEqual(out3.metrics["episode_return"], 3.0) + + def test_info_namespace(self): + self.adapter.reset(self.scenario) + outcome = self.adapter.step(0, self.scenario) + isaac_info = outcome.info["isaac"] + self.assertIn("terminated", isaac_info) + self.assertIn("truncated", isaac_info) + self.assertIn("raw_info", isaac_info) + self.assertIn("batch_index", isaac_info) + self.assertEqual(isaac_info["batch_index"], 0) + + def test_terminal_fires_at_expected_step(self): + self.adapter.reset(self.scenario) + outcomes = [] + for _ in range(self.env._terminate_at_step + 2): + outcomes.append(self.adapter.step(0, self.scenario)) + if outcomes[-1].terminal: + break + self.assertTrue(outcomes[-1].terminal) + self.assertIn("episode_terminated", outcomes[-1].events) + # episode_return after termination + self.assertEqual( + outcomes[-1].metrics["episode_return"], float(self.env._terminate_at_step) + ) + + def test_bare_tensor_obs_gets_wrapped(self): + env = MockIsaacEnv(num_envs=1, obs_as_dict=False) + adapter = IsaacEnvironmentAdapter(env=env) + state = adapter.reset(Scenario("s", {"seed": 0}, max_steps=10)) + self.assertIn("observation", state) + self.assertEqual(state["observation"].shape, (4,)) + + def test_options_refusal_is_handled(self): + env = MockIsaacEnvWithOptionsRefusal(num_envs=1) + adapter = IsaacEnvironmentAdapter(env=env) + scenario = Scenario( + "s", {"seed": 1}, max_steps=10, metadata={"reset_options": {"foo": 1}} + ) + # Should not raise even though env refuses options + state = adapter.reset(scenario) + self.assertIn("policy", state) + + +# ───────────────────────────────────────────────────────────────────────── +# Hook overrides +# ───────────────────────────────────────────────────────────────────────── + + +@unittest.skipUnless(HAS_TORCH, "torch required") +class IsaacEnvironmentAdapterHookOverrideTest(unittest.TestCase): + def test_outcome_hook_override(self): + env = MockIsaacEnv(num_envs=1, terminate_at_step=5) + + def always_goal(reward, terminated, truncated, info): + return ("goal_reached", "") + + adapter = IsaacEnvironmentAdapter(env=env, outcome_from_step=always_goal) + adapter.reset(Scenario("s", {"seed": 0}, max_steps=10)) + outcome = adapter.step(0, Scenario("s", {"seed": 0}, max_steps=10)) + self.assertEqual(outcome.outcome, "goal_reached") + + def test_action_hook_translates_vocabulary(self): + env = MockIsaacEnv(num_envs=1, terminate_at_step=10) + vocab = {"push_left": 0, "push_right": 1} + + adapter = IsaacEnvironmentAdapter( + env=env, action_from_decision=lambda a: vocab[a] if isinstance(a, str) else a + ) + adapter.reset(Scenario("s", {"seed": 0}, max_steps=10)) + outcome = adapter.step("push_left", Scenario("s", {"seed": 0}, max_steps=10)) + # No crash means the action got translated and accepted + self.assertFalse(outcome.terminal) + + def test_info_keys_allowlist(self): + env = MockIsaacEnv(num_envs=1) + adapter = IsaacEnvironmentAdapter(env=env, info_keys=["only_this_key"]) + adapter.reset(Scenario("s", {"seed": 0}, max_steps=10)) + outcome = adapter.step(0, Scenario("s", {"seed": 0}, max_steps=10)) + # MockIsaacEnv's info is {} so allowlist returns {} + self.assertEqual(outcome.info["isaac"]["raw_info"], {}) + + +# ───────────────────────────────────────────────────────────────────────── +# Action shape normalization +# ───────────────────────────────────────────────────────────────────────── + + +@unittest.skipUnless(HAS_TORCH, "torch required") +class ActionShapingTest(unittest.TestCase): + """Action shape normalization tests. + + Two realistic scenarios covered: + 1. Discrete-style env (action_dim=1) + scalar action → (1, 1) + 2. Continuous env (action_dim=2) + 1-D array action → (1, 2) + + A scalar action against a multi-dim continuous env is ambiguous user + error; we don't broadcast the scalar across all action dims because + that's almost never what the user intended. + """ + + def test_scalar_action_for_single_dim_env(self): + env = MockIsaacEnv(num_envs=1, action_dim=1) + adapter = IsaacEnvironmentAdapter(env=env) + adapter.reset(Scenario("s", {"seed": 0}, max_steps=10)) + adapter.step(1, Scenario("s", {"seed": 0}, max_steps=10)) + action = env._last_action + self.assertIsNotNone(action) + self.assertTrue(isinstance(action, torch.Tensor)) + self.assertEqual(action.shape, (1, 1)) + + def test_1d_array_action_for_multi_dim_env(self): + env = MockIsaacEnv(num_envs=1, action_dim=2) + adapter = IsaacEnvironmentAdapter(env=env) + adapter.reset(Scenario("s", {"seed": 0}, max_steps=10)) + adapter.step( + np.array([0.5, -0.3]), Scenario("s", {"seed": 0}, max_steps=10) + ) + action = env._last_action + self.assertEqual(action.shape, (1, 2)) + + +# ───────────────────────────────────────────────────────────────────────── +# Serialization round-trip +# ───────────────────────────────────────────────────────────────────────── + + +@unittest.skipUnless(HAS_TORCH, "torch required") +class SerializationTest(unittest.TestCase): + def test_step_outcome_is_json_safe(self): + env = MockIsaacEnv(num_envs=1, terminate_at_step=10) + adapter = IsaacEnvironmentAdapter(env=env) + adapter.reset(Scenario("s", {"seed": 0}, max_steps=10)) + outcome = adapter.step(0, Scenario("s", {"seed": 0}, max_steps=10)) + # All of these should round-trip through to_serializable + json.dumps + for field in ( + outcome.next_state, + outcome.metrics, + outcome.events, + outcome.info, + ): + payload = to_serializable(field) + json.dumps(payload) + + +# ───────────────────────────────────────────────────────────────────────── +# Edge case tests — non-standard inputs from real Isaac envs +# ───────────────────────────────────────────────────────────────────────── + + +class MockIsaacEnvWithRichInfo(MockIsaacEnv): + """Mock that returns a non-empty info dict each step.""" + + def step(self, action): + obs, reward, terminated, truncated, _info = super().step(action) + info = { + "is_success": bool(terminated[0].item()), + "step_count": self._step, + "task_reward": float(reward[0].item()), + "huge_tensor": torch.zeros(100, 100), + } + return obs, reward, terminated, truncated, info + + +class MockIsaacEnvWithNonDictInfo(MockIsaacEnv): + """Mock that returns a list info (some Isaac envs do this per-env).""" + + def step(self, action): + obs, reward, terminated, truncated, _info = super().step(action) + # Some Isaac envs return a list of per-env infos instead of a dict + return obs, reward, terminated, truncated, ["per_env_info_0"] + + +@unittest.skipUnless(HAS_TORCH, "torch required") +class EdgeCaseTest(unittest.TestCase): + def test_rich_info_passes_through_via_to_serializable(self): + env = MockIsaacEnvWithRichInfo(num_envs=1) + adapter = IsaacEnvironmentAdapter(env=env) + adapter.reset(Scenario("s", {"seed": 0}, max_steps=10)) + outcome = adapter.step(0, Scenario("s", {"seed": 0}, max_steps=10)) + raw_info = outcome.info["isaac"]["raw_info"] + # All keys present (default no allowlist) + self.assertIn("is_success", raw_info) + self.assertIn("step_count", raw_info) + self.assertIn("task_reward", raw_info) + self.assertIn("huge_tensor", raw_info) + # JSON round-trip must succeed even with the tensor + json.dumps(raw_info) + + def test_info_keys_allowlist_drops_unwanted_keys(self): + env = MockIsaacEnvWithRichInfo(num_envs=1) + adapter = IsaacEnvironmentAdapter(env=env, info_keys=["is_success", "step_count"]) + adapter.reset(Scenario("s", {"seed": 0}, max_steps=10)) + outcome = adapter.step(0, Scenario("s", {"seed": 0}, max_steps=10)) + raw_info = outcome.info["isaac"]["raw_info"] + self.assertIn("is_success", raw_info) + self.assertIn("step_count", raw_info) + # Filtered out + self.assertNotIn("huge_tensor", raw_info) + self.assertNotIn("task_reward", raw_info) + + def test_non_dict_info_does_not_crash(self): + env = MockIsaacEnvWithNonDictInfo(num_envs=1) + adapter = IsaacEnvironmentAdapter(env=env) + adapter.reset(Scenario("s", {"seed": 0}, max_steps=10)) + outcome = adapter.step(0, Scenario("s", {"seed": 0}, max_steps=10)) + # Non-dict info gets filtered to {}, but doesn't crash + self.assertEqual(outcome.info["isaac"]["raw_info"], {}) + + def test_custom_outcome_hook_reads_info(self): + """Manipulation-style outcome detection: use info['is_success'].""" + env = MockIsaacEnvWithRichInfo(num_envs=1, terminate_at_step=3) + + def info_based_outcome(reward, terminated, truncated, info): + if info.get("is_success"): + return ("goal_reached", "") + if terminated: + return ("terminated_failure", "did_not_succeed") + if truncated: + return ("truncated", "timeout") + return ("progress", "") + + adapter = IsaacEnvironmentAdapter(env=env, outcome_from_step=info_based_outcome) + adapter.reset(Scenario("s", {"seed": 0}, max_steps=10)) + # Step until terminal + for _ in range(5): + outcome = adapter.step(0, Scenario("s", {"seed": 0}, max_steps=10)) + if outcome.terminal: + break + # is_success becomes True when terminated → outcome should be goal_reached + self.assertEqual(outcome.outcome, "goal_reached") + + def test_multiple_episodes_via_reset(self): + """Verify the adapter is re-entrant across multiple reset/step cycles.""" + env = MockIsaacEnv(num_envs=1, terminate_at_step=5) + adapter = IsaacEnvironmentAdapter(env=env) + + for episode_i in range(3): + adapter.reset(Scenario("s", {"seed": episode_i}, max_steps=10)) + total_steps = 0 + for _ in range(20): + outcome = adapter.step(0, Scenario("s", {"seed": episode_i}, max_steps=10)) + total_steps += 1 + if outcome.terminal: + break + # Each episode should terminate at step 5 + self.assertEqual(total_steps, 5) + self.assertEqual(outcome.metrics["episode_return"], 5.0) + + def test_pre_batched_2d_action_passes_through(self): + """If the user already provides a (1, action_dim) tensor, don't re-shape.""" + env = MockIsaacEnv(num_envs=1, action_dim=2) + adapter = IsaacEnvironmentAdapter(env=env) + adapter.reset(Scenario("s", {"seed": 0}, max_steps=10)) + # Pre-shaped tensor + action = torch.tensor([[0.5, -0.3]]) + adapter.step(action, Scenario("s", {"seed": 0}, max_steps=10)) + sent_action = env._last_action + self.assertEqual(sent_action.shape, (1, 2)) + # Values preserved + self.assertAlmostEqual(sent_action[0, 0].item(), 0.5, places=5) + self.assertAlmostEqual(sent_action[0, 1].item(), -0.3, places=5) + + +# ───────────────────────────────────────────────────────────────────────── +# Full EvalRunner integration (mocked Isaac env) +# ───────────────────────────────────────────────────────────────────────── + + +@unittest.skipUnless(HAS_TORCH, "torch required") +class EvalRunnerIntegrationTest(unittest.TestCase): + def test_eval_runner_produces_report(self): + def naive_policy(state): + return {"action": 0} + + env = MockIsaacEnv(num_envs=1, terminate_at_step=20) + adapter = IsaacEnvironmentAdapter(env=env, name="mock_isaac") + + ruleset = Ruleset( + [ + require_metric("episode_return", ">=", 5.0, name="balance_5_steps"), + ] + ) + + report = EvalRunner( + policies=[naive_policy], + scenarios=[Scenario("smoke", {"seed": 0}, max_steps=30)], + ruleset=ruleset, + baseline_policy="naive_policy", + environment=adapter, + ).run() + + self.assertEqual(len(report.episodes), 1) + self.assertIn("naive_policy", report.metric_summary) + self.assertIn("episode_return", report.metric_summary["naive_policy"]) + # Should succeed since 20 steps of reward=1.0 >> 5 + self.assertTrue(report.episodes[0].success) + + def test_eval_runner_with_multiple_policies(self): + """Test the full comparison path: 2 policies, 2 scenarios, regression detection.""" + def good_policy(state): + return {"action": 0, "debug_info": {"version": "good"}} + + def bad_policy(state): + return {"action": 0, "debug_info": {"version": "bad"}} + + # Different envs with different terminate-at-step thresholds to simulate + # one passing and one failing the rule + env = MockIsaacEnv(num_envs=1, terminate_at_step=10) + adapter = IsaacEnvironmentAdapter(env=env, name="mock_isaac") + + ruleset = Ruleset( + [ + require_metric("episode_return", ">=", 5.0, name="balance_5_steps"), + ] + ) + + report = EvalRunner( + policies=[good_policy, bad_policy], + scenarios=[ + Scenario("s1", {"seed": 0}, max_steps=20), + Scenario("s2", {"seed": 1}, max_steps=20), + ], + ruleset=ruleset, + baseline_policy="good_policy", + environment=adapter, + ).run() + + # 2 policies × 2 scenarios = 4 episodes + self.assertEqual(len(report.episodes), 4) + # Both policies should succeed (10 steps > 5 threshold) + for ep in report.episodes: + self.assertTrue(ep.success) + # Metric summary should track both policies + self.assertIn("good_policy", report.metric_summary) + self.assertIn("bad_policy", report.metric_summary) + # Action divergences should be empty since both policies do action=0 + self.assertEqual(len(report.action_divergences), 0) + + +if __name__ == "__main__": + unittest.main() From a986bda009ab1feb24cd630dd10a02e1ad91ee7b Mon Sep 17 00:00:00 2001 From: arryan0112 Date: Mon, 8 Jun 2026 14:30:57 +0530 Subject: [PATCH 2/6] Add roboeval/batched core: types, environment Protocol, policy adapter - BatchedState (list[State]) and BatchedStepOutcome (per-slot lists with length validation) - BatchedEnvironmentAdapter Protocol (reset/step/reset_slots) for N-env vectorized rollouts - BatchedPolicy Protocol, BatchedPolicyAdapter, and from_single() shim for single-state policies - Tests for type construction, validation, slot extraction, and policy normalization --- roboeval/batched/__init__.py | 27 +++++ roboeval/batched/environment.py | 64 +++++++++++ roboeval/batched/policy.py | 102 +++++++++++++++++ roboeval/batched/tests/__init__.py | 0 roboeval/batched/tests/test_policy.py | 158 ++++++++++++++++++++++++++ roboeval/batched/tests/test_types.py | 144 +++++++++++++++++++++++ roboeval/batched/types.py | 77 +++++++++++++ 7 files changed, 572 insertions(+) create mode 100644 roboeval/batched/__init__.py create mode 100644 roboeval/batched/environment.py create mode 100644 roboeval/batched/policy.py create mode 100644 roboeval/batched/tests/__init__.py create mode 100644 roboeval/batched/tests/test_policy.py create mode 100644 roboeval/batched/tests/test_types.py create mode 100644 roboeval/batched/types.py diff --git a/roboeval/batched/__init__.py b/roboeval/batched/__init__.py new file mode 100644 index 0000000..353af69 --- /dev/null +++ b/roboeval/batched/__init__.py @@ -0,0 +1,27 @@ +"""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 .types import BatchedState, BatchedStepOutcome + + +__all__ = [ + "BatchedEnvironmentAdapter", + "BatchedPolicy", + "BatchedPolicyAdapter", + "BatchedState", + "BatchedStepOutcome", + "from_single", + "normalize_batched_policy", +] diff --git a/roboeval/batched/environment.py b/roboeval/batched/environment.py new file mode 100644 index 0000000..ae61310 --- /dev/null +++ b/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. + """ + ... diff --git a/roboeval/batched/policy.py b/roboeval/batched/policy.py new file mode 100644 index 0000000..346541d --- /dev/null +++ b/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) diff --git a/roboeval/batched/tests/__init__.py b/roboeval/batched/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/roboeval/batched/tests/test_policy.py b/roboeval/batched/tests/test_policy.py new file mode 100644 index 0000000..ed24f42 --- /dev/null +++ b/roboeval/batched/tests/test_policy.py @@ -0,0 +1,158 @@ +"""Tests for BatchedPolicyAdapter normalization and the from_single shim.""" + +from __future__ import annotations + +import unittest + +from roboeval.batched.policy import ( + BatchedPolicyAdapter, + from_single, + normalize_batched_policy, +) +from roboeval.core import Decision + + +class TestBatchedPolicyAdapterNormalizesItems(unittest.TestCase): + def test_normalizes_dict_returns(self) -> None: + def policy(states): + return [{"action": i, "debug_info": {"slot": i}} for i in range(len(states))] + + adapter = BatchedPolicyAdapter(name="dict_policy", policy=policy) + decisions = adapter.decide([{}, {}, {}]) + self.assertEqual(len(decisions), 3) + self.assertTrue(all(isinstance(d, Decision) for d in decisions)) + self.assertEqual(decisions[0].action, 0) + self.assertEqual(decisions[2].debug_info["slot"], 2) + + def test_normalizes_tuple_returns(self) -> None: + def policy(states): + return [(i, {"version": "v1"}) for i in range(len(states))] + + adapter = BatchedPolicyAdapter(name="tuple_policy", policy=policy) + decisions = adapter.decide([{}, {}]) + self.assertEqual(decisions[0].action, 0) + self.assertEqual(decisions[1].debug_info, {"version": "v1"}) + + def test_normalizes_raw_action_returns(self) -> None: + def policy(states): + return [0 for _ in states] + + adapter = BatchedPolicyAdapter(name="raw_policy", policy=policy) + decisions = adapter.decide([{}, {}]) + self.assertEqual(decisions[0].action, 0) + self.assertEqual(decisions[0].debug_info, {}) + + def test_normalizes_decision_returns_unchanged(self) -> None: + canned = [Decision(action=42, debug_info={"k": "v"})] + + def policy(states): + return canned + + adapter = BatchedPolicyAdapter(name="decision_policy", policy=policy) + decisions = adapter.decide([{}]) + self.assertEqual(decisions[0].action, 42) + self.assertEqual(decisions[0].debug_info, {"k": "v"}) + + def test_mixed_return_shapes_in_one_batch(self) -> None: + def policy(states): + return [{"action": 1}, (2, {"version": "v2"}), 3] + + adapter = BatchedPolicyAdapter(name="mixed", policy=policy) + decisions = adapter.decide([{}, {}, {}]) + self.assertEqual([d.action for d in decisions], [1, 2, 3]) + self.assertEqual(decisions[1].debug_info, {"version": "v2"}) + + +class TestBatchedPolicyAdapterValidation(unittest.TestCase): + def test_non_list_return_raises(self) -> None: + def bad_policy(states): + return {"oops": "dict not list"} + + adapter = BatchedPolicyAdapter(name="bad", policy=bad_policy) + with self.assertRaisesRegex(TypeError, r"must return a list"): + adapter.decide([{}, {}]) + + def test_length_mismatch_raises(self) -> None: + def short_policy(states): + return [0] + + adapter = BatchedPolicyAdapter(name="short", policy=short_policy) + with self.assertRaisesRegex(ValueError, r"returned 1 decisions for 3 input"): + adapter.decide([{}, {}, {}]) + + def test_non_callable_no_decide_raises(self) -> None: + adapter = BatchedPolicyAdapter(name="bad", policy=object()) + with self.assertRaisesRegex(TypeError, r"is not callable and has no decide"): + adapter.decide([{}]) + + def test_object_with_decide_method_works(self) -> None: + class CtrlPolicy: + def decide(self, states): + return [{"action": "noop"} for _ in states] + + adapter = BatchedPolicyAdapter(name="ctrl", policy=CtrlPolicy()) + decisions = adapter.decide([{}, {}]) + self.assertEqual(decisions[0].action, "noop") + + +class TestNormalizeBatchedPolicy(unittest.TestCase): + def test_passes_through_existing_adapter(self) -> None: + def f(states): + return [0] * len(states) + + original = BatchedPolicyAdapter(name="orig", policy=f) + self.assertIs(normalize_batched_policy(original), original) + + def test_uses_function_name_when_no_name_given(self) -> None: + def my_named_policy(states): + return [0] * len(states) + + adapter = normalize_batched_policy(my_named_policy) + self.assertEqual(adapter.name, "my_named_policy") + + def test_uses_version_attribute_if_present(self) -> None: + class VersionedPolicy: + version = "v3.1.4" + + def __call__(self, states): + return [0] * len(states) + + adapter = normalize_batched_policy(VersionedPolicy()) + self.assertEqual(adapter.name, "v3.1.4") + + +class TestFromSingleShim(unittest.TestCase): + def test_wraps_single_state_policy(self) -> None: + def single(state): + return {"action": state["x"] * 2} + + batched = from_single(single) + decisions = batched.decide([{"x": 1}, {"x": 2}, {"x": 5}]) + self.assertEqual([d.action for d in decisions], [2, 4, 10]) + + def test_inherits_function_name(self) -> None: + def baseline_policy(state): + return 0 + + batched = from_single(baseline_policy) + self.assertEqual(batched.name, "baseline_policy") + + def test_explicit_name_wins(self) -> None: + def fn(state): + return 0 + + batched = from_single(fn, name="renamed") + self.assertEqual(batched.name, "renamed") + + def test_returns_normalized_decisions(self) -> None: + def returns_tuple(state): + return (state["v"], {"x": "y"}) + + batched = from_single(returns_tuple) + decisions = batched.decide([{"v": 7}]) + self.assertEqual(decisions[0].action, 7) + self.assertEqual(decisions[0].debug_info, {"x": "y"}) + + +if __name__ == "__main__": + unittest.main() diff --git a/roboeval/batched/tests/test_types.py b/roboeval/batched/tests/test_types.py new file mode 100644 index 0000000..fb52af3 --- /dev/null +++ b/roboeval/batched/tests/test_types.py @@ -0,0 +1,144 @@ +"""Tests for BatchedStepOutcome construction, validation, and slot extraction.""" + +from __future__ import annotations + +import unittest + +from roboeval.batched.types import BatchedStepOutcome +from roboeval.environment import StepOutcome + + +class TestBatchedStepOutcomeMinimalConstruction(unittest.TestCase): + def test_constructs_with_required_fields_only(self) -> None: + outcome = BatchedStepOutcome( + next_states=[{"obs": 1}, {"obs": 2}, {"obs": 3}], + outcomes=["progress", "progress", "terminated_success"], + failure_labels=["", "", ""], + terminals=[False, False, True], + ) + self.assertEqual(outcome.num_envs, 3) + self.assertIsNone(outcome.metrics) + self.assertIsNone(outcome.events) + self.assertIsNone(outcome.artifacts) + self.assertIsNone(outcome.info) + + def test_num_envs_zero_is_legal(self) -> None: + outcome = BatchedStepOutcome( + next_states=[], outcomes=[], failure_labels=[], terminals=[] + ) + self.assertEqual(outcome.num_envs, 0) + + +class TestBatchedStepOutcomeFullConstruction(unittest.TestCase): + def test_constructs_with_all_fields(self) -> None: + outcome = BatchedStepOutcome( + next_states=[{"obs": 1}, {"obs": 2}], + outcomes=["progress", "terminated_failure"], + failure_labels=["", "collision"], + terminals=[False, True], + metrics=[{"reward": 1.0}, {"reward": -1.0}], + events=[["step_logged"], ["episode_terminated"]], + artifacts=[{}, {"frame": b"\x00\x01"}], + info=[{"env": "a"}, {"env": "b", "error": "collision_detected"}], + ) + self.assertEqual(outcome.num_envs, 2) + self.assertEqual(outcome.metrics[1]["reward"], -1.0) + self.assertEqual(outcome.events[1], ["episode_terminated"]) + self.assertEqual(outcome.info[1]["error"], "collision_detected") + + +class TestBatchedStepOutcomeValidation(unittest.TestCase): + def test_outcomes_length_mismatch_raises(self) -> None: + with self.assertRaisesRegex(ValueError, r"outcomes.*length 2.*expected 3"): + BatchedStepOutcome( + next_states=[{"a": 1}, {"a": 2}, {"a": 3}], + outcomes=["progress", "progress"], + failure_labels=["", "", ""], + terminals=[False, False, False], + ) + + def test_terminals_length_mismatch_raises(self) -> None: + with self.assertRaisesRegex(ValueError, r"terminals.*length 1.*expected 2"): + BatchedStepOutcome( + next_states=[{"a": 1}, {"a": 2}], + outcomes=["progress", "progress"], + failure_labels=["", ""], + terminals=[False], + ) + + def test_metrics_length_mismatch_raises(self) -> None: + with self.assertRaisesRegex(ValueError, r"metrics.*length 1.*expected 2"): + BatchedStepOutcome( + next_states=[{"a": 1}, {"a": 2}], + outcomes=["progress", "progress"], + failure_labels=["", ""], + terminals=[False, False], + metrics=[{"reward": 1.0}], + ) + + def test_failure_labels_length_mismatch_raises(self) -> None: + with self.assertRaisesRegex(ValueError, r"failure_labels.*length 0.*expected 1"): + BatchedStepOutcome( + next_states=[{"a": 1}], + outcomes=["progress"], + failure_labels=[], + terminals=[False], + ) + + def test_info_length_mismatch_raises(self) -> None: + with self.assertRaisesRegex(ValueError, r"info.*length 3.*expected 2"): + BatchedStepOutcome( + next_states=[{"a": 1}, {"a": 2}], + outcomes=["progress", "progress"], + failure_labels=["", ""], + terminals=[False, False], + info=[{}, {}, {}], + ) + + +class TestBatchedStepOutcomeSlotExtraction(unittest.TestCase): + def test_slot_returns_single_env_step_outcome(self) -> None: + outcome = BatchedStepOutcome( + next_states=[{"obs": 1}, {"obs": 2}, {"obs": 3}], + outcomes=["progress", "progress", "terminated_success"], + failure_labels=["", "", ""], + terminals=[False, False, True], + metrics=[{"reward": 1.0}, {"reward": 0.5}, {"reward": 10.0}], + events=[["a"], ["b"], ["episode_terminated"]], + info=[{"x": 1}, {"x": 2}, {"x": 3}], + ) + slot2 = outcome.slot(2) + self.assertIsInstance(slot2, StepOutcome) + self.assertEqual(slot2.next_state, {"obs": 3}) + self.assertEqual(slot2.outcome, "terminated_success") + self.assertEqual(slot2.terminal, True) + self.assertEqual(slot2.metrics, {"reward": 10.0}) + self.assertEqual(slot2.events, ["episode_terminated"]) + self.assertEqual(slot2.info, {"x": 3}) + + def test_slot_with_none_optional_fields(self) -> None: + outcome = BatchedStepOutcome( + next_states=[{"obs": 1}], + outcomes=["progress"], + failure_labels=[""], + terminals=[False], + ) + slot0 = outcome.slot(0) + self.assertIsNone(slot0.metrics) + self.assertIsNone(slot0.events) + self.assertIsNone(slot0.artifacts) + self.assertIsNone(slot0.info) + + def test_slot_out_of_range_raises_index_error(self) -> None: + outcome = BatchedStepOutcome( + next_states=[{"obs": 1}], + outcomes=["progress"], + failure_labels=[""], + terminals=[False], + ) + with self.assertRaises(IndexError): + outcome.slot(1) + + +if __name__ == "__main__": + unittest.main() diff --git a/roboeval/batched/types.py b/roboeval/batched/types.py new file mode 100644 index 0000000..b01e09a --- /dev/null +++ b/roboeval/batched/types.py @@ -0,0 +1,77 @@ +"""Batched analogs of the single-env State and StepOutcome. + +Vectorized envs (Isaac Lab, gym.vector, Brax, MJX) run N envs in parallel and +return batched results. This module mirrors the single-env types from +roboeval.core and roboeval.environment, but every payload is a per-slot list of +length num_envs. + +Design notes: + * BatchedState is a list[State] alias rather than a dataclass so callers can + slice, append, and reorder slots with plain list operations. + * BatchedStepOutcome is a dataclass with strict length checks; mismatched + list lengths are the most common source of bugs at the SDK/sim boundary. + * The .slot(i) helper produces a single-env StepOutcome for that slot. The + runner uses it to fan the batch back into per-slot episode buffers, which + keeps the report format identical between single-env and batched runs. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from roboeval.core import State +from roboeval.environment import StepOutcome + + +BatchedState = list[State] + + +@dataclass +class BatchedStepOutcome: + next_states: list[State] + outcomes: list[str] + failure_labels: list[str] + terminals: list[bool] + metrics: list[dict[str, float | int] | None] | None = None + events: list[list[str] | None] | None = None + artifacts: list[dict[str, object] | None] | None = None + info: list[dict[str, object] | None] | None = None + + def __post_init__(self) -> None: + n = len(self.next_states) + _check_length("outcomes", self.outcomes, n) + _check_length("failure_labels", self.failure_labels, n) + _check_length("terminals", self.terminals, n) + if self.metrics is not None: + _check_length("metrics", self.metrics, n) + if self.events is not None: + _check_length("events", self.events, n) + if self.artifacts is not None: + _check_length("artifacts", self.artifacts, n) + if self.info is not None: + _check_length("info", self.info, n) + + @property + def num_envs(self) -> int: + return len(self.next_states) + + def slot(self, i: int) -> StepOutcome: + """Return slot ``i`` as a single-env StepOutcome.""" + return StepOutcome( + next_state=self.next_states[i], + outcome=self.outcomes[i], + failure_label=self.failure_labels[i], + terminal=self.terminals[i], + metrics=self.metrics[i] if self.metrics is not None else None, + events=self.events[i] if self.events is not None else None, + artifacts=self.artifacts[i] if self.artifacts is not None else None, + info=self.info[i] if self.info is not None else None, + ) + + +def _check_length(field_name: str, value: list, expected: int) -> None: + if len(value) != expected: + raise ValueError( + f"BatchedStepOutcome.{field_name} has length {len(value)}, " + f"expected {expected} (num_envs)." + ) From 81d89dd64731e66840fba86de1543282b7903e18 Mon Sep 17 00:00:00 2001 From: arryan0112 Date: Mon, 8 Jun 2026 14:31:42 +0530 Subject: [PATCH 3/6] Add SlotScheduler and BatchedEvalRunner with refill semantics - SlotScheduler: maps (Scenario, replica) tasks to N env slots with refill on slot completion; replication support (D1) for running each scenario K times - BatchedEvalRunner: drives one policy at a time through the scheduler, fans per-slot outcomes into normal EpisodeResults, reuses _build_report from EvalRunner so reports are byte-identical to the single-env path - Tests: scheduler init/refill/replication/idle handling; runner smoke, refill, idle slots, multi-policy, replication scenario_name suffix, max_steps enforcement, ruleset flow, report shape --- roboeval/batched/__init__.py | 5 + roboeval/batched/runner.py | 276 ++++++++++++++++ roboeval/batched/scheduler.py | 136 ++++++++ roboeval/batched/tests/test_runner.py | 384 +++++++++++++++++++++++ roboeval/batched/tests/test_scheduler.py | 143 +++++++++ 5 files changed, 944 insertions(+) create mode 100644 roboeval/batched/runner.py create mode 100644 roboeval/batched/scheduler.py create mode 100644 roboeval/batched/tests/test_runner.py create mode 100644 roboeval/batched/tests/test_scheduler.py diff --git a/roboeval/batched/__init__.py b/roboeval/batched/__init__.py index 353af69..656c603 100644 --- a/roboeval/batched/__init__.py +++ b/roboeval/batched/__init__.py @@ -13,15 +13,20 @@ 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", ] diff --git a/roboeval/batched/runner.py b/roboeval/batched/runner.py new file mode 100644 index 0000000..f565bd7 --- /dev/null +++ b/roboeval/batched/runner.py @@ -0,0 +1,276 @@ +"""BatchedEvalRunner — the vectorized counterpart of EvalRunner. + +Drives a BatchedEnvironmentAdapter through (policies × scenarios × replicas) +rollouts using a SlotScheduler to keep all N env slots busy. Per-slot +episodes are buffered into normal single-env StepRecord lists, then fed into +the same _build_report() that the single-env runner uses — so reports come +out byte-identical between the two paths (decision D5). + +Execution model: + * One policy at a time. The runner does all of policy P's rollouts before + moving to policy P+1. This keeps things simple and is how lerobot_eval + works. Mixing policies across slots is a future optimization. + * Idle slot handling: when the queue drains but other slots are still + rolling, idle slots receive a placeholder action (copied from the first + active slot) and their outcomes are discarded. + * max_steps enforcement: per-slot step counter. If a slot reaches + scenario.max_steps without env-terminating, the runner force-terminates + that episode with terminal_outcome="max_steps_reached". + +Replication (D1): pass ``replicas=N`` to run each scenario N times. With +replicas > 1, EpisodeResult.scenario_name is suffixed with ``#r{i}`` so each +replica produces a distinct report row. +""" + +from __future__ import annotations + +from typing import Any, Iterable + +from roboeval.core import ( + EpisodeContext, + EpisodeResult, + EvalReport, + Ruleset, + Scenario, + StepRecord, + SuccessCriteria, +) +from roboeval.runner import _build_report, _environment_name, _validate_step_outcome + +from .environment import BatchedEnvironmentAdapter +from .policy import BatchedPolicyAdapter, normalize_batched_policy +from .scheduler import SlotScheduler, SlotTask +from .types import BatchedStepOutcome + + +class BatchedEvalRunner: + """Vectorized eval runner. Produces the same EvalReport as EvalRunner.""" + + def __init__( + self, + policies: Iterable[Any], + scenarios: Iterable[Scenario], + environment: BatchedEnvironmentAdapter, + ruleset: Ruleset | None = None, + baseline_policy: str | None = None, + replicas: int = 1, + ) -> None: + self.policies: list[BatchedPolicyAdapter] = [ + normalize_batched_policy(policy) for policy in policies + ] + self.scenarios: list[Scenario] = list(scenarios) + self.environment = environment + self.environment_name = _environment_name(environment) + self.ruleset = ruleset or SuccessCriteria().to_ruleset() + self.baseline_policy = baseline_policy or ( + self.policies[0].name if self.policies else "" + ) + if replicas <= 0: + raise ValueError(f"replicas must be positive, got {replicas}.") + self.replicas = replicas + if not self.policies: + raise ValueError("BatchedEvalRunner needs at least one policy.") + if not self.scenarios: + raise ValueError("BatchedEvalRunner needs at least one scenario.") + if not hasattr(environment, "num_envs"): + raise TypeError( + "BatchedEvalRunner.environment must expose num_envs " + "(BatchedEnvironmentAdapter Protocol)." + ) + + # ----- public ----- + + def run(self) -> EvalReport: + episodes: list[EpisodeResult] = [] + for policy in self.policies: + episodes.extend(self._run_policy(policy)) + return _build_report(episodes, self.baseline_policy, self.environment_name) + + # ----- per-policy rollout ----- + + def _run_policy(self, policy: BatchedPolicyAdapter) -> list[EpisodeResult]: + n = self.environment.num_envs + sched = SlotScheduler( + num_envs=n, scenarios=self.scenarios, replicas=self.replicas + ) + initial_scenarios = sched.initialize() + + if sched.is_done(): + return [] + + # Build the reset batch: every slot needs a Scenario for env.reset. + # Idle slots get a placeholder (any active scenario), their outcomes + # are ignored throughout. + placeholder_scenario = next(s for s in initial_scenarios if s is not None) + reset_batch = [s if s is not None else placeholder_scenario for s in initial_scenarios] + slot_states = list(self.environment.reset(reset_batch)) + if len(slot_states) != n: + raise ValueError( + f"environment.reset returned {len(slot_states)} states, expected {n}." + ) + + # Per-slot rollout buffers + slot_step_counts: list[int] = [0] * n + slot_logs: list[list[StepRecord]] = [[] for _ in range(n)] + completed: list[EpisodeResult] = [] + + while not sched.is_done(): + active = sched.active_slots() + active_states = [slot_states[i] for i in active] + decisions = policy.decide(active_states) + slot_decisions = [None] * n + for j, slot in enumerate(active): + slot_decisions[slot] = decisions[j] + + # Build full action batch (idle slots get a placeholder action). + placeholder_action = slot_decisions[active[0]].action + actions = [ + slot_decisions[i].action if slot_decisions[i] is not None else placeholder_action + for i in range(n) + ] + + outcome_batch = self.environment.step(actions) + _validate_batched_step_outcome(outcome_batch, n) + + slots_to_refill_idx: list[int] = [] + scenarios_to_refill: list[Scenario] = [] + + for slot in active: + task = sched.current_task(slot) + step_outcome = outcome_batch.slot(slot) + # Per-slot validation reuses the single-env validator + _validate_step_outcome(step_outcome) + decision = slot_decisions[slot] + step_idx = slot_step_counts[slot] + + episode_id, scenario_name = _episode_identity( + task, policy.name, self.replicas + ) + + slot_logs[slot].append( + StepRecord( + episode_id=episode_id, + scenario_name=scenario_name, + policy_version=policy.name, + step=step_idx, + state=dict(slot_states[slot]), + action=decision.action, + outcome=step_outcome.outcome, + failure_label=step_outcome.failure_label, + next_state=dict(step_outcome.next_state), + is_terminal=bool(step_outcome.terminal), + debug_info=dict(decision.debug_info), + metrics=dict(step_outcome.metrics or {}), + events=list(step_outcome.events or []), + artifacts=dict(step_outcome.artifacts or {}), + info=dict(step_outcome.info or {}), + ) + ) + + slot_states[slot] = step_outcome.next_state + slot_step_counts[slot] = step_idx + 1 + + env_term = bool(step_outcome.terminal) + max_steps_hit = slot_step_counts[slot] >= task.scenario.max_steps + if env_term or max_steps_hit: + terminal_outcome = ( + step_outcome.outcome if env_term else "max_steps_reached" + ) + completed.append( + _finalize_episode( + task=task, + policy_name=policy.name, + replicas=self.replicas, + logs=slot_logs[slot], + terminal_outcome=terminal_outcome, + ruleset=self.ruleset, + ) + ) + refill = sched.complete_slot(slot) + if refill is not None: + slots_to_refill_idx.append(slot) + scenarios_to_refill.append(refill.scenario) + slot_step_counts[slot] = 0 + slot_logs[slot] = [] + + if slots_to_refill_idx: + new_states = self.environment.reset_slots( + slots_to_refill_idx, scenarios_to_refill + ) + if len(new_states) != len(slots_to_refill_idx): + raise ValueError( + f"reset_slots returned {len(new_states)} states for " + f"{len(slots_to_refill_idx)} requested slots." + ) + for j, slot in enumerate(slots_to_refill_idx): + slot_states[slot] = new_states[j] + + return completed + + +# ----- helpers ----- + + +def _episode_identity( + task: SlotTask, policy_name: str, replicas: int +) -> tuple[str, str]: + """Return (episode_id, scenario_name) for the task. + + With replicas > 1, scenario_name carries a #r{idx} suffix so each replica + produces a distinct EpisodeResult row in the report. + """ + if replicas > 1: + scenario_name = f"{task.scenario.name}#r{task.replica_idx}" + else: + scenario_name = task.scenario.name + episode_id = f"{scenario_name}:{policy_name}" + return episode_id, scenario_name + + +def _finalize_episode( + *, + task: SlotTask, + policy_name: str, + replicas: int, + logs: list[StepRecord], + terminal_outcome: str, + ruleset: Ruleset, +) -> EpisodeResult: + episode_id, scenario_name = _episode_identity(task, policy_name, replicas) + context = EpisodeContext( + episode_id=episode_id, + scenario=task.scenario, + policy_version=policy_name, + logs=logs, + terminal_outcome=terminal_outcome, + ) + rule_results = ruleset.evaluate(context) + first_failure = next((r for r in rule_results if not r.passed), None) + success = first_failure is None + failure_label = "" if success else first_failure.name + return EpisodeResult( + episode_id=episode_id, + scenario_name=scenario_name, + policy_version=policy_name, + success=success, + terminal_outcome=terminal_outcome, + failure_label=failure_label, + steps=len(logs), + logs=logs, + rule_results=rule_results, + first_failure_step=first_failure.step if first_failure else None, + scenario_metadata=dict(task.scenario.metadata), + ) + + +def _validate_batched_step_outcome(outcome: Any, expected_num_envs: int) -> None: + if not isinstance(outcome, BatchedStepOutcome): + raise TypeError( + f"BatchedEnvironmentAdapter.step must return BatchedStepOutcome, " + f"got {type(outcome).__name__}." + ) + if outcome.num_envs != expected_num_envs: + raise ValueError( + f"BatchedStepOutcome.num_envs is {outcome.num_envs}, expected " + f"{expected_num_envs} (matches environment.num_envs)." + ) diff --git a/roboeval/batched/scheduler.py b/roboeval/batched/scheduler.py new file mode 100644 index 0000000..cbf14df --- /dev/null +++ b/roboeval/batched/scheduler.py @@ -0,0 +1,136 @@ +"""SlotScheduler — maps a queue of (scenario, replica) tasks to N env slots. + +Owns the bookkeeping for vectorized rollouts: + * which slot is running which scenario, and which replica index of it + * which slots are active (have a task) vs idle (queue drained) + * when a slot terminates, whether the runner should refill it from the + pending queue or leave it idle + +The scheduler is intentionally environment- and policy-agnostic so it can be +unit-tested in isolation. The runner calls into it; the scheduler never calls +back into the runner. + +Task identity: + Each (Scenario, replica_idx) pair is one task. When ``replicas > 1`` the + replica_idx differentiates rollouts of the same Scenario; downstream + reporting suffixes ``#r{i}`` to ``scenario_name`` so each replica produces + a distinct EpisodeResult row. +""" + +from __future__ import annotations + +from collections import deque +from dataclasses import dataclass + +from roboeval.core import Scenario + + +@dataclass(frozen=True) +class SlotTask: + """One pending or in-flight task: which scenario, which replica.""" + scenario: Scenario + replica_idx: int + + +class SlotScheduler: + """Assigns (Scenario, replica) tasks to a fixed bank of ``num_envs`` slots. + + Lifecycle: + sched = SlotScheduler(num_envs=8, scenarios=[...], replicas=4) + initial_scenarios = sched.initialize() # len == num_envs, may include None entries + # ... runner resets the env with initial_scenarios as a batch ... + while not sched.is_done(): + active = sched.active_slots() + # ... runner calls policy + env.step ... + for slot in active: + if terminal[slot] or max_steps_hit[slot]: + refilled = sched.complete_slot(slot) # returns Optional[SlotTask] + if refilled: + # ... runner calls env.reset_slots([slot], [refilled.scenario]) ... + + Args: + num_envs: width of the env batch — number of parallel slots + scenarios: ordered list of scenarios to evaluate + replicas: how many times each scenario is rolled out (D1: replication mode) + """ + + def __init__( + self, + num_envs: int, + scenarios: list[Scenario], + replicas: int = 1, + ) -> None: + if num_envs <= 0: + raise ValueError(f"num_envs must be positive, got {num_envs}.") + if replicas <= 0: + raise ValueError(f"replicas must be positive, got {replicas}.") + self.num_envs = num_envs + self.replicas = replicas + self._pending: deque[SlotTask] = deque( + SlotTask(scenario=sc, replica_idx=r) + for sc in scenarios + for r in range(replicas) + ) + self._slot_tasks: list[SlotTask | None] = [None] * num_envs + + # ----- queries ----- + + @property + def total_tasks(self) -> int: + """Total tasks scheduled across the whole run (used for sizing reports).""" + return len(self._slot_tasks) + len(self._pending) - sum( + 1 for t in self._slot_tasks if t is None + ) + + def active_slots(self) -> list[int]: + return [i for i, t in enumerate(self._slot_tasks) if t is not None] + + def current_task(self, slot: int) -> SlotTask: + """Task currently assigned to ``slot``. Raises if slot is idle.""" + task = self._slot_tasks[slot] + if task is None: + raise ValueError(f"Slot {slot} is idle — no task currently assigned.") + return task + + def is_done(self) -> bool: + return not self._pending and all(t is None for t in self._slot_tasks) + + # ----- mutations ----- + + def initialize(self) -> list[Scenario | None]: + """Pull tasks off the queue and assign to every slot. + + Returns one entry per slot: + * a Scenario for slots that received a task + * None for slots left idle (queue smaller than num_envs) + + The runner uses this list to construct the env.reset batch. Idle + slots can be padded with any placeholder; the runner ignores their + outcomes. + """ + initial: list[Scenario | None] = [] + for i in range(self.num_envs): + if self._pending: + task = self._pending.popleft() + self._slot_tasks[i] = task + initial.append(task.scenario) + else: + self._slot_tasks[i] = None + initial.append(None) + return initial + + def complete_slot(self, slot: int) -> SlotTask | None: + """Finish the current task on ``slot``, pull the next from the queue. + + Returns the new SlotTask if the queue had pending work, or None if + the queue is now empty (slot becomes idle). The runner uses the + returned task to call env.reset_slots([slot], [task.scenario]). + """ + if self._slot_tasks[slot] is None: + raise ValueError(f"Slot {slot} has no task to complete.") + if self._pending: + next_task = self._pending.popleft() + self._slot_tasks[slot] = next_task + return next_task + self._slot_tasks[slot] = None + return None diff --git a/roboeval/batched/tests/test_runner.py b/roboeval/batched/tests/test_runner.py new file mode 100644 index 0000000..661f242 --- /dev/null +++ b/roboeval/batched/tests/test_runner.py @@ -0,0 +1,384 @@ +"""Tests for BatchedEvalRunner — uses a deterministic in-memory mock adapter. + +Coverage: + * single-policy / single-scenario smoke + * num_envs == num_tasks (no refill) + * num_envs < num_tasks (refill required, scheduler exercises queue) + * num_envs > num_tasks (idle slots — placeholder action handling) + * multi-policy iteration + * replication mode (replicas > 1) and scenario_name #r suffix + * max_steps enforcement (force-terminate without env terminal) + * Ruleset pass/fail flows through to EpisodeResult + * EvalReport shape matches single-env EvalRunner expectations +""" + +from __future__ import annotations + +import unittest +from typing import Any + +from roboeval.batched.runner import BatchedEvalRunner +from roboeval.batched.types import BatchedStepOutcome +from roboeval.core import Ruleset, Scenario, require_metric, require_outcome + + +# ─── Mock vectorized environment ──────────────────────────────────────────── + + +class MockBatchedEnv: + """Deterministic mock that terminates per-scenario at a configurable step. + + Args: + num_envs: batch width + terminate_at: dict mapping scenario.name -> step count after which the + slot returns ``terminal=True``. Missing entries never terminate + (force-truncate via scenario.max_steps). + outcome_on_terminate: which ``outcome`` string to emit on terminal step + """ + + def __init__( + self, + num_envs: int, + terminate_at: dict[str, int] | None = None, + outcome_on_terminate: str = "terminated_success", + ) -> None: + self.num_envs = num_envs + self._terminate_at = terminate_at or {} + self._outcome_on_terminate = outcome_on_terminate + self._slot_step_counts: list[int] = [0] * num_envs + self._slot_scenarios: list[Scenario | None] = [None] * num_envs + self.step_calls = 0 + self.reset_calls = 0 + self.reset_slots_calls = 0 + + def reset(self, scenarios: list[Scenario]) -> list[dict[str, Any]]: + assert len(scenarios) == self.num_envs + self.reset_calls += 1 + self._slot_step_counts = [0] * self.num_envs + self._slot_scenarios = list(scenarios) + return [{"obs": [0.0], "slot_scenario": s.name} for s in scenarios] + + def step(self, actions: list[Any]) -> BatchedStepOutcome: + assert len(actions) == self.num_envs + self.step_calls += 1 + next_states = [] + outcomes = [] + failure_labels = [] + terminals = [] + metrics = [] + for i in range(self.num_envs): + self._slot_step_counts[i] += 1 + sc = self._slot_scenarios[i] + term_at = self._terminate_at.get(sc.name) if sc is not None else None + is_term = term_at is not None and self._slot_step_counts[i] >= term_at + next_states.append( + {"obs": [float(self._slot_step_counts[i])], "slot_scenario": sc.name} + ) + outcomes.append(self._outcome_on_terminate if is_term else "progress") + failure_labels.append("") + terminals.append(is_term) + metrics.append({"reward": 1.0, "episode_return": float(self._slot_step_counts[i])}) + return BatchedStepOutcome( + next_states=next_states, + outcomes=outcomes, + failure_labels=failure_labels, + terminals=terminals, + metrics=metrics, + ) + + def reset_slots( + self, slots: list[int], scenarios: list[Scenario] + ) -> list[dict[str, Any]]: + assert len(slots) == len(scenarios) + self.reset_slots_calls += 1 + new_states = [] + for slot, sc in zip(slots, scenarios): + self._slot_step_counts[slot] = 0 + self._slot_scenarios[slot] = sc + new_states.append({"obs": [0.0], "slot_scenario": sc.name}) + return new_states + + +# ─── Helpers ───────────────────────────────────────────────────────────────── + + +def constant_policy(action: int): + """Return a single-state policy that always emits the same action.""" + + def _policy(state): + return {"action": action, "debug_info": {}} + + _policy.__name__ = f"const_{action}" + return _policy + + +def batched_constant_policy(action: int, name: str | None = None): + """A natively-batched constant policy.""" + + def _batched(states): + return [{"action": action} for _ in states] + + _batched.__name__ = name or f"batched_const_{action}" + return _batched + + +def scenarios_named(*names: str, max_steps: int = 50) -> list[Scenario]: + return [Scenario(name=n, initial_state={"seed": i}, max_steps=max_steps) for i, n in enumerate(names)] + + +def from_single(fn): + from roboeval.batched.policy import from_single as _fs + + return _fs(fn) + + +# ─── Tests ─────────────────────────────────────────────────────────────────── + + +class TestBatchedRunnerSmoke(unittest.TestCase): + def test_single_policy_single_scenario_single_slot(self) -> None: + env = MockBatchedEnv(num_envs=1, terminate_at={"only": 3}) + report = BatchedEvalRunner( + policies=[from_single(constant_policy(0))], + scenarios=scenarios_named("only", max_steps=10), + environment=env, + ).run() + self.assertEqual(len(report.episodes), 1) + ep = report.episodes[0] + self.assertEqual(ep.scenario_name, "only") + self.assertEqual(ep.steps, 3) + self.assertEqual(ep.terminal_outcome, "terminated_success") + + def test_num_envs_equals_num_scenarios_no_refill(self) -> None: + env = MockBatchedEnv( + num_envs=3, terminate_at={"a": 2, "b": 3, "c": 4} + ) + report = BatchedEvalRunner( + policies=[from_single(constant_policy(0))], + scenarios=scenarios_named("a", "b", "c", max_steps=10), + environment=env, + ).run() + self.assertEqual(len(report.episodes), 3) + # No slot refill — only the bulk reset happens + self.assertEqual(env.reset_calls, 1) + self.assertEqual(env.reset_slots_calls, 0) + + +class TestBatchedRunnerRefill(unittest.TestCase): + def test_num_envs_less_than_scenarios_triggers_refill(self) -> None: + env = MockBatchedEnv( + num_envs=2, + terminate_at={n: 2 for n in ("a", "b", "c", "d", "e")}, + ) + report = BatchedEvalRunner( + policies=[from_single(constant_policy(0))], + scenarios=scenarios_named("a", "b", "c", "d", "e", max_steps=10), + environment=env, + ).run() + self.assertEqual(len(report.episodes), 5) + # Started with 2 in flight; 3 refills needed to drain queue + self.assertGreaterEqual(env.reset_slots_calls, 1) + # Every scenario name appears exactly once across episodes + seen = sorted(ep.scenario_name for ep in report.episodes) + self.assertEqual(seen, ["a", "b", "c", "d", "e"]) + + def test_idle_slots_when_more_envs_than_tasks(self) -> None: + env = MockBatchedEnv(num_envs=5, terminate_at={"a": 2, "b": 3}) + report = BatchedEvalRunner( + policies=[from_single(constant_policy(0))], + scenarios=scenarios_named("a", "b", max_steps=10), + environment=env, + ).run() + self.assertEqual(len(report.episodes), 2) + # 3 slots were idle the whole time; their outcomes should be ignored. + seen = sorted(ep.scenario_name for ep in report.episodes) + self.assertEqual(seen, ["a", "b"]) + + +class TestBatchedRunnerMultiPolicy(unittest.TestCase): + def test_two_policies_two_scenarios_produces_four_episodes(self) -> None: + env = MockBatchedEnv(num_envs=2, terminate_at={"a": 3, "b": 3}) + report = BatchedEvalRunner( + policies=[ + from_single(constant_policy(0)), + from_single(constant_policy(1)), + ], + scenarios=scenarios_named("a", "b", max_steps=10), + environment=env, + baseline_policy="const_0", + ).run() + self.assertEqual(len(report.episodes), 4) + self.assertIn("const_0", report.policy_summary) + self.assertIn("const_1", report.policy_summary) + # Each policy gets its own bulk reset + self.assertEqual(env.reset_calls, 2) + + +class TestBatchedRunnerReplication(unittest.TestCase): + def test_replicas_produce_suffixed_scenario_names(self) -> None: + env = MockBatchedEnv(num_envs=2, terminate_at={"a": 2, "b": 2}) + report = BatchedEvalRunner( + policies=[from_single(constant_policy(0))], + scenarios=scenarios_named("a", "b", max_steps=10), + environment=env, + replicas=3, + ).run() + # 2 scenarios * 3 replicas = 6 episodes + self.assertEqual(len(report.episodes), 6) + seen = sorted(ep.scenario_name for ep in report.episodes) + expected = sorted([f"a#r{i}" for i in range(3)] + [f"b#r{i}" for i in range(3)]) + self.assertEqual(seen, expected) + + def test_replicas_one_keeps_plain_scenario_names(self) -> None: + env = MockBatchedEnv(num_envs=2, terminate_at={"a": 2, "b": 2}) + report = BatchedEvalRunner( + policies=[from_single(constant_policy(0))], + scenarios=scenarios_named("a", "b"), + environment=env, + replicas=1, + ).run() + seen = sorted(ep.scenario_name for ep in report.episodes) + self.assertEqual(seen, ["a", "b"]) + + +class TestBatchedRunnerMaxSteps(unittest.TestCase): + def test_max_steps_force_terminates_without_env_terminal(self) -> None: + # terminate_at empty -> env never returns terminal=True + env = MockBatchedEnv(num_envs=1, terminate_at={}) + report = BatchedEvalRunner( + policies=[from_single(constant_policy(0))], + scenarios=scenarios_named("never_terminates", max_steps=4), + environment=env, + ).run() + self.assertEqual(len(report.episodes), 1) + ep = report.episodes[0] + self.assertEqual(ep.steps, 4) + self.assertEqual(ep.terminal_outcome, "max_steps_reached") + + def test_env_terminal_wins_over_max_steps(self) -> None: + env = MockBatchedEnv(num_envs=1, terminate_at={"a": 3}) + report = BatchedEvalRunner( + policies=[from_single(constant_policy(0))], + scenarios=scenarios_named("a", max_steps=10), + environment=env, + ).run() + ep = report.episodes[0] + self.assertEqual(ep.steps, 3) + self.assertEqual(ep.terminal_outcome, "terminated_success") + + +class TestBatchedRunnerRulesetFlow(unittest.TestCase): + def test_ruleset_failures_appear_in_episode_result(self) -> None: + env = MockBatchedEnv(num_envs=1, terminate_at={"x": 5}) + # Ruleset requires episode_return >= 100 — way above what 5 steps gives + report = BatchedEvalRunner( + policies=[from_single(constant_policy(0))], + scenarios=scenarios_named("x", max_steps=10), + environment=env, + ruleset=Ruleset([ + require_metric("episode_return", ">=", 100.0, name="needs_100"), + ]), + ).run() + ep = report.episodes[0] + self.assertFalse(ep.success) + self.assertEqual(ep.failure_label, "needs_100") + + def test_ruleset_success_flows_through(self) -> None: + env = MockBatchedEnv(num_envs=1, terminate_at={"x": 3}) + report = BatchedEvalRunner( + policies=[from_single(constant_policy(0))], + scenarios=scenarios_named("x", max_steps=10), + environment=env, + ruleset=Ruleset([require_outcome("terminated_success")]), + ).run() + ep = report.episodes[0] + self.assertTrue(ep.success) + + +class TestBatchedRunnerReportShape(unittest.TestCase): + def test_report_contains_expected_fields(self) -> None: + env = MockBatchedEnv( + num_envs=2, terminate_at={"a": 3, "b": 3} + ) + report = BatchedEvalRunner( + policies=[ + from_single(constant_policy(0)), + from_single(constant_policy(1)), + ], + scenarios=scenarios_named("a", "b"), + environment=env, + baseline_policy="const_0", + ).run() + # _build_report fills these in + self.assertEqual(report.baseline_policy, "const_0") + self.assertEqual(len(report.episodes), 4) + self.assertIsInstance(report.policy_summary, dict) + self.assertIsInstance(report.regressions, list) + self.assertIsInstance(report.improvements, list) + + def test_steprecord_logs_carry_per_slot_info(self) -> None: + env = MockBatchedEnv(num_envs=2, terminate_at={"a": 3, "b": 5}) + report = BatchedEvalRunner( + policies=[from_single(constant_policy(0))], + scenarios=scenarios_named("a", "b"), + environment=env, + ).run() + for ep in report.episodes: + for log in ep.logs: + self.assertIn("slot_scenario", log.next_state) + self.assertEqual(log.next_state["slot_scenario"], ep.scenario_name) + + +class TestBatchedRunnerNativeBatchedPolicy(unittest.TestCase): + def test_natively_batched_policy_runs(self) -> None: + env = MockBatchedEnv(num_envs=2, terminate_at={"a": 2, "b": 2}) + report = BatchedEvalRunner( + policies=[batched_constant_policy(0, name="native_v0")], + scenarios=scenarios_named("a", "b"), + environment=env, + ).run() + self.assertEqual(len(report.episodes), 2) + self.assertIn("native_v0", report.policy_summary) + + +class TestBatchedRunnerInputValidation(unittest.TestCase): + def test_no_policies_raises(self) -> None: + env = MockBatchedEnv(num_envs=1) + with self.assertRaisesRegex(ValueError, r"at least one policy"): + BatchedEvalRunner(policies=[], scenarios=scenarios_named("a"), environment=env) + + def test_no_scenarios_raises(self) -> None: + env = MockBatchedEnv(num_envs=1) + with self.assertRaisesRegex(ValueError, r"at least one scenario"): + BatchedEvalRunner( + policies=[from_single(constant_policy(0))], + scenarios=[], + environment=env, + ) + + def test_environment_without_num_envs_raises(self) -> None: + class Bad: + def reset(self, scenarios): return [] + def step(self, actions): return None + def reset_slots(self, slots, scenarios): return [] + + with self.assertRaisesRegex(TypeError, r"num_envs"): + BatchedEvalRunner( + policies=[from_single(constant_policy(0))], + scenarios=scenarios_named("a"), + environment=Bad(), + ) + + def test_zero_replicas_raises(self) -> None: + env = MockBatchedEnv(num_envs=1) + with self.assertRaisesRegex(ValueError, r"replicas must be positive"): + BatchedEvalRunner( + policies=[from_single(constant_policy(0))], + scenarios=scenarios_named("a"), + environment=env, + replicas=0, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/roboeval/batched/tests/test_scheduler.py b/roboeval/batched/tests/test_scheduler.py new file mode 100644 index 0000000..0578bc5 --- /dev/null +++ b/roboeval/batched/tests/test_scheduler.py @@ -0,0 +1,143 @@ +"""Tests for SlotScheduler — task assignment, refill, replication, idle handling.""" + +from __future__ import annotations + +import unittest + +from roboeval.batched.scheduler import SlotScheduler, SlotTask +from roboeval.core import Scenario + + +def _scenarios(*names: str) -> list[Scenario]: + return [Scenario(name=n, initial_state={"seed": i}, max_steps=5) for i, n in enumerate(names)] + + +class TestSchedulerInitialization(unittest.TestCase): + def test_rejects_zero_or_negative_num_envs(self) -> None: + with self.assertRaisesRegex(ValueError, r"num_envs must be positive"): + SlotScheduler(num_envs=0, scenarios=_scenarios("a")) + with self.assertRaisesRegex(ValueError, r"num_envs must be positive"): + SlotScheduler(num_envs=-2, scenarios=_scenarios("a")) + + def test_rejects_zero_replicas(self) -> None: + with self.assertRaisesRegex(ValueError, r"replicas must be positive"): + SlotScheduler(num_envs=4, scenarios=_scenarios("a"), replicas=0) + + def test_initialize_fills_slots_when_tasks_ge_slots(self) -> None: + sched = SlotScheduler(num_envs=2, scenarios=_scenarios("a", "b", "c")) + initial = sched.initialize() + self.assertEqual([sc.name for sc in initial], ["a", "b"]) + self.assertEqual(sched.active_slots(), [0, 1]) + self.assertFalse(sched.is_done()) + + def test_initialize_pads_with_none_when_tasks_lt_slots(self) -> None: + sched = SlotScheduler(num_envs=5, scenarios=_scenarios("a", "b")) + initial = sched.initialize() + self.assertEqual([sc.name if sc else None for sc in initial], ["a", "b", None, None, None]) + self.assertEqual(sched.active_slots(), [0, 1]) + + def test_initialize_with_zero_scenarios_is_legal(self) -> None: + sched = SlotScheduler(num_envs=4, scenarios=[]) + initial = sched.initialize() + self.assertEqual(initial, [None, None, None, None]) + self.assertTrue(sched.is_done()) + self.assertEqual(sched.active_slots(), []) + + +class TestSchedulerReplication(unittest.TestCase): + def test_three_scenarios_with_two_replicas_yields_six_tasks(self) -> None: + sched = SlotScheduler(num_envs=2, scenarios=_scenarios("a", "b", "c"), replicas=2) + initial = sched.initialize() + # First 2 tasks: (a, r=0), (a, r=1) + self.assertEqual(sched.current_task(0), SlotTask(scenario=initial[0], replica_idx=0)) + self.assertEqual(sched.current_task(1), SlotTask(scenario=initial[1], replica_idx=1)) + self.assertEqual(initial[0].name, "a") + self.assertEqual(initial[1].name, "a") # second replica of "a" + # 4 more tasks pending + sched.complete_slot(0) # -> (b, r=0) + sched.complete_slot(1) # -> (b, r=1) + sched.complete_slot(0) # -> (c, r=0) + sched.complete_slot(1) # -> (c, r=1) + # next two complete drain the queue + next_task = sched.complete_slot(0) + self.assertIsNone(next_task) + next_task = sched.complete_slot(1) + self.assertIsNone(next_task) + self.assertTrue(sched.is_done()) + + def test_replica_indices_increment_per_scenario(self) -> None: + sched = SlotScheduler(num_envs=1, scenarios=_scenarios("solo"), replicas=4) + sched.initialize() + seen_indices = [sched.current_task(0).replica_idx] + for _ in range(3): + task = sched.complete_slot(0) + self.assertIsNotNone(task) + seen_indices.append(task.replica_idx) + self.assertEqual(seen_indices, [0, 1, 2, 3]) + self.assertIsNone(sched.complete_slot(0)) + + +class TestSchedulerRefill(unittest.TestCase): + def test_complete_slot_pulls_next_pending(self) -> None: + sched = SlotScheduler(num_envs=2, scenarios=_scenarios("a", "b", "c", "d")) + sched.initialize() + refilled = sched.complete_slot(0) + self.assertIsNotNone(refilled) + self.assertEqual(refilled.scenario.name, "c") + self.assertEqual(sched.current_task(0).scenario.name, "c") + + def test_complete_slot_returns_none_when_queue_empty(self) -> None: + sched = SlotScheduler(num_envs=4, scenarios=_scenarios("a", "b")) + sched.initialize() + # Only 2 tasks, 4 slots — slots 2,3 are already idle + self.assertEqual(sched.active_slots(), [0, 1]) + self.assertIsNone(sched.complete_slot(0)) + self.assertEqual(sched.active_slots(), [1]) + self.assertIsNone(sched.complete_slot(1)) + self.assertEqual(sched.active_slots(), []) + self.assertTrue(sched.is_done()) + + def test_complete_slot_on_idle_slot_raises(self) -> None: + sched = SlotScheduler(num_envs=4, scenarios=_scenarios("a")) + sched.initialize() + with self.assertRaisesRegex(ValueError, r"Slot 3 has no task"): + sched.complete_slot(3) + + def test_current_task_on_idle_slot_raises(self) -> None: + sched = SlotScheduler(num_envs=4, scenarios=_scenarios("a")) + sched.initialize() + with self.assertRaisesRegex(ValueError, r"Slot 2 is idle"): + sched.current_task(2) + + +class TestSchedulerEndToEnd(unittest.TestCase): + def test_drains_queue_correctly_with_mixed_completion_order(self) -> None: + """Slots terminate out of order — scheduler must still drain the queue + without losing or duplicating tasks.""" + sched = SlotScheduler(num_envs=3, scenarios=_scenarios("a", "b", "c", "d", "e", "f", "g")) + sched.initialize() + # Started: slot0=a, slot1=b, slot2=c. Pending: d,e,f,g. + + seen = ["a", "b", "c"] # initial assignments + + # slot 2 finishes first + t = sched.complete_slot(2); seen.append(t.scenario.name) # d + # slot 0 finishes + t = sched.complete_slot(0); seen.append(t.scenario.name) # e + # slot 2 again + t = sched.complete_slot(2); seen.append(t.scenario.name) # f + # slot 1 finishes + t = sched.complete_slot(1); seen.append(t.scenario.name) # g + + # Now queue empty. Remaining slot completions return None. + self.assertIsNone(sched.complete_slot(2)) + self.assertIsNone(sched.complete_slot(0)) + self.assertIsNone(sched.complete_slot(1)) + self.assertTrue(sched.is_done()) + + # Every scenario ran exactly once + self.assertEqual(sorted(seen), ["a", "b", "c", "d", "e", "f", "g"]) + + +if __name__ == "__main__": + unittest.main() From bc48ee6235710b8dd070213ee5de562115be86b1 Mon Sep 17 00:00:00 2001 From: arryan0112 Date: Mon, 8 Jun 2026 14:32:43 +0530 Subject: [PATCH 4/6] Add BatchedGymnasiumEnvironmentAdapter for gym.vector.SyncVectorEnv - Wraps gym.vector.SyncVectorEnv as a BatchedEnvironmentAdapter; supports any class exposing env.envs (AsyncVectorEnv not yet supported) - Reuses the six default hooks from the single-env adapter for translation parity at the slot level - Handles Gymnasium 1.x NEXT_STEP autoreset; reset_slots overrides the implicit auto-reset by calling env.envs[i].reset(seed=...) for deterministic per-slot refills - Tests: construction, reset/step shape, per-slot episode_return accumulation, reset_slots determinism, action conversion, end-to-end against real CartPole-v1 with BatchedEvalRunner --- roboeval/integrations/gymnasium/__init__.py | 2 + .../integrations/gymnasium/batched_adapter.py | 309 +++++++++++ .../gymnasium/tests/test_batched_adapter.py | 482 ++++++++++++++++++ 3 files changed, 793 insertions(+) create mode 100644 roboeval/integrations/gymnasium/batched_adapter.py create mode 100644 roboeval/integrations/gymnasium/tests/test_batched_adapter.py diff --git a/roboeval/integrations/gymnasium/__init__.py b/roboeval/integrations/gymnasium/__init__.py index fb02cd9..dda7413 100644 --- a/roboeval/integrations/gymnasium/__init__.py +++ b/roboeval/integrations/gymnasium/__init__.py @@ -13,8 +13,10 @@ default_outcome_from_step, default_seed_from_scenario, ) +from .batched_adapter import BatchedGymnasiumEnvironmentAdapter __all__ = [ + "BatchedGymnasiumEnvironmentAdapter", "GymnasiumEnvironmentAdapter", "default_action_from_decision", "default_events_from_step", diff --git a/roboeval/integrations/gymnasium/batched_adapter.py b/roboeval/integrations/gymnasium/batched_adapter.py new file mode 100644 index 0000000..1f0a466 --- /dev/null +++ b/roboeval/integrations/gymnasium/batched_adapter.py @@ -0,0 +1,309 @@ +"""Batched Gymnasium ↔ roboeval adapter. + +Wraps ``gymnasium.vector.SyncVectorEnv`` (and any subclass exposing +``env.envs``) as a ``BatchedEnvironmentAdapter`` so ``BatchedEvalRunner`` can +drive N parallel Gymnasium environments in one step call. + +Translation parity with the single-env adapter +---------------------------------------------- +The six per-axis translation hooks are imported directly from the single-env +``adapter`` module so behavior matches byte-for-byte at the slot level: + + * ``default_observation_to_state`` + * ``default_action_from_decision`` + * ``default_outcome_from_step`` + * ``default_events_from_step`` + * ``default_seed_from_scenario`` + * ``default_options_from_scenario`` + +A user who customized hooks on the single-env adapter drops them in here +unchanged. + +Gymnasium 1.x autoreset +----------------------- +Default ``AutoresetMode.NEXT_STEP``: when slot i terminates on step N, the +returned ``obs[i]`` is the *terminal* observation. On step N+1, slot i has +been auto-reset internally to a fresh episode. Our runner intercepts the +terminal between those steps and calls ``reset_slots([i], [new_scenario])``, +which explicitly re-seeds slot i via ``env.envs[i].reset(seed=...)`` — that +overrides the implicit auto-reset with the deterministic, scenario-driven +one. ``SAME_STEP`` mode is not currently supported. + +AsyncVectorEnv is refused because ``env.envs`` is not directly addressable +across the subprocess boundary; supporting it requires ``env.call("reset", ...)`` +which has different return-value semantics. Future enhancement. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable + +import gymnasium as gym +import numpy as np + +from roboeval.batched.types import BatchedStepOutcome +from roboeval.core import Action, Scenario, State, to_serializable + +from .adapter import ( + default_action_from_decision, + default_events_from_step, + default_observation_to_state, + default_options_from_scenario, + default_outcome_from_step, + default_seed_from_scenario, +) + + +@dataclass +class BatchedGymnasiumEnvironmentAdapter: + """Wraps ``gym.vector.SyncVectorEnv`` as a ``BatchedEnvironmentAdapter``. + + The same six hooks as the single-env adapter are exposed; they're applied + per slot. Per-slot ``episode_return`` is tracked internally and reset on + each slot's terminal step. + + Parameters + ---------- + env : + ``gym.vector.SyncVectorEnv`` (or any class exposing ``env.envs`` as + the addressable list of per-slot envs). + name : + Display name for reports. + info_keys : + Optional allowlist for keys passed through to + ``info["gymnasium"]["raw_info"]`` per slot. + coerce_observations : + When ``True``, observations are coerced to JSON-safe Python types + before being returned (off by default; the runner coerces at write time). + """ + + env: gym.vector.VectorEnv + name: str = "batched_gymnasium_env" + + observation_to_state: Callable[[Any], dict[str, Any]] | None = None + action_from_decision: Callable[[Action], Any] | None = None + outcome_from_step: ( + Callable[[float, bool, bool, dict], tuple[str, str]] | None + ) = None + events_from_step: ( + Callable[[float, bool, bool, dict], list[str]] | None + ) = None + seed_from_scenario: Callable[[Scenario], int | None] | None = None + options_from_scenario: Callable[[Scenario], dict | None] | None = None + + info_keys: list[str] | None = None + coerce_observations: bool = False + + def __post_init__(self) -> None: + if not isinstance(self.env, gym.vector.VectorEnv): + raise TypeError( + "BatchedGymnasiumEnvironmentAdapter requires gym.vector.VectorEnv; " + f"got {type(self.env).__name__}." + ) + if not hasattr(self.env, "envs"): + raise TypeError( + "Vector env must expose an addressable .envs list " + "(SyncVectorEnv supports this; AsyncVectorEnv does not yet)." + ) + autoreset_mode = self.env.metadata.get("autoreset_mode") if hasattr(self.env, "metadata") else None + if autoreset_mode is not None and getattr(autoreset_mode, "value", str(autoreset_mode)) not in ( + "NextStep", "next-step", "NEXT_STEP" + ): + # We tolerate envs that don't declare autoreset_mode (older or custom + # vector wrappers). We refuse only when we KNOW it's SAME_STEP. + if "Same" in str(autoreset_mode): + raise NotImplementedError( + f"BatchedGymnasiumEnvironmentAdapter currently supports " + f"NEXT_STEP autoreset only; got {autoreset_mode}." + ) + + self.num_envs = self.env.num_envs + + if self.observation_to_state is None: + self.observation_to_state = default_observation_to_state + if self.action_from_decision is None: + self.action_from_decision = default_action_from_decision + if self.outcome_from_step is None: + self.outcome_from_step = default_outcome_from_step + if self.events_from_step is None: + self.events_from_step = default_events_from_step + if self.seed_from_scenario is None: + self.seed_from_scenario = default_seed_from_scenario + if self.options_from_scenario is None: + self.options_from_scenario = default_options_from_scenario + + self._episode_returns: list[float] = [0.0] * self.num_envs + + # --- BatchedEnvironmentAdapter Protocol --------------------------------- + + def reset(self, scenarios: list[Scenario]) -> list[State]: + if len(scenarios) != self.num_envs: + raise ValueError( + f"reset() requires {self.num_envs} scenarios (num_envs), " + f"got {len(scenarios)}." + ) + # Build a per-slot seed list. Gymnasium VectorEnv.reset() accepts + # ``seed: int | list[int] | None``. If any slot has no seed we pass + # None (relinquishing determinism) — otherwise we forward the full list. + seeds = [self.seed_from_scenario(sc) for sc in scenarios] + seed_arg: int | list[int] | None + if all(s is None for s in seeds): + seed_arg = None + else: + seed_arg = [int(s) if s is not None else 0 for s in seeds] + + # gym.vector.reset() does not accept per-slot options; if any scenario + # specifies options we apply them per-slot by calling env.envs[i].reset. + per_slot_options = [self.options_from_scenario(sc) for sc in scenarios] + + obs, _info = self.env.reset(seed=seed_arg) + self._episode_returns = [0.0] * self.num_envs + + states = [self._state_from_obs(obs[i]) for i in range(self.num_envs)] + # Re-reset any slot that requested non-None options + for i, opts in enumerate(per_slot_options): + if opts is not None: + slot_seed = seeds[i] if seeds[i] is not None else None + obs_i, _info_i = self.env.envs[i].reset(seed=slot_seed, options=opts) + states[i] = self._state_from_obs(obs_i) + return states + + def step(self, actions: list[Action]) -> BatchedStepOutcome: + if len(actions) != self.num_envs: + raise ValueError( + f"step() requires {self.num_envs} actions (num_envs), " + f"got {len(actions)}." + ) + gym_actions = [self.action_from_decision(a) for a in actions] + action_array = self._actions_to_array(gym_actions) + + obs, rewards, terminateds, truncateds, infos = self.env.step(action_array) + + next_states: list[State] = [] + outcomes: list[str] = [] + failure_labels: list[str] = [] + terminals: list[bool] = [] + metrics: list[dict[str, float | int] | None] = [] + events: list[list[str] | None] = [] + info_list: list[dict[str, object] | None] = [] + + for i in range(self.num_envs): + reward_i = float(rewards[i]) + term_i = bool(terminateds[i]) + trunc_i = bool(truncateds[i]) + slot_info = _per_slot_info(infos, i, self.num_envs, self.info_keys) + self._episode_returns[i] += reward_i + + next_states.append(self._state_from_obs(obs[i])) + outcome, failure_label = self.outcome_from_step(reward_i, term_i, trunc_i, slot_info) + outcomes.append(outcome) + failure_labels.append(failure_label) + terminals.append(term_i or trunc_i) + metrics.append({ + "reward": reward_i, + "episode_return": float(self._episode_returns[i]), + }) + events.append(self.events_from_step(reward_i, term_i, trunc_i, slot_info)) + info_list.append({ + "gymnasium": { + "terminated": term_i, + "truncated": trunc_i, + "raw_info": to_serializable(slot_info), + "slot": i, + } + }) + + # Reset return so the next episode for this slot starts at 0. + # The runner's reset_slots() call comes between steps; this keeps + # state consistent if the runner force-truncates via max_steps too. + if term_i or trunc_i: + self._episode_returns[i] = 0.0 + + return BatchedStepOutcome( + next_states=next_states, + outcomes=outcomes, + failure_labels=failure_labels, + terminals=terminals, + metrics=metrics, + events=events, + info=info_list, + ) + + def reset_slots( + self, slots: list[int], scenarios: list[Scenario] + ) -> list[State]: + if len(slots) != len(scenarios): + raise ValueError( + f"reset_slots: slots and scenarios must be same length, " + f"got {len(slots)} and {len(scenarios)}." + ) + new_states: list[State] = [] + for slot, sc in zip(slots, scenarios): + seed = self.seed_from_scenario(sc) + options = self.options_from_scenario(sc) + obs, _info = self.env.envs[slot].reset(seed=seed, options=options) + self._episode_returns[slot] = 0.0 + new_states.append(self._state_from_obs(obs)) + return new_states + + def close(self) -> None: + """Forward close to the wrapped env. Optional — not part of the Protocol.""" + close = getattr(self.env, "close", None) + if callable(close): + close() + + # --- Internals ---------------------------------------------------------- + + def _state_from_obs(self, obs: Any) -> State: + state = self.observation_to_state(obs) + if self.coerce_observations: + state = to_serializable(state) + return state + + def _actions_to_array(self, actions: list[Any]) -> Any: + """Stack per-slot actions into the shape gym.vector expects. + + For Discrete actions (int) -> ndarray shape (num_envs,). + For Box actions (np.ndarray) -> ndarray shape (num_envs, *action_dim). + For Dict/Tuple action spaces, the array form differs; the safest path + is to let numpy figure it out via np.array(actions), falling back to + a Python list for spaces that error on that conversion. + """ + try: + return np.array(actions) + except (TypeError, ValueError): + return actions + + +def _per_slot_info( + infos: dict[str, Any], + slot: int, + num_envs: int, + info_keys: list[str] | None, +) -> dict[str, Any]: + """Extract the slot-specific entries from a batched info dict. + + Gymnasium's batched info has per-key arrays of length ``num_envs`` for + most entries. ``_final_observation`` / ``final_observation`` (when + present) are skipped — they're handled separately by the autoreset + contract. Scalar entries (rare) are forwarded unchanged. + """ + slot_info: dict[str, Any] = {} + keys = info_keys if info_keys is not None else list(infos.keys()) + for key in keys: + if key in ("final_observation", "_final_observation"): + continue + if key not in infos: + continue + value = infos[key] + # Strings/bytes have __len__ too but are scalar info; don't index into them. + if hasattr(value, "__len__") and not isinstance(value, (str, bytes)): + try: + if len(value) == num_envs: + slot_info[key] = value[slot] + continue + except TypeError: + pass + slot_info[key] = value + return slot_info diff --git a/roboeval/integrations/gymnasium/tests/test_batched_adapter.py b/roboeval/integrations/gymnasium/tests/test_batched_adapter.py new file mode 100644 index 0000000..af91007 --- /dev/null +++ b/roboeval/integrations/gymnasium/tests/test_batched_adapter.py @@ -0,0 +1,482 @@ +"""Tests for BatchedGymnasiumEnvironmentAdapter. + +Covers adapter-level behavior (state shape, action conversion, reset_slots, +per-slot return tracking, info filtering) and end-to-end integration with +BatchedEvalRunner on a real CartPole-v1 vector env. +""" + +from __future__ import annotations + +import unittest +from typing import Any + +import gymnasium as gym +import numpy as np + +from roboeval.batched.runner import BatchedEvalRunner +from roboeval.batched.types import BatchedStepOutcome +from roboeval.core import Ruleset, Scenario, require_metric, require_outcome +from roboeval.integrations.gymnasium.batched_adapter import ( + BatchedGymnasiumEnvironmentAdapter, + _per_slot_info, +) + + +def _cartpole_vector(num_envs: int) -> gym.vector.SyncVectorEnv: + return gym.vector.SyncVectorEnv( + [lambda: gym.make("CartPole-v1") for _ in range(num_envs)] + ) + + +def _seeded_scenarios(*seeds: int, max_steps: int = 200) -> list[Scenario]: + return [Scenario(f"seed_{s}", {"seed": s}, max_steps=max_steps) for s in seeds] + + +# ─── Construction + input validation ──────────────────────────────────────── + + +class TestBatchedAdapterConstruction(unittest.TestCase): + def test_rejects_non_vector_env(self) -> None: + single_env = gym.make("CartPole-v1") + with self.assertRaisesRegex(TypeError, r"requires gym.vector.VectorEnv"): + BatchedGymnasiumEnvironmentAdapter(env=single_env) + single_env.close() + + def test_num_envs_attribute_set_from_vector_env(self) -> None: + env = _cartpole_vector(4) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + self.assertEqual(adapter.num_envs, 4) + finally: + env.close() + + +# ─── reset() ──────────────────────────────────────────────────────────────── + + +class TestBatchedAdapterReset(unittest.TestCase): + def test_reset_returns_num_envs_states(self) -> None: + env = _cartpole_vector(3) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + states = adapter.reset(_seeded_scenarios(1, 2, 3)) + self.assertEqual(len(states), 3) + for s in states: + self.assertIn("observation", s) + self.assertEqual(len(s["observation"]), 4) + finally: + env.close() + + def test_reset_requires_num_envs_scenarios(self) -> None: + env = _cartpole_vector(3) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + with self.assertRaisesRegex(ValueError, r"3 scenarios"): + adapter.reset(_seeded_scenarios(1, 2)) + finally: + env.close() + + def test_seeded_reset_is_deterministic(self) -> None: + # Two adapters, same seeds → same initial obs (CartPole is deterministic given a seed). + env_a = _cartpole_vector(2) + env_b = _cartpole_vector(2) + try: + adapter_a = BatchedGymnasiumEnvironmentAdapter(env=env_a) + adapter_b = BatchedGymnasiumEnvironmentAdapter(env=env_b) + states_a = adapter_a.reset(_seeded_scenarios(42, 43)) + states_b = adapter_b.reset(_seeded_scenarios(42, 43)) + np.testing.assert_array_equal(states_a[0]["observation"], states_b[0]["observation"]) + np.testing.assert_array_equal(states_a[1]["observation"], states_b[1]["observation"]) + finally: + env_a.close() + env_b.close() + + def test_reset_zeroes_episode_returns(self) -> None: + env = _cartpole_vector(2) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + adapter.reset(_seeded_scenarios(1, 2)) + # Step a few times to accumulate returns + for _ in range(3): + adapter.step([0, 0]) + self.assertGreater(sum(adapter._episode_returns), 0.0) + adapter.reset(_seeded_scenarios(1, 2)) + self.assertEqual(adapter._episode_returns, [0.0, 0.0]) + finally: + env.close() + + +# ─── step() ───────────────────────────────────────────────────────────────── + + +class TestBatchedAdapterStep(unittest.TestCase): + def test_step_returns_batched_step_outcome_with_correct_shape(self) -> None: + env = _cartpole_vector(2) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + adapter.reset(_seeded_scenarios(1, 2)) + outcome = adapter.step([0, 1]) + self.assertIsInstance(outcome, BatchedStepOutcome) + self.assertEqual(outcome.num_envs, 2) + self.assertEqual(len(outcome.next_states), 2) + self.assertEqual(len(outcome.outcomes), 2) + self.assertEqual(len(outcome.terminals), 2) + self.assertEqual(len(outcome.metrics), 2) + finally: + env.close() + + def test_step_requires_num_envs_actions(self) -> None: + env = _cartpole_vector(3) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + adapter.reset(_seeded_scenarios(1, 2, 3)) + with self.assertRaisesRegex(ValueError, r"3 actions"): + adapter.step([0]) + finally: + env.close() + + def test_per_slot_episode_return_accumulates(self) -> None: + env = _cartpole_vector(2) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + adapter.reset(_seeded_scenarios(1, 2)) + # CartPole rewards 1 per step until terminal. + outcome = adapter.step([0, 0]) + outcome = adapter.step([0, 0]) + outcome = adapter.step([0, 0]) + # After 3 steps (no terminal), episode_return == 3 + self.assertEqual(outcome.metrics[0]["episode_return"], 3.0) + self.assertEqual(outcome.metrics[1]["episode_return"], 3.0) + finally: + env.close() + + def test_terminal_step_records_terminated(self) -> None: + env = _cartpole_vector(1) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + adapter.reset([Scenario("only", {"seed": 42}, max_steps=600)]) + # Always push left — pole will fall + terminated = False + for _ in range(600): + outcome = adapter.step([0]) + if outcome.terminals[0]: + terminated = True + self.assertEqual(outcome.info[0]["gymnasium"]["terminated"], True) + self.assertEqual(outcome.info[0]["gymnasium"]["truncated"], False) + self.assertIn("episode_terminated", outcome.events[0]) + break + self.assertTrue(terminated, "expected CartPole to terminate within 600 steps with action=0") + finally: + env.close() + + def test_info_namespaces_gymnasium_payload(self) -> None: + env = _cartpole_vector(2) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + adapter.reset(_seeded_scenarios(1, 2)) + outcome = adapter.step([0, 1]) + for slot in (0, 1): + self.assertIn("gymnasium", outcome.info[slot]) + gym_info = outcome.info[slot]["gymnasium"] + self.assertIn("terminated", gym_info) + self.assertIn("truncated", gym_info) + self.assertIn("raw_info", gym_info) + self.assertEqual(gym_info["slot"], slot) + finally: + env.close() + + def test_metrics_contain_reward_and_episode_return(self) -> None: + env = _cartpole_vector(2) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + adapter.reset(_seeded_scenarios(1, 2)) + outcome = adapter.step([0, 1]) + for slot in (0, 1): + self.assertIn("reward", outcome.metrics[slot]) + self.assertIn("episode_return", outcome.metrics[slot]) + finally: + env.close() + + +# ─── reset_slots() ────────────────────────────────────────────────────────── + + +class TestBatchedAdapterResetSlots(unittest.TestCase): + def test_reset_slots_returns_one_state_per_slot(self) -> None: + env = _cartpole_vector(3) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + adapter.reset(_seeded_scenarios(1, 2, 3)) + new = adapter.reset_slots( + [0, 2], _seeded_scenarios(100, 200) + ) + self.assertEqual(len(new), 2) + self.assertIn("observation", new[0]) + self.assertEqual(len(new[0]["observation"]), 4) + finally: + env.close() + + def test_reset_slots_zeroes_episode_return_for_those_slots(self) -> None: + env = _cartpole_vector(3) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + adapter.reset(_seeded_scenarios(1, 2, 3)) + for _ in range(5): + adapter.step([0, 0, 0]) + # All slots had accumulated returns + before = list(adapter._episode_returns) + self.assertTrue(all(r > 0 for r in before)) + adapter.reset_slots([1], _seeded_scenarios(999)) + after = list(adapter._episode_returns) + self.assertEqual(after[1], 0.0) + self.assertEqual(after[0], before[0]) + self.assertEqual(after[2], before[2]) + finally: + env.close() + + def test_reset_slots_with_seed_is_deterministic(self) -> None: + env_a = _cartpole_vector(2) + env_b = _cartpole_vector(2) + try: + adapter_a = BatchedGymnasiumEnvironmentAdapter(env=env_a) + adapter_b = BatchedGymnasiumEnvironmentAdapter(env=env_b) + adapter_a.reset(_seeded_scenarios(1, 2)) + adapter_b.reset(_seeded_scenarios(1, 2)) + new_a = adapter_a.reset_slots([0], _seeded_scenarios(777)) + new_b = adapter_b.reset_slots([0], _seeded_scenarios(777)) + np.testing.assert_array_equal(new_a[0]["observation"], new_b[0]["observation"]) + finally: + env_a.close() + env_b.close() + + def test_reset_slots_validates_length_mismatch(self) -> None: + env = _cartpole_vector(2) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + adapter.reset(_seeded_scenarios(1, 2)) + with self.assertRaisesRegex(ValueError, r"same length"): + adapter.reset_slots([0, 1], _seeded_scenarios(99)) + finally: + env.close() + + +# ─── Action conversion ────────────────────────────────────────────────────── + + +class TestBatchedAdapterActionConversion(unittest.TestCase): + def test_int_actions_become_ndarray(self) -> None: + env = _cartpole_vector(3) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + adapter.reset(_seeded_scenarios(1, 2, 3)) + adapter.step([0, 1, 0]) + finally: + env.close() + + def test_custom_action_from_decision_runs(self) -> None: + # Translate "left"/"right" → 0/1 via hook + def custom_action(decision_action: Any) -> int: + return {"left": 0, "right": 1}[decision_action] + + env = _cartpole_vector(2) + try: + adapter = BatchedGymnasiumEnvironmentAdapter( + env=env, action_from_decision=custom_action + ) + adapter.reset(_seeded_scenarios(1, 2)) + outcome = adapter.step(["left", "right"]) + self.assertEqual(outcome.num_envs, 2) + finally: + env.close() + + +# ─── _per_slot_info helper ────────────────────────────────────────────────── + + +class TestPerSlotInfo(unittest.TestCase): + def test_extracts_per_key_arrays(self) -> None: + infos = {"score": np.array([10, 20, 30])} + result = _per_slot_info(infos, slot=1, num_envs=3, info_keys=None) + self.assertEqual(result, {"score": 20}) + + def test_skips_final_observation_keys(self) -> None: + infos = { + "final_observation": np.array([1, 2, 3]), + "_final_observation": np.array([True, False, False]), + "score": np.array([10, 20, 30]), + } + result = _per_slot_info(infos, slot=0, num_envs=3, info_keys=None) + self.assertNotIn("final_observation", result) + self.assertNotIn("_final_observation", result) + self.assertEqual(result, {"score": 10}) + + def test_info_keys_allowlist_filters(self) -> None: + infos = { + "score": np.array([10, 20]), + "big_blob": np.array([[1, 2], [3, 4]]), + } + result = _per_slot_info(infos, slot=0, num_envs=2, info_keys=["score"]) + self.assertEqual(result, {"score": 10}) + + def test_scalar_info_passed_through(self) -> None: + infos = {"env_version": "v1"} + result = _per_slot_info(infos, slot=0, num_envs=2, info_keys=None) + self.assertEqual(result, {"env_version": "v1"}) + + def test_empty_info_returns_empty(self) -> None: + result = _per_slot_info({}, slot=0, num_envs=4, info_keys=None) + self.assertEqual(result, {}) + + +# ─── End-to-end: BatchedEvalRunner + real CartPole vector env ─────────────── + + +def _balance_policy(state: dict) -> dict: + """Push left if pole leans left.""" + pole_angle = float(state["observation"][2]) + return {"action": 0 if pole_angle < 0 else 1, "debug_info": {"version": "balance"}} + + +def _always_left(state: dict) -> dict: + return {"action": 0, "debug_info": {"version": "left_only"}} + + +def _always_right(state: dict) -> dict: + return {"action": 1, "debug_info": {"version": "right_only"}} + + +class TestBatchedRunnerWithRealCartPole(unittest.TestCase): + def test_three_policies_four_scenarios_eight_envs_smokes(self) -> None: + """The headline integration test: 3 policies × 4 scenarios, 8-env batch. + Verifies refill, scheduler, episode termination, and report assembly.""" + env = gym.vector.SyncVectorEnv([lambda: gym.make("CartPole-v1") for _ in range(8)]) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env, name="cartpole_x8") + scenarios = _seeded_scenarios(1, 2, 42, 100, max_steps=200) + from roboeval.batched.policy import from_single + report = BatchedEvalRunner( + policies=[ + from_single(_balance_policy), + from_single(_always_left), + from_single(_always_right), + ], + scenarios=scenarios, + environment=adapter, + ruleset=Ruleset([require_metric("episode_return", ">=", 50.0, name="balance_50")]), + baseline_policy="_balance_policy", + ).run() + + # 3 policies × 4 scenarios = 12 episodes + self.assertEqual(len(report.episodes), 12) + # All three policies are present in the summary + self.assertIn("_balance_policy", report.policy_summary) + self.assertIn("_always_left", report.policy_summary) + self.assertIn("_always_right", report.policy_summary) + # Every episode should have steps > 0 + for ep in report.episodes: + self.assertGreater(ep.steps, 0) + finally: + env.close() + + def test_replication_yields_replica_suffixed_scenario_names(self) -> None: + env = gym.vector.SyncVectorEnv([lambda: gym.make("CartPole-v1") for _ in range(4)]) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + from roboeval.batched.policy import from_single + report = BatchedEvalRunner( + policies=[from_single(_balance_policy)], + scenarios=_seeded_scenarios(1, 2, max_steps=50), + environment=adapter, + replicas=3, + ruleset=Ruleset([require_outcome("terminated_failure")]), + ).run() + # 2 scenarios × 3 replicas = 6 episodes + self.assertEqual(len(report.episodes), 6) + names = sorted(ep.scenario_name for ep in report.episodes) + self.assertEqual( + names, + sorted([f"seed_1#r{i}" for i in range(3)] + [f"seed_2#r{i}" for i in range(3)]), + ) + finally: + env.close() + + def test_max_steps_truncates_when_pole_does_not_fall(self) -> None: + # The balance policy on CartPole should easily exceed 30 steps for several seeds. + env = gym.vector.SyncVectorEnv([lambda: gym.make("CartPole-v1") for _ in range(2)]) + try: + adapter = BatchedGymnasiumEnvironmentAdapter(env=env) + from roboeval.batched.policy import from_single + report = BatchedEvalRunner( + policies=[from_single(_balance_policy)], + scenarios=_seeded_scenarios(1, 2, max_steps=10), + environment=adapter, + ruleset=Ruleset([require_outcome("terminated_success")]), + ).run() + # If the policy keeps the pole up for 10 steps, the episode is truncated + # by max_steps; runner labels it max_steps_reached. + outcomes = {ep.terminal_outcome for ep in report.episodes} + # At least one of the seeds should hit max_steps_reached with this policy + self.assertTrue( + any(o == "max_steps_reached" for o in outcomes) + or any(o.startswith("terminated") for o in outcomes), + f"unexpected outcomes set: {outcomes}", + ) + finally: + env.close() + + def test_batched_and_single_env_report_shapes_match(self) -> None: + """Cross-check: BatchedEvalRunner and EvalRunner produce report objects + with the same top-level field set and types.""" + from roboeval.batched.policy import from_single + from roboeval.integrations.gymnasium.adapter import GymnasiumEnvironmentAdapter + from roboeval.runner import EvalRunner + + # Single-env path + single_env = gym.make("CartPole-v1") + single_adapter = GymnasiumEnvironmentAdapter(env=single_env) + single_report = EvalRunner( + policies=[_balance_policy], + scenarios=_seeded_scenarios(42, max_steps=20), + environment=single_adapter, + ruleset=Ruleset([require_metric("episode_return", ">=", 10.0)]), + ).run() + single_env.close() + + # Batched path + batched_env = gym.vector.SyncVectorEnv([lambda: gym.make("CartPole-v1") for _ in range(1)]) + try: + batched_adapter = BatchedGymnasiumEnvironmentAdapter(env=batched_env) + batched_report = BatchedEvalRunner( + policies=[from_single(_balance_policy)], + scenarios=_seeded_scenarios(42, max_steps=20), + environment=batched_adapter, + ruleset=Ruleset([require_metric("episode_return", ">=", 10.0)]), + ).run() + + # Top-level field types match + for attr in ("policy_summary", "episodes", "regressions", "improvements", "baseline_policy"): + self.assertEqual( + type(getattr(single_report, attr)), + type(getattr(batched_report, attr)), + f"field {attr} type mismatch", + ) + # Episode count matches + self.assertEqual(len(single_report.episodes), len(batched_report.episodes)) + # Episode result shape matches + s_ep = single_report.episodes[0] + b_ep = batched_report.episodes[0] + for attr in ( + "scenario_name", "policy_version", "success", + "terminal_outcome", "failure_label", "steps", + ): + self.assertEqual( + type(getattr(s_ep, attr)), + type(getattr(b_ep, attr)), + f"episode attr {attr} type mismatch", + ) + finally: + batched_env.close() + + +if __name__ == "__main__": + unittest.main() From 831ebc27c32ed74dfb35fd5ffed6f8284523365d Mon Sep 17 00:00:00 2001 From: arryan0112 Date: Mon, 8 Jun 2026 14:33:51 +0530 Subject: [PATCH 5/6] Add BatchedIsaacEnvironmentAdapter with wrapper-aware num_envs/device lookup - Reads the full Isaac GPU batch (no batch_index=0 slicing); stacks per-slot actions into (N, action_dim) tensors and fans (N,) reward/terminal tensors back into per-slot lists - Selective reset via three-tier fallback: env.reset(env_ids=...), env.unwrapped._reset_idx(env_ids), or full env.reset() with a one-time warning - num_envs and device lookups fall back to env.unwrapped to handle gym.make() wrappers (OrderEnforcing, TimeLimit) that don't proxy these attributes - Tests: construction, reset/step shape, per-slot terminals, episode_return tracking, action stacking, selective reset paths, end-to-end with BatchedEvalRunner against a mock that uses real torch tensors; regression test for the wrapper-aware num_envs lookup --- roboeval/integrations/isaac/__init__.py | 2 + .../integrations/isaac/batched_adapter.py | 444 +++++++++++++ .../isaac/tests/test_batched_adapter.py | 601 ++++++++++++++++++ 3 files changed, 1047 insertions(+) create mode 100644 roboeval/integrations/isaac/batched_adapter.py create mode 100644 roboeval/integrations/isaac/tests/test_batched_adapter.py diff --git a/roboeval/integrations/isaac/__init__.py b/roboeval/integrations/isaac/__init__.py index 0a52f58..37359b6 100644 --- a/roboeval/integrations/isaac/__init__.py +++ b/roboeval/integrations/isaac/__init__.py @@ -36,8 +36,10 @@ default_seed_from_scenario, tensor_to_numpy, ) +from .batched_adapter import BatchedIsaacEnvironmentAdapter __all__ = [ + "BatchedIsaacEnvironmentAdapter", "IsaacEnvironmentAdapter", "default_action_from_decision", "default_events_from_step", diff --git a/roboeval/integrations/isaac/batched_adapter.py b/roboeval/integrations/isaac/batched_adapter.py new file mode 100644 index 0000000..1560912 --- /dev/null +++ b/roboeval/integrations/isaac/batched_adapter.py @@ -0,0 +1,444 @@ +"""Batched Isaac Lab ↔ roboeval adapter. + +Wraps a vectorized Isaac Lab env (``ManagerBasedRLEnv`` or the gymnasium-style +wrapper around it) as a ``BatchedEnvironmentAdapter`` so ``BatchedEvalRunner`` +can read all N parallel slots — the way Isaac actually wants to be used. + +Difference from ``IsaacEnvironmentAdapter`` (single-env) +-------------------------------------------------------- +The single-env adapter takes ``batch_index=0`` and slices, discarding the +other 99% of the GPU's work. This batched adapter reads every slot: + + * ``reset(scenarios)`` — bulk reset all num_envs slots + * ``step(actions)`` — stack per-slot actions into a (N, *) tensor, + step the env, fan the (N,) reward/terminal + tensors back into per-slot lists + * ``reset_slots(slots, scs)`` — selective reset using ``env_ids`` if the + underlying env supports it, falling back to + a full reset with a warning + +The translation hooks (``observation_to_state``, ``outcome_from_step``, etc.) +match the single-env Isaac adapter so user customizations port over unchanged. + +Per-slot seeding limitation +--------------------------- +Isaac Lab's reset accepts a single ``seed`` — there's no public per-env seed +API. ``reset(scenarios)`` uses the first scenario's seed as the global seed +and proceeds. For deterministic per-replica variance, run the same scenario +with different reset seeds across separate ``run()`` calls or use Isaac's +domain randomization config. + +Selective reset +--------------- +Isaac Lab's ``ManagerBasedRLEnv`` exposes ``_reset_idx(env_ids)`` for in-place +per-env reset. The gym wrapper sometimes does too. We try, in order: + + 1. ``env.reset(env_ids=...)`` (gym wrapper, if implemented) + 2. ``env.unwrapped._reset_idx(env_ids)`` (ManagerBasedRLEnv direct) + 3. Full ``env.reset()`` (fallback — warns once) + +In all cases we re-read the obs to capture the post-reset state for the +just-reset slots. +""" + +from __future__ import annotations + +import warnings +from dataclasses import dataclass +from typing import Any, Callable + +from roboeval.batched.types import BatchedStepOutcome +from roboeval.core import Action, Scenario, State, to_serializable + +from .adapter import ( + default_action_from_decision, + default_events_from_step, + default_observation_to_state, + default_options_from_scenario, + default_outcome_from_step, + default_seed_from_scenario, + tensor_to_numpy, +) + + +@dataclass +class BatchedIsaacEnvironmentAdapter: + """Wraps a vectorized Isaac Lab env as a ``BatchedEnvironmentAdapter``. + + Parameters + ---------- + env : + The Isaac Lab vectorized env. Must expose ``num_envs``. + name : + Display name for reports. + observation_to_state, action_from_decision, outcome_from_step, + events_from_step, seed_from_scenario, options_from_scenario : + Override hooks. ``None`` selects the single-env Isaac adapter's + defaults. + info_keys : + Optional allowlist for keys passed through to + ``info["isaac"]["raw_info"]`` per slot. + """ + + env: Any + name: str = "batched_isaac_env" + + observation_to_state: Callable[[Any], dict[str, Any]] | None = None + action_from_decision: Callable[[Action], Any] | None = None + outcome_from_step: ( + Callable[[float, bool, bool, dict], tuple[str, str]] | None + ) = None + events_from_step: ( + Callable[[float, bool, bool, dict], list[str]] | None + ) = None + seed_from_scenario: Callable[[Scenario], int | None] | None = None + options_from_scenario: Callable[[Scenario], dict | None] | None = None + + info_keys: list[str] | None = None + + def __post_init__(self) -> None: + # Gymnasium wrappers (OrderEnforcing, TimeLimit, etc.) wrap Isaac Lab + # envs after gym.make(...) and don't proxy num_envs. Fall back to + # env.unwrapped — discovered while validating against real Isaac. + num_envs = getattr(self.env, "num_envs", None) + if num_envs is None: + unwrapped = getattr(self.env, "unwrapped", None) + num_envs = getattr(unwrapped, "num_envs", None) if unwrapped is not None else None + if num_envs is None: + raise TypeError( + "BatchedIsaacEnvironmentAdapter requires the wrapped env to " + "expose .num_envs (checked self.env and self.env.unwrapped). " + "Got an env with no num_envs attribute." + ) + if num_envs <= 0: + raise ValueError(f"env.num_envs must be positive, got {num_envs}.") + self.num_envs = int(num_envs) + + if self.observation_to_state is None: + self.observation_to_state = default_observation_to_state + if self.action_from_decision is None: + self.action_from_decision = default_action_from_decision + if self.outcome_from_step is None: + self.outcome_from_step = default_outcome_from_step + if self.events_from_step is None: + self.events_from_step = default_events_from_step + if self.seed_from_scenario is None: + self.seed_from_scenario = default_seed_from_scenario + if self.options_from_scenario is None: + self.options_from_scenario = default_options_from_scenario + + self._episode_returns: list[float] = [0.0] * self.num_envs + self._device: str | None = None + self._warned_full_reset = False + + # ── BatchedEnvironmentAdapter Protocol ────────────────────────────────── + + def reset(self, scenarios: list[Scenario]) -> list[State]: + if len(scenarios) != self.num_envs: + raise ValueError( + f"reset() requires {self.num_envs} scenarios (num_envs), " + f"got {len(scenarios)}." + ) + seeds = [self.seed_from_scenario(sc) for sc in scenarios] + seed = next((s for s in seeds if s is not None), None) + options = self.options_from_scenario(scenarios[0]) + + try: + obs, _info = self.env.reset(seed=seed, options=options) + except TypeError: + obs, _info = self.env.reset(seed=seed) + self._episode_returns = [0.0] * self.num_envs + return [self._state_from_obs_slot(obs, i) for i in range(self.num_envs)] + + def step(self, actions: list[Action]) -> BatchedStepOutcome: + if len(actions) != self.num_envs: + raise ValueError( + f"step() requires {self.num_envs} actions (num_envs), " + f"got {len(actions)}." + ) + decision_actions = [self.action_from_decision(a) for a in actions] + batched_action = self._stack_actions(decision_actions) + obs, reward, terminated, truncated, info = self.env.step(batched_action) + + rewards_np = tensor_to_numpy(reward) + terminateds_np = tensor_to_numpy(terminated) + truncateds_np = tensor_to_numpy(truncated) + + next_states: list[State] = [] + outcomes: list[str] = [] + failure_labels: list[str] = [] + terminals: list[bool] = [] + metrics: list[dict[str, float | int] | None] = [] + events: list[list[str] | None] = [] + info_list: list[dict[str, object] | None] = [] + + for i in range(self.num_envs): + reward_i = float(_scalar_at(rewards_np, i, default=0.0)) + term_i = bool(_scalar_at(terminateds_np, i, default=False)) + trunc_i = bool(_scalar_at(truncateds_np, i, default=False)) + slot_info = _per_slot_isaac_info(info, i, self.num_envs, self.info_keys) + self._episode_returns[i] += reward_i + + next_states.append(self._state_from_obs_slot(obs, i)) + outcome, failure_label = self.outcome_from_step(reward_i, term_i, trunc_i, slot_info) + outcomes.append(outcome) + failure_labels.append(failure_label) + terminals.append(term_i or trunc_i) + metrics.append({ + "reward": reward_i, + "episode_return": float(self._episode_returns[i]), + }) + events.append(self.events_from_step(reward_i, term_i, trunc_i, slot_info)) + info_list.append({ + "isaac": { + "terminated": term_i, + "truncated": trunc_i, + "raw_info": to_serializable(slot_info), + "slot": i, + } + }) + + if term_i or trunc_i: + self._episode_returns[i] = 0.0 + + return BatchedStepOutcome( + next_states=next_states, + outcomes=outcomes, + failure_labels=failure_labels, + terminals=terminals, + metrics=metrics, + events=events, + info=info_list, + ) + + def reset_slots( + self, slots: list[int], scenarios: list[Scenario] + ) -> list[State]: + if len(slots) != len(scenarios): + raise ValueError( + f"reset_slots: slots and scenarios must be same length, " + f"got {len(slots)} and {len(scenarios)}." + ) + if not slots: + return [] + + seed = next( + (self.seed_from_scenario(sc) for sc in scenarios + if self.seed_from_scenario(sc) is not None), + None, + ) + + env_ids = self._env_ids_from_slots(slots) + post_obs = self._try_selective_reset(env_ids, seed) + if post_obs is None: + post_obs = self._fallback_full_reset(seed) + + new_states: list[State] = [] + for slot in slots: + self._episode_returns[slot] = 0.0 + new_states.append(self._state_from_obs_slot(post_obs, slot)) + return new_states + + def close(self) -> None: + """Forward close to the wrapped env.""" + close = getattr(self.env, "close", None) + if callable(close): + close() + + # ── Internals ────────────────────────────────────────────────────────── + + def _state_from_obs_slot(self, obs: Any, slot: int) -> State: + sliced = self._slice_obs_at_slot(obs, slot) + return self.observation_to_state(sliced) + + def _slice_obs_at_slot(self, obs: Any, slot: int) -> Any: + """Index batched obs at ``slot`` and coerce tensors to numpy. + + Handles two common Isaac obs shapes: + * dict like ``{"policy": tensor(N, *)}`` → slice each value at slot + * tensor ``(N, *)`` directly → slice at slot + """ + if isinstance(obs, dict): + return {k: tensor_to_numpy(_index_safe(v, slot)) for k, v in obs.items()} + return tensor_to_numpy(_index_safe(obs, slot)) + + def _stack_actions(self, actions: list[Any]) -> Any: + """Stack per-slot actions into a single (N, *action_dim) tensor. + + Accepts ints, floats, lists, numpy arrays, or torch tensors per slot. + Output device matches the env's device when available. + """ + try: + import torch + except ImportError as exc: + raise ImportError( + "PyTorch is required for the Isaac adapter. Install via " + "`pip install torch` or as part of your Isaac Lab install." + ) from exc + + device = self._resolve_device(actions) + as_tensors: list[Any] = [] + for a in actions: + if isinstance(a, torch.Tensor): + t = a + else: + t = torch.as_tensor(a) + if t.device != torch.device(device): + t = t.to(device) + # Ensure each slot's action is 1-D (vector); a scalar becomes (1,) + if t.ndim == 0: + t = t.view(1) + as_tensors.append(t) + stacked = torch.stack(as_tensors, dim=0) + return stacked + + def _resolve_device(self, sample_actions: list[Any]) -> str: + if self._device is not None: + return self._device + + # Same wrapper-proxy issue as num_envs — check env.unwrapped too. + unwrapped = getattr(self.env, "unwrapped", self.env) + env_device = ( + getattr(self.env, "device", None) + or getattr(self.env, "sim_device", None) + or getattr(unwrapped, "device", None) + or getattr(unwrapped, "sim_device", None) + ) + if env_device is not None: + self._device = str(env_device) + return self._device + + try: + import torch + except ImportError: + self._device = "cpu" + return self._device + + for a in sample_actions: + if hasattr(a, "device"): + self._device = str(a.device) + return self._device + + self._device = "cuda" if torch.cuda.is_available() else "cpu" + return self._device + + def _env_ids_from_slots(self, slots: list[int]) -> Any: + """Return slots as a torch tensor on the env's device when possible.""" + try: + import torch + device = self._resolve_device([]) + return torch.tensor(slots, dtype=torch.long, device=device) + except ImportError: + return slots + + def _try_selective_reset(self, env_ids: Any, seed: int | None) -> Any: + """Attempt selective reset via two known Isaac code paths. + + Returns the post-reset obs (batched, with all envs included) if a + selective reset succeeded, or None if neither path was available. + """ + # Path 1: gym wrapper with env_ids kwarg + try: + obs, _info = self.env.reset(seed=seed, env_ids=env_ids) + return obs + except TypeError: + pass + + # Path 2: ManagerBasedRLEnv._reset_idx + unwrapped = getattr(self.env, "unwrapped", self.env) + reset_idx = getattr(unwrapped, "_reset_idx", None) + if callable(reset_idx): + try: + reset_idx(env_ids) + # Read post-reset obs without stepping + obs = self._observe_unwrapped(unwrapped) + if obs is not None: + return obs + except Exception: # noqa: BLE001 — any error → fall through + pass + + return None + + def _observe_unwrapped(self, unwrapped: Any) -> Any: + """Pull current obs without stepping. Tries common Isaac entry points.""" + get_obs = getattr(unwrapped, "_get_observations", None) + if callable(get_obs): + try: + return get_obs() + except Exception: # noqa: BLE001 + return None + obs_buf = getattr(unwrapped, "obs_buf", None) + if obs_buf is not None: + return obs_buf + return None + + def _fallback_full_reset(self, seed: int | None) -> Any: + if not self._warned_full_reset: + warnings.warn( + "BatchedIsaacEnvironmentAdapter could not perform a selective " + "reset (no env_ids kwarg and no _reset_idx). Falling back to " + "full env.reset() — this resets ALL slots including ones still " + "rolling. For deterministic per-slot refills, wire the env to " + "support env_ids or expose _reset_idx.", + stacklevel=3, + ) + self._warned_full_reset = True + try: + obs, _info = self.env.reset(seed=seed) + except TypeError: + obs, _info = self.env.reset() + return obs + + +# ── module helpers ────────────────────────────────────────────────────────── + + +def _index_safe(value: Any, slot: int) -> Any: + try: + return value[slot] + except (TypeError, IndexError, KeyError): + return value + + +def _scalar_at(value: Any, slot: int, default: Any) -> Any: + """Read a scalar from a 1-D numpy/tensor at ``slot``, falling back gracefully.""" + sliced = _index_safe(value, slot) + if hasattr(sliced, "item"): + try: + return sliced.item() + except (ValueError, RuntimeError): + pass + try: + return float(sliced) if not isinstance(sliced, bool) else bool(sliced) + except (TypeError, ValueError): + return default + + +def _per_slot_isaac_info( + info: Any, + slot: int, + num_envs: int, + info_keys: list[str] | None, +) -> dict[str, Any]: + """Pluck slot-specific entries from a batched Isaac info dict. + + Isaac's info is a dict where per-env entries are tensors of shape + ``(num_envs, *)``; scalar entries apply to the whole batch. + """ + if not isinstance(info, dict): + return {} + slot_info: dict[str, Any] = {} + keys = info_keys if info_keys is not None else list(info.keys()) + for key in keys: + if key not in info: + continue + value = info[key] + if hasattr(value, "__len__") and not isinstance(value, (str, bytes)): + try: + if len(value) == num_envs: + slot_info[key] = tensor_to_numpy(value[slot]) + continue + except TypeError: + pass + slot_info[key] = tensor_to_numpy(value) if hasattr(value, "is_cuda") else value + return slot_info diff --git a/roboeval/integrations/isaac/tests/test_batched_adapter.py b/roboeval/integrations/isaac/tests/test_batched_adapter.py new file mode 100644 index 0000000..e40f8db --- /dev/null +++ b/roboeval/integrations/isaac/tests/test_batched_adapter.py @@ -0,0 +1,601 @@ +"""Unit + integration tests for BatchedIsaacEnvironmentAdapter. + +Runs on Mac without Isaac installed. Uses a MockBatchedIsaacEnv that fakes +the Isaac Lab vectorized shape — real torch tensors, batched obs/reward/ +terminal flags, and configurable per-slot termination steps. +""" + +from __future__ import annotations + +import unittest +import warnings +from typing import Any + +import numpy as np + +try: + import torch +except ImportError: # pragma: no cover + torch = None # type: ignore[assignment] + +from roboeval.batched.runner import BatchedEvalRunner +from roboeval.batched.types import BatchedStepOutcome +from roboeval.core import Ruleset, Scenario, require_metric, require_outcome +from roboeval.integrations.isaac.batched_adapter import ( + BatchedIsaacEnvironmentAdapter, + _per_slot_isaac_info, + _scalar_at, +) + + +HAS_TORCH = torch is not None + + +# ─── Mock Isaac vectorized env ────────────────────────────────────────────── + + +class MockBatchedIsaacEnv: + """Fakes a vectorized Isaac Lab env with real torch tensors. + + Each env in the batch has its own step counter and terminates when it + hits its assigned ``terminate_at_step``. Supports ``_reset_idx`` for + selective reset (the recommended Isaac path) and ``_get_observations`` + so the adapter can read post-reset obs without stepping. + """ + + def __init__( + self, + num_envs: int = 4, + obs_dim: int = 4, + action_dim: int = 2, + terminate_at_steps: list[int] | int | None = None, + obs_as_dict: bool = True, + device: str = "cpu", + rich_info: bool = False, + ) -> None: + if not HAS_TORCH: + raise RuntimeError("MockBatchedIsaacEnv requires torch") + self.num_envs = num_envs + self._obs_dim = obs_dim + self._action_dim = action_dim + if terminate_at_steps is None: + terminate_at_steps = [10] * num_envs + elif isinstance(terminate_at_steps, int): + terminate_at_steps = [terminate_at_steps] * num_envs + self._terminate_at = list(terminate_at_steps) + self._obs_as_dict = obs_as_dict + self.device = device + self.rich_info = rich_info + self._step_counts = [0] * num_envs + self._last_action: Any = None + self.reset_calls = 0 + self.reset_idx_calls = 0 + self.step_calls = 0 + self._last_obs: Any = None + + def _make_obs(self) -> Any: + # Per-slot obs is the step count, broadcast across obs_dim + flat = torch.tensor( + [[float(c)] * self._obs_dim for c in self._step_counts], + dtype=torch.float32, device=self.device, + ) + if self._obs_as_dict: + return {"policy": flat} + return flat + + def reset(self, seed: int | None = None, options: dict | None = None): + self.reset_calls += 1 + self._step_counts = [0] * self.num_envs + self._last_obs = self._make_obs() + return self._last_obs, {} + + def step(self, action): + if hasattr(action, "shape") and action.shape[0] != self.num_envs: + raise ValueError( + f"action batch dim {action.shape[0]} != num_envs {self.num_envs}" + ) + self.step_calls += 1 + self._last_action = action + for i in range(self.num_envs): + self._step_counts[i] += 1 + obs = self._make_obs() + self._last_obs = obs + reward = torch.ones(self.num_envs, dtype=torch.float32, device=self.device) + terminated = torch.tensor( + [self._step_counts[i] >= self._terminate_at[i] for i in range(self.num_envs)], + dtype=torch.bool, device=self.device, + ) + truncated = torch.zeros(self.num_envs, dtype=torch.bool, device=self.device) + info: dict[str, Any] = {} + if self.rich_info: + info["task_reward"] = torch.tensor( + [float(c) * 2 for c in self._step_counts], + dtype=torch.float32, device=self.device, + ) + info["step_count"] = torch.tensor( + self._step_counts, dtype=torch.long, device=self.device + ) + info["sim_version"] = "mock_v1" # scalar info + return obs, reward, terminated, truncated, info + + # ── Isaac Lab Manager API surface (for selective reset) ───────────── + @property + def unwrapped(self): + return self + + def _reset_idx(self, env_ids) -> None: + self.reset_idx_calls += 1 + if hasattr(env_ids, "tolist"): + env_ids = env_ids.tolist() + for i in env_ids: + self._step_counts[i] = 0 + # Refresh obs buffer so _get_observations returns post-reset state + self._last_obs = self._make_obs() + + def _get_observations(self): + return self._last_obs + + def close(self): + pass + + +class MockBatchedIsaacEnvNoSelectiveReset: + """Mock that doesn't expose _reset_idx — exercises the full-reset fallback.""" + + def __init__(self, num_envs: int = 2): + if not HAS_TORCH: + raise RuntimeError("torch required") + self.num_envs = num_envs + self.device = "cpu" + self._step_counts = [0] * num_envs + self.reset_calls = 0 + + @property + def unwrapped(self): + # No _reset_idx, no _get_observations + return object() + + def reset(self, seed: int | None = None, options: dict | None = None): + self.reset_calls += 1 + self._step_counts = [0] * self.num_envs + return {"policy": torch.zeros(self.num_envs, 4)}, {} + + def step(self, action): + for i in range(self.num_envs): + self._step_counts[i] += 1 + obs = {"policy": torch.tensor([[float(c)] * 4 for c in self._step_counts])} + reward = torch.ones(self.num_envs) + terminated = torch.tensor([c >= 3 for c in self._step_counts]) + truncated = torch.zeros(self.num_envs, dtype=torch.bool) + return obs, reward, terminated, truncated, {} + + +# ─── Construction + validation ────────────────────────────────────────────── + + +@unittest.skipUnless(HAS_TORCH, "torch not installed") +class TestBatchedIsaacConstruction(unittest.TestCase): + def test_rejects_env_without_num_envs(self) -> None: + class NoNumEnvs: + def reset(self): return None, {} + def step(self, a): return None, None, None, None, {} + + with self.assertRaisesRegex(TypeError, r"num_envs"): + BatchedIsaacEnvironmentAdapter(env=NoNumEnvs()) + + def test_rejects_zero_num_envs(self) -> None: + class ZeroNumEnvs: + num_envs = 0 + def reset(self, **k): return None, {} + def step(self, a): return None, None, None, None, {} + + with self.assertRaisesRegex(ValueError, r"num_envs must be positive"): + BatchedIsaacEnvironmentAdapter(env=ZeroNumEnvs()) + + def test_num_envs_set_from_env(self) -> None: + env = MockBatchedIsaacEnv(num_envs=8) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + self.assertEqual(adapter.num_envs, 8) + + def test_num_envs_falls_back_to_unwrapped(self) -> None: + """gym.make() wraps Isaac envs and the outer wrapper doesn't proxy + num_envs. Adapter must read it from env.unwrapped — regression check + for the bug discovered while validating against real Isaac-Cartpole.""" + inner = MockBatchedIsaacEnv(num_envs=8) + + class GymStyleWrapper: + """Simulates the gym.OrderEnforcing/TimeLimit wrapper layer that + sits between gym.make()'s return value and the underlying env.""" + def __init__(self, inner_env): + self.unwrapped = inner_env + + def reset(self, **kwargs): + return self.unwrapped.reset(**kwargs) + + def step(self, action): + return self.unwrapped.step(action) + + wrapped = GymStyleWrapper(inner) + # The wrapped env has NO .num_envs attribute, only .unwrapped.num_envs + self.assertFalse(hasattr(wrapped, "num_envs")) + adapter = BatchedIsaacEnvironmentAdapter(env=wrapped) + self.assertEqual(adapter.num_envs, 8) + + +# ─── reset() ───────────────────────────────────────────────────────────────── + + +@unittest.skipUnless(HAS_TORCH, "torch not installed") +class TestBatchedIsaacReset(unittest.TestCase): + def test_reset_returns_num_envs_states(self) -> None: + env = MockBatchedIsaacEnv(num_envs=4) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + scenarios = [Scenario(f"s{i}", {"seed": i}, max_steps=20) for i in range(4)] + states = adapter.reset(scenarios) + self.assertEqual(len(states), 4) + for s in states: + self.assertIn("policy", s) + # Each state's policy obs is numpy (coerced from torch) + self.assertIsInstance(s["policy"], np.ndarray) + self.assertEqual(s["policy"].shape, (4,)) + + def test_reset_requires_num_envs_scenarios(self) -> None: + env = MockBatchedIsaacEnv(num_envs=3) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + with self.assertRaisesRegex(ValueError, r"3 scenarios"): + adapter.reset([Scenario("only", {}, max_steps=10)]) + + def test_reset_zeroes_episode_returns(self) -> None: + env = MockBatchedIsaacEnv(num_envs=2, terminate_at_steps=100) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario("a", {}, max_steps=50), Scenario("b", {}, max_steps=50)]) + # Step a few times to accumulate + for _ in range(3): + adapter.step([0, 0]) + self.assertGreater(sum(adapter._episode_returns), 0.0) + adapter.reset([Scenario("a", {}, max_steps=50), Scenario("b", {}, max_steps=50)]) + self.assertEqual(adapter._episode_returns, [0.0, 0.0]) + + def test_reset_tolerates_env_without_options_kwarg(self) -> None: + class NoOptionsEnv(MockBatchedIsaacEnv): + def reset(self, seed=None): + return super().reset(seed=seed) + + env = NoOptionsEnv(num_envs=2) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + states = adapter.reset([Scenario("a", {}, max_steps=10), Scenario("b", {}, max_steps=10)]) + self.assertEqual(len(states), 2) + + +# ─── step() ───────────────────────────────────────────────────────────────── + + +@unittest.skipUnless(HAS_TORCH, "torch not installed") +class TestBatchedIsaacStep(unittest.TestCase): + def test_step_returns_batched_step_outcome(self) -> None: + env = MockBatchedIsaacEnv(num_envs=3, terminate_at_steps=[2, 4, 6]) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario(f"s{i}", {}, max_steps=10) for i in range(3)]) + outcome = adapter.step([0, 1, 0]) + self.assertIsInstance(outcome, BatchedStepOutcome) + self.assertEqual(outcome.num_envs, 3) + + def test_step_requires_num_envs_actions(self) -> None: + env = MockBatchedIsaacEnv(num_envs=2) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario("a", {}, max_steps=5), Scenario("b", {}, max_steps=5)]) + with self.assertRaisesRegex(ValueError, r"2 actions"): + adapter.step([0]) + + def test_per_slot_terminals_independent(self) -> None: + """Slot 0 terminates at step 2, slot 1 at step 4. Both must surface correctly.""" + env = MockBatchedIsaacEnv(num_envs=2, terminate_at_steps=[2, 4]) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario("a", {}, max_steps=10), Scenario("b", {}, max_steps=10)]) + # Step 1: nobody terminates + out = adapter.step([0, 0]) + self.assertEqual(out.terminals, [False, False]) + # Step 2: slot 0 terminates + out = adapter.step([0, 0]) + self.assertEqual(out.terminals, [True, False]) + # Step 3: slot 0 just got auto-stepped (mock keeps counting); only slot 1 matters + out = adapter.step([0, 0]) + self.assertEqual(out.terminals[1], False) + # Step 4: slot 1 terminates + out = adapter.step([0, 0]) + self.assertTrue(out.terminals[1]) + + def test_per_slot_episode_return_accumulates(self) -> None: + env = MockBatchedIsaacEnv(num_envs=2, terminate_at_steps=100) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario("a", {}, max_steps=10), Scenario("b", {}, max_steps=10)]) + for _ in range(3): + out = adapter.step([0, 0]) + # Each slot got reward=1 per step for 3 steps + self.assertEqual(out.metrics[0]["episode_return"], 3.0) + self.assertEqual(out.metrics[1]["episode_return"], 3.0) + + def test_terminal_zeroes_episode_return_for_that_slot(self) -> None: + env = MockBatchedIsaacEnv(num_envs=2, terminate_at_steps=[2, 10]) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario("a", {}, max_steps=20), Scenario("b", {}, max_steps=20)]) + adapter.step([0, 0]) # step 1 + adapter.step([0, 0]) # step 2 — slot 0 terminates; episode_return zeroed afterward + self.assertEqual(adapter._episode_returns[0], 0.0) + # Slot 1 keeps accumulating + self.assertGreater(adapter._episode_returns[1], 0.0) + + def test_info_namespace_isaac(self) -> None: + env = MockBatchedIsaacEnv(num_envs=2, terminate_at_steps=10, rich_info=True) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario("a", {}, max_steps=10), Scenario("b", {}, max_steps=10)]) + out = adapter.step([0, 0]) + for slot in (0, 1): + self.assertIn("isaac", out.info[slot]) + isaac_info = out.info[slot]["isaac"] + self.assertIn("terminated", isaac_info) + self.assertIn("truncated", isaac_info) + self.assertIn("raw_info", isaac_info) + self.assertEqual(isaac_info["slot"], slot) + # rich_info batched values land per-slot + self.assertIn("task_reward", isaac_info["raw_info"]) + self.assertIn("step_count", isaac_info["raw_info"]) + # scalar info is forwarded + self.assertEqual(isaac_info["raw_info"]["sim_version"], "mock_v1") + + def test_metrics_carry_reward_and_episode_return(self) -> None: + env = MockBatchedIsaacEnv(num_envs=2) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario("a", {}, max_steps=10), Scenario("b", {}, max_steps=10)]) + out = adapter.step([0, 0]) + for slot in (0, 1): + self.assertIn("reward", out.metrics[slot]) + self.assertIn("episode_return", out.metrics[slot]) + + def test_obs_with_bare_tensor_handled(self) -> None: + env = MockBatchedIsaacEnv(num_envs=2, obs_as_dict=False) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + states = adapter.reset([Scenario("a", {}, max_steps=10), Scenario("b", {}, max_steps=10)]) + # Bare tensor → state has "observation" key + for s in states: + self.assertIn("observation", s) + + +# ─── Action stacking ──────────────────────────────────────────────────────── + + +@unittest.skipUnless(HAS_TORCH, "torch not installed") +class TestActionStacking(unittest.TestCase): + def test_int_actions_become_batched_tensor(self) -> None: + env = MockBatchedIsaacEnv(num_envs=3, action_dim=1) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario(f"s{i}", {}, max_steps=10) for i in range(3)]) + adapter.step([0, 1, 0]) + # Mock stores last action — verify shape + last = env._last_action + self.assertEqual(tuple(last.shape), (3, 1)) + + def test_list_actions_become_batched_tensor(self) -> None: + env = MockBatchedIsaacEnv(num_envs=2, action_dim=3) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario(f"s{i}", {}, max_steps=10) for i in range(2)]) + adapter.step([[0.1, 0.2, 0.3], [-0.1, 0.0, 0.5]]) + last = env._last_action + self.assertEqual(tuple(last.shape), (2, 3)) + + def test_tensor_actions_stack(self) -> None: + env = MockBatchedIsaacEnv(num_envs=2, action_dim=2) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario(f"s{i}", {}, max_steps=10) for i in range(2)]) + adapter.step([torch.tensor([0.1, 0.2]), torch.tensor([0.3, 0.4])]) + last = env._last_action + self.assertEqual(tuple(last.shape), (2, 2)) + + +# ─── reset_slots() ────────────────────────────────────────────────────────── + + +@unittest.skipUnless(HAS_TORCH, "torch not installed") +class TestBatchedIsaacResetSlots(unittest.TestCase): + def test_selective_reset_via_reset_idx(self) -> None: + """When env exposes ``_reset_idx``, we hit that path (not full reset).""" + env = MockBatchedIsaacEnv(num_envs=4, terminate_at_steps=[2, 100, 100, 100]) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario(f"s{i}", {}, max_steps=20) for i in range(4)]) + # Step until slot 0 has accumulated some step count + for _ in range(3): + adapter.step([0, 0, 0, 0]) + before_full_resets = env.reset_calls + # Refill slot 0 with a new scenario + new_states = adapter.reset_slots([0], [Scenario("refilled", {"seed": 99}, max_steps=20)]) + self.assertEqual(len(new_states), 1) + # The selective path was used + self.assertGreaterEqual(env.reset_idx_calls, 1) + # No additional full reset was needed + self.assertEqual(env.reset_calls, before_full_resets) + # Other slots' step counts are NOT zeroed (proves selectivity) + self.assertEqual(env._step_counts[0], 0) + self.assertGreater(env._step_counts[1], 0) + self.assertGreater(env._step_counts[2], 0) + self.assertGreater(env._step_counts[3], 0) + + def test_fallback_full_reset_when_no_selective_path(self) -> None: + """When env has no ``_reset_idx``, the adapter warns and full-resets.""" + env = MockBatchedIsaacEnvNoSelectiveReset(num_envs=2) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario("a", {}, max_steps=10), Scenario("b", {}, max_steps=10)]) + before = env.reset_calls + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + adapter.reset_slots([0], [Scenario("refill", {"seed": 1}, max_steps=10)]) + # Full reset was called + self.assertGreater(env.reset_calls, before) + # And a warning was emitted + messages = [str(w.message) for w in caught] + self.assertTrue(any("selective reset" in m for m in messages)) + + def test_reset_slots_zeroes_per_slot_episode_return(self) -> None: + env = MockBatchedIsaacEnv(num_envs=2, terminate_at_steps=100) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario("a", {}, max_steps=20), Scenario("b", {}, max_steps=20)]) + for _ in range(3): + adapter.step([0, 0]) + before = list(adapter._episode_returns) + self.assertTrue(all(r > 0 for r in before)) + adapter.reset_slots([0], [Scenario("refilled", {}, max_steps=20)]) + self.assertEqual(adapter._episode_returns[0], 0.0) + self.assertEqual(adapter._episode_returns[1], before[1]) + + def test_reset_slots_length_mismatch_raises(self) -> None: + env = MockBatchedIsaacEnv(num_envs=2) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario("a", {}, max_steps=10), Scenario("b", {}, max_steps=10)]) + with self.assertRaisesRegex(ValueError, r"same length"): + adapter.reset_slots([0, 1], [Scenario("only", {}, max_steps=10)]) + + def test_reset_slots_empty_input_is_noop(self) -> None: + env = MockBatchedIsaacEnv(num_envs=2) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + adapter.reset([Scenario("a", {}, max_steps=10), Scenario("b", {}, max_steps=10)]) + out = adapter.reset_slots([], []) + self.assertEqual(out, []) + + +# ─── Helper functions ─────────────────────────────────────────────────────── + + +@unittest.skipUnless(HAS_TORCH, "torch not installed") +class TestScalarAt(unittest.TestCase): + def test_extracts_from_tensor(self) -> None: + t = torch.tensor([1.0, 2.0, 3.0]) + self.assertEqual(_scalar_at(t, 1, default=0.0), 2.0) + + def test_extracts_from_numpy(self) -> None: + arr = np.array([True, False, True]) + self.assertEqual(_scalar_at(arr, 0, default=False), True) + + def test_fallback_on_bad_index(self) -> None: + # Scalar value gets passed through + self.assertEqual(_scalar_at(1.5, 0, default=0.0), 1.5) + + +@unittest.skipUnless(HAS_TORCH, "torch not installed") +class TestPerSlotIsaacInfo(unittest.TestCase): + def test_extracts_per_slot_tensor_values(self) -> None: + info = {"task_reward": torch.tensor([10.0, 20.0, 30.0])} + result = _per_slot_isaac_info(info, slot=1, num_envs=3, info_keys=None) + self.assertEqual(result["task_reward"], 20.0) + + def test_passes_through_scalar_info(self) -> None: + info = {"sim_version": "v1"} + result = _per_slot_isaac_info(info, slot=0, num_envs=4, info_keys=None) + self.assertEqual(result, {"sim_version": "v1"}) + + def test_info_keys_allowlist_filters(self) -> None: + info = { + "task_reward": torch.tensor([1.0, 2.0]), + "big_tensor": torch.tensor([[1, 2], [3, 4]]), + } + result = _per_slot_isaac_info(info, slot=0, num_envs=2, info_keys=["task_reward"]) + self.assertEqual(set(result.keys()), {"task_reward"}) + + def test_non_dict_info_returns_empty(self) -> None: + result = _per_slot_isaac_info(None, slot=0, num_envs=2, info_keys=None) + self.assertEqual(result, {}) + + +# ─── End-to-end: BatchedEvalRunner + MockBatchedIsaacEnv ──────────────────── + + +def _balance_policy(state): + """Use slot's policy obs first value to pick an action.""" + if "policy" in state: + v = float(state["policy"][0]) + else: + v = float(state["observation"][0]) + return {"action": [0.5 if v < 5 else -0.5], "debug_info": {"v": v}} + + +def _zero_policy(state): + return {"action": [0.0], "debug_info": {"version": "zero"}} + + +@unittest.skipUnless(HAS_TORCH, "torch not installed") +class TestEndToEndBatchedIsaacWithRunner(unittest.TestCase): + def test_three_policies_four_scenarios_eight_envs(self) -> None: + env = MockBatchedIsaacEnv(num_envs=8, action_dim=1, terminate_at_steps=15) + adapter = BatchedIsaacEnvironmentAdapter(env=env, name="mock_isaac_x8") + scenarios = [ + Scenario(f"seed_{s}", {"seed": s}, max_steps=30, + metadata={"tags": ["mock_isaac"]}) + for s in (1, 2, 42, 100) + ] + from roboeval.batched.policy import from_single + report = BatchedEvalRunner( + policies=[ + from_single(_balance_policy), + from_single(_zero_policy), + from_single(_balance_policy, name="balance_v2"), + ], + scenarios=scenarios, + environment=adapter, + ruleset=Ruleset([require_metric("episode_return", ">=", 10.0)]), + baseline_policy="_balance_policy", + ).run() + # 3 policies × 4 scenarios = 12 episodes + self.assertEqual(len(report.episodes), 12) + for ep in report.episodes: + self.assertGreater(ep.steps, 0) + # Each step record carries the isaac namespace + for log in ep.logs: + self.assertIn("isaac", log.info) + + def test_replication_with_isaac_mock(self) -> None: + env = MockBatchedIsaacEnv(num_envs=4, action_dim=1, terminate_at_steps=10) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + from roboeval.batched.policy import from_single + report = BatchedEvalRunner( + policies=[from_single(_zero_policy)], + scenarios=[Scenario("only", {"seed": 0}, max_steps=20)], + environment=adapter, + replicas=4, + ruleset=Ruleset([require_outcome("terminated_success")]), + ).run() + self.assertEqual(len(report.episodes), 4) + names = sorted(ep.scenario_name for ep in report.episodes) + self.assertEqual(names, ["only#r0", "only#r1", "only#r2", "only#r3"]) + + def test_ruleset_flow_through(self) -> None: + env = MockBatchedIsaacEnv(num_envs=2, action_dim=1, terminate_at_steps=5) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + from roboeval.batched.policy import from_single + # episode_return = 5 (5 steps * reward=1); ruleset requires >= 100 + report = BatchedEvalRunner( + policies=[from_single(_zero_policy)], + scenarios=[Scenario("a", {}, max_steps=20), Scenario("b", {}, max_steps=20)], + environment=adapter, + ruleset=Ruleset([require_metric("episode_return", ">=", 100.0, name="hard_target")]), + ).run() + self.assertEqual(len(report.episodes), 2) + for ep in report.episodes: + self.assertFalse(ep.success) + self.assertEqual(ep.failure_label, "hard_target") + + def test_reset_slots_called_during_refill(self) -> None: + """num_envs=2 with 4 scenarios → refill kicks in twice, exercising reset_slots.""" + env = MockBatchedIsaacEnv(num_envs=2, action_dim=1, terminate_at_steps=3) + adapter = BatchedIsaacEnvironmentAdapter(env=env) + from roboeval.batched.policy import from_single + report = BatchedEvalRunner( + policies=[from_single(_zero_policy)], + scenarios=[Scenario(f"s{i}", {}, max_steps=10) for i in range(4)], + environment=adapter, + ruleset=Ruleset([require_metric("episode_return", ">=", 1.0)]), + ).run() + self.assertEqual(len(report.episodes), 4) + # Each refill triggers a _reset_idx call on the underlying mock + self.assertGreaterEqual(env.reset_idx_calls, 1) + + +if __name__ == "__main__": + unittest.main() From 3467b9f8dfd47ac87ea666893a153a2db8b306ae Mon Sep 17 00:00:00 2001 From: arryan0112 Date: Mon, 8 Jun 2026 16:18:12 +0530 Subject: [PATCH 6/6] Fix ruff E702 in test_scheduler.py: split semicolon-joined statements --- roboeval/batched/tests/test_scheduler.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/roboeval/batched/tests/test_scheduler.py b/roboeval/batched/tests/test_scheduler.py index 0578bc5..73ee376 100644 --- a/roboeval/batched/tests/test_scheduler.py +++ b/roboeval/batched/tests/test_scheduler.py @@ -121,13 +121,17 @@ def test_drains_queue_correctly_with_mixed_completion_order(self) -> None: seen = ["a", "b", "c"] # initial assignments # slot 2 finishes first - t = sched.complete_slot(2); seen.append(t.scenario.name) # d + t = sched.complete_slot(2) + seen.append(t.scenario.name) # d # slot 0 finishes - t = sched.complete_slot(0); seen.append(t.scenario.name) # e + t = sched.complete_slot(0) + seen.append(t.scenario.name) # e # slot 2 again - t = sched.complete_slot(2); seen.append(t.scenario.name) # f + t = sched.complete_slot(2) + seen.append(t.scenario.name) # f # slot 1 finishes - t = sched.complete_slot(1); seen.append(t.scenario.name) # g + t = sched.complete_slot(1) + seen.append(t.scenario.name) # g # Now queue empty. Remaining slot completions return None. self.assertIsNone(sched.complete_slot(2))