Add JaxARC environment setup and documentation (#188)

Author: aadimator

README.md

@@ -95,7 +95,7 @@ Stoix currently offers the following building blocks for Single-Agent RL researc
 ### Environment Wrappers 🍬
 Stoix offers wrappers for:
 
-- **JAX environments:** [Gymnax][gymnax], [Jumanji][jumanji], [Brax][brax], [XMinigrid][xminigrid], [Craftax][craftax], [POPJym][popjym], [Navix][navix] and more.
+- **JAX environments:** [Gymnax][gymnax], [Jumanji][jumanji], [Brax][brax], [XMinigrid][xminigrid], [Craftax][craftax], [POPJym][popjym], [Navix][navix], [JaxARC][jaxarc] and more.
 - **Non-JAX environments:** [Envpool][envpool] and [Gymnasium][gymnasium].
 
 ### Statistically Robust Evaluation 🧪
@@ -270,3 +270,4 @@ We would like to thank the authors and developers of [Mava][mava] as this was es
 [envpool]: https://github.com/sail-sg/envpool/
 [gymnasium]: https://github.com/Farama-Foundation/Gymnasium
 [cleanba]: https://github.com/vwxyzjn/cleanba
+[jaxarc]: https://github.com/aadimator/JaxARC

docs/envs/jaxarc.md

@@ -0,0 +1,82 @@
+# JaxARC
+
+[JaxARC](https://github.com/aadimator/JaxARC) is a hardware-accelerated ARC (Abstraction and Reasoning Corpus) environment built on JAX. It enables massively parallel RL training on ARC puzzle tasks with full JIT compilation support.
+
+## Basic Usage
+
+To run PPO with JaxARC on the MiniARC dataset:
+
+```bash
+python stoix/systems/ppo/anakin/ff_ppo.py env=jaxarc/default
+```
+
+This uses the `Mini` dataset (70 tasks) with point-based actions — a good starting point for prototyping.
+
+To train on the full ARC-AGI-1 training set (400 tasks):
+
+```bash
+python stoix/systems/ppo/anakin/ff_ppo.py env=jaxarc/agi1_train
+```
+
+## Available Configs
+
+| Config | Dataset | Tasks | Action Mode | Description |
+|--------|---------|-------|-------------|-------------|
+| `jaxarc/default` | MiniARC | 70 | point | Small dataset for fast prototyping |
+| `jaxarc/agi1_train` | ARC-AGI-1 | 400 | point | Official ARC training benchmark |
+| `jaxarc/agi1_bbox` | ARC-AGI-1 | 400 | bbox | Bounding-box actions for structured edits |
+| `jaxarc/concept` | ConceptARC | ~40 | point | Tasks grouped by reasoning concept |
+
+## Customising Behaviour
+
+### Datasets
+
+JaxARC supports four ARC datasets. Set the dataset via the `env.scenario.name` config:
+
+- **Mini** (`env.scenario.name=Mini`): 70 simplified tasks from MiniARC. Good for debugging and prototyping.
+- **AGI1** (`env.scenario.name=AGI1-train`): 400 official ARC-AGI-1 training tasks. The standard RL benchmark. Use `AGI1-eval` for the evaluation split.
+- **AGI2** (`env.scenario.name=AGI2-train`): ARC-AGI-2 competition tasks. Use `AGI2-eval` for the evaluation split.
+- **ConceptARC** (`env.scenario.name=Concept-SameDifferent`): Tasks grouped by abstract reasoning concepts. Replace `SameDifferent` with any concept name.
+
+Available ConceptARC concepts: `AboveBelow`, `Center`, `CleanUp`, `CompleteShape`, `Copy`, `Count`, `ExtendToBoundary`, `ExtractObjects`, `FilledNotFilled`, `HorizontalVertical`, `InsideOutside`, `MoveToBoundary`, `Order`, `SameDifferent`, `TopBottom2D`, `TopBottom3D`.
+
+### Action Modes
+
+JaxARC offers three action modes. Override via `env.action.mode`:
+
+```bash
+# Point actions (default) — select individual cells
+python stoix/systems/ppo/anakin/ff_ppo.py env=jaxarc/default env.action.mode=point
+
+# Bounding-box actions — select rectangular regions
+python stoix/systems/ppo/anakin/ff_ppo.py env=jaxarc/default env.action.mode=bbox
+```
+
+- **point**: Agent selects one cell at a time (row, col, color). Simple but large action space.
+- **bbox**: Agent selects rectangular regions (top-left, bottom-right, color). Enables structured edits.
+- **mask**: Full grid mask specification. Most expressive but largest action space.
+
+### Observation Wrappers
+
+Three observation components can be toggled independently:
+
+```yaml
+observation_wrappers:
+  answer_grid: true   # Include the current answer grid state
+  input_grid: true    # Include the input grid (task specification)
+  contextual: true    # Include task context (episode step, etc.)
+```
+
+## Key Metrics
+
+JaxARC provides domain-specific metrics alongside the standard Stoix `episode_return` and `episode_length`:
+
+| Metric | Description |
+|--------|-------------|
+| `best_similarity` | Peak grid similarity achieved during the episode (0–1) |
+| `solved` | Whether the agent produced a perfect solution |
+| `steps_to_solve` | Number of steps taken to first solve (if solved) |
+| `final_similarity` | Grid similarity at episode termination (0–1) |
+| `was_truncated` | Whether the episode hit the step limit without solving |
+
+These are emitted via JaxARC's `ExtendedMetrics` wrapper and appear in `timestep.extras["episode_metrics"]`.

pyproject.toml

@@ -141,6 +141,7 @@ dependencies = [
     "gymnasium>=1.1.1",
     "huggingface-hub>=0.30.1",
     "hydra-core>=1.3.2",
+    "jaxarc>=1.0.3",
     "jax>=0.4.25,<0.6.0",
     "jaxlib>=0.5.3",
     "jumanji>=1.0.0",

stoix/configs/env/jaxarc/agi1_bbox.yaml

@@ -0,0 +1,20 @@
+# ARC-AGI-1 with bounding-box actions — demonstrates bbox action mode.
+# Bbox actions let the agent select rectangular regions instead of individual cells,
+# reducing the action space as compared to mask selection and enabling more structured edits.
+env_name: jaxarc
+
+scenario:
+  name: AGI1-train
+
+action:
+  mode: bbox  # point | bbox | mask
+
+observation_wrappers:
+  answer_grid: true
+  input_grid: true
+  contextual: true
+
+eval_metric: episode_return
+
+wrapper:
+  _target_: stoa.FlattenObservationWrapper

stoix/configs/env/jaxarc/agi1_train.yaml

@@ -0,0 +1,19 @@
+# ARC-AGI-1 Training Set — Standard ARC benchmark for RL research.
+# Uses the official training split (400 tasks) with point-based actions.
+env_name: jaxarc
+
+scenario:
+  name: AGI1-train
+
+action:
+  mode: point  # point | bbox | mask
+
+observation_wrappers:
+  answer_grid: true
+  input_grid: true
+  contextual: true
+
+eval_metric: episode_return
+
+wrapper:
+  _target_: stoa.FlattenObservationWrapper

stoix/configs/env/jaxarc/concept.yaml

@@ -0,0 +1,23 @@
+# ConceptARC — SameDifferent concept group.
+# ConceptARC organises tasks by abstract reasoning concepts.
+# Use "Concept-{ConceptName}" as the scenario to focus on a single concept:
+#   AboveBelow, Center, CleanUp, CompleteShape, Copy, Count,
+#   ExtendToBoundary, ExtractObjects, FilledNotFilled, HorizontalVertical,
+#   InsideOutside, MoveToBoundary, Order, SameDifferent, TopBottom2D, TopBottom3D
+env_name: jaxarc
+
+scenario:
+  name: Concept-SameDifferent
+
+action:
+  mode: point  # point | bbox | mask
+
+observation_wrappers:
+  answer_grid: true
+  input_grid: true
+  contextual: true
+
+eval_metric: episode_return
+
+wrapper:
+  _target_: stoa.FlattenObservationWrapper

stoix/configs/env/jaxarc/default.yaml

@@ -0,0 +1,23 @@
+# ---Environment Configs---
+env_name: jaxarc
+
+scenario:
+  name: "Mini-Colony_Expansion_2_l6d2q1yx5npw0qdrnzn"
+  task_name: mini
+
+action:
+  mode: point  # point | bbox | mask
+
+observation_wrappers:
+  answer_grid: true
+  input_grid: true
+  contextual: true
+
+# Defines the metric that will be used to evaluate the performance of the agent.
+# This metric is returned at the end of an experiment and can be used for hyperparameter tuning.
+eval_metric: episode_return
+
+# Flatten multi-dimensional observation (4, 5, 5, 13) -> (1300,) for MLP networks.
+# Remove this wrapper when using CNN networks (e.g., network=jaxarc/arc_shallow_cnn).
+wrapper:
+  _target_: stoa.FlattenObservationWrapper

stoix/utils/make_env.py

@@ -304,6 +304,50 @@ def make_debug_env(scenario_name: str, config: DictConfig) -> Tuple[Environment,
     return env, eval_env
 
 
+def make_jaxarc_env(scenario_name: str, config: DictConfig) -> Tuple[Environment, Environment]:
+    """Creates and wraps a JaxARC environment for ARC puzzle tasks.
+
+    JaxARC environments are natively Stoa-compatible, so no adapter is needed.
+    Domain-specific wrapper application (action, observation) is delegated to
+    JaxARC's factory function. ExtendedMetrics is applied after RecordEpisodeMetrics
+    so that domain-specific metrics merge correctly with episode metrics.
+    """
+    try:
+        from jaxarc.stoix_adapter import make_jaxarc_env as _make_jaxarc
+        from jaxarc.wrappers import ExtendedMetrics
+    except ImportError as e:
+        raise ImportError(
+            "JaxARC is required for 'jaxarc' environments. Install it with: pip install jaxarc"
+        ) from e
+
+    env, eval_env = _make_jaxarc(config)
+    env, eval_env = apply_optional_wrappers((env, eval_env), config)
+
+    # Apply core wrappers with ExtendedMetrics injected after RecordEpisodeMetrics.
+    # ExtendedMetrics must come after REM so it can merge its fields into the
+    # episode_metrics dict that REM creates, rather than being overwritten by it.
+    # Note: only train_env gets core wrappers (standard Stoix pattern).
+    env = AddRNGKey(env)
+    env = RecordEpisodeMetrics(env)
+    env = ExtendedMetrics(env)
+
+    if config.env.get("use_optimistic_reset", False):
+        env = OptimisticResetVmapWrapper(
+            env,
+            config.arch.num_envs,
+            min(config.env.get("reset_ratio", 16), config.arch.num_envs),
+            next_obs_in_extras=True,
+        )
+    else:
+        if config.env.get("use_cached_auto_reset", False):
+            env = CachedAutoResetWrapper(env, next_obs_in_extras=True)
+        else:
+            env = AutoResetWrapper(env, next_obs_in_extras=True)
+        env = VmapWrapper(env)
+
+    return env, eval_env
+
+
 def make_popjym_env(scenario_name: str, config: DictConfig) -> Tuple[Environment, Environment]:
     """Creates and wraps a POPJym environment."""
     import popjym
@@ -378,6 +422,7 @@ def make_playground_env(scenario_name: str, config: DictConfig) -> Tuple[Environ
     "gymnax": make_gymnax_env,
     "brax": make_brax_env,
     "craftax": make_craftax_env,
+    "jaxarc": make_jaxarc_env,
     "popgym_arcade": make_popgym_arcade_env,
     "xland_minigrid": make_xland_minigrid_env,
     "popjym": make_popjym_env,