diff --git a/roboeval/batched/__init__.py b/roboeval/batched/__init__.py new file mode 100644 index 0000000..656c603 --- /dev/null +++ b/roboeval/batched/__init__.py @@ -0,0 +1,32 @@ +"""Vectorized-environment namespace for roboeval. + +Single-env users should keep importing from ``roboeval`` directly. This +subpackage adds parallel types/Protocols for N-env vectorized rollouts +(Isaac Lab, gym.vector, Brax, MJX), strictly isolated from the single-env +public API so existing callers see no behavior change. +""" + +from .environment import BatchedEnvironmentAdapter +from .policy import ( + BatchedPolicy, + BatchedPolicyAdapter, + from_single, + normalize_batched_policy, +) +from .runner import BatchedEvalRunner +from .scheduler import SlotScheduler, SlotTask +from .types import BatchedState, BatchedStepOutcome + + +__all__ = [ + "BatchedEnvironmentAdapter", + "BatchedEvalRunner", + "BatchedPolicy", + "BatchedPolicyAdapter", + "BatchedState", + "BatchedStepOutcome", + "SlotScheduler", + "SlotTask", + "from_single", + "normalize_batched_policy", +] 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/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/__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_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..73ee376 --- /dev/null +++ b/roboeval/batched/tests/test_scheduler.py @@ -0,0 +1,147 @@ +"""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() 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)." + ) 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() 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..37359b6 --- /dev/null +++ b/roboeval/integrations/isaac/__init__.py @@ -0,0 +1,51 @@ +"""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, +) +from .batched_adapter import BatchedIsaacEnvironmentAdapter + +__all__ = [ + "BatchedIsaacEnvironmentAdapter", + "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/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/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() 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()