Spike/vectorized env - #6
Open
arryan0112 wants to merge 6 commits into
Open
Conversation
Wraps any single-env Isaac Lab environment (gym.make("Isaac-...-v0",
num_envs=1)) into the EnvironmentAdapter Protocol so policies can be
evaluated against Isaac Lab tasks through EvalRunner. Zero changes to
existing SDK files.
The adapter handles three Isaac-specific concerns that distinguish it
from the existing Gymnasium adapter:
1. Always-batched envs. Isaac envs subclass gym.vector.VectorEnv and
return tensors with a leading batch dimension even at num_envs=1.
The adapter slices batch_index=0 (configurable) to expose
single-episode semantics to the runner.
2. GPU tensors. Observations come back as torch.Tensor on cuda:0.
The tensor_to_numpy helper moves any tensor to CPU and converts to
numpy before it reaches the runner's JSON writer.
3. Tensor actions on the env device. Policies typically return numpy
arrays or Python scalars; the adapter wraps them in a batched torch
tensor on the env's device with the expected (num_envs, action_dim)
shape.
Vectorized eval (num_envs > 1 producing N independent rollouts) is
explicitly out of scope for this spike. The adapter accepts num_envs > 1
but only reads batch_index=0; a clear warning fires explaining the
parallelism cost. True vector eval requires SDK runner changes documented
as Phase 3+ work.
Files added:
- roboeval/integrations/isaac/__init__.py (public re-exports)
- roboeval/integrations/isaac/adapter.py (adapter + 6 default hooks
+ tensor_to_numpy helper + action shape normalization)
- roboeval/integrations/isaac/demo_rollout.py (Isaac-Cartpole-Direct-v0
manual rollout demo)
- roboeval/integrations/isaac/README.md (usage, mapping table, cloud
GPU setup workflow)
- roboeval/integrations/isaac/notes.md (design rationale, 12 known
edge cases, real-Isaac validation plan)
- roboeval/integrations/isaac/requirements.txt (Isaac Sim + Isaac Lab
install pointers; not a single pip line)
- roboeval/integrations/isaac/tests/__init__.py
- roboeval/integrations/isaac/tests/test_adapter.py (46 unit +
integration tests using mock Isaac envs; runs on Mac, no Isaac
install needed)
What's verified by the test suite:
- All 46 tests pass on Mac in 23ms via mocked Isaac envs (real torch
tensors, no Isaac dependency)
- Adapter integrates cleanly with EvalRunner + Ruleset
- info["isaac"] namespace preserved through to StepRecord
- Rich env info flows through to_serializable + JSON write
- Hook overrides work (outcome, action, info_keys, etc.)
- Multi-episode reset behavior
- Action shape normalization across 4 input shapes (scalar, 1-D array
matching action_dim, 1-D with batch shape, pre-batched 2-D tensor)
What's NOT verified (requires cloud GPU pass):
- Real Isaac-Cartpole-Direct-v0 smoke test
- GPU->CPU tensor sync against actual cuda:0 tensors
- env.device / env.sim_device attribute resolution on real Isaac
- env.reset(options=...) acceptance across Isaac task variants
See notes.md section 12 for the cloud GPU validation plan.
Validate:
python -m unittest discover -s roboeval/integrations/isaac/tests
- BatchedState (list[State]) and BatchedStepOutcome (per-slot lists with length validation) - BatchedEnvironmentAdapter Protocol (reset/step/reset_slots) for N-env vectorized rollouts - BatchedPolicy Protocol, BatchedPolicyAdapter, and from_single() shim for single-state policies - Tests for type construction, validation, slot extraction, and policy normalization
- SlotScheduler: maps (Scenario, replica) tasks to N env slots with refill on slot completion; replication support (D1) for running each scenario K times - BatchedEvalRunner: drives one policy at a time through the scheduler, fans per-slot outcomes into normal EpisodeResults, reuses _build_report from EvalRunner so reports are byte-identical to the single-env path - Tests: scheduler init/refill/replication/idle handling; runner smoke, refill, idle slots, multi-policy, replication scenario_name suffix, max_steps enforcement, ruleset flow, report shape
- Wraps gym.vector.SyncVectorEnv as a BatchedEnvironmentAdapter; supports any class exposing env.envs (AsyncVectorEnv not yet supported) - Reuses the six default hooks from the single-env adapter for translation parity at the slot level - Handles Gymnasium 1.x NEXT_STEP autoreset; reset_slots overrides the implicit auto-reset by calling env.envs[i].reset(seed=...) for deterministic per-slot refills - Tests: construction, reset/step shape, per-slot episode_return accumulation, reset_slots determinism, action conversion, end-to-end against real CartPole-v1 with BatchedEvalRunner
… lookup - Reads the full Isaac GPU batch (no batch_index=0 slicing); stacks per-slot actions into (N, action_dim) tensors and fans (N,) reward/terminal tensors back into per-slot lists - Selective reset via three-tier fallback: env.reset(env_ids=...), env.unwrapped._reset_idx(env_ids), or full env.reset() with a one-time warning - num_envs and device lookups fall back to env.unwrapped to handle gym.make() wrappers (OrderEnforcing, TimeLimit) that don't proxy these attributes - Tests: construction, reset/step shape, per-slot terminals, episode_return tracking, action stacking, selective reset paths, end-to-end with BatchedEvalRunner against a mock that uses real torch tensors; regression test for the wrapper-aware num_envs lookup
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds a parallel roboeval/batched/ namespace that lets EvalRunner-style evals run against vectorized environments — Isaac Lab, gym.vector.SyncVectorEnv, and any other sim that exposes an N-env batch interface. The single-env EvalRunner and its API are not changed; batched code is a strict superset added alongside.
Validated end-to-end on a real NVIDIA Isaac Sim 4.5 deployment (cloud A40 GPU): the same eval pipeline that runs on gymnasium.CartPole-v1 now runs on Isaac-Cartpole-Direct-v0 at num_envs=1024, yielding 108,136 env-steps/sec — a 186× throughput improvement over single-env baseline.