XLand-MiniGrid

XLand-MiniGrid provides a JAX-based interface that combines elements of dm_env and gymnax. The API is designed to be fully compatible with JAX transformations like jit.

Key workflow steps:

1. Load a benchmark: Use xminigrid.load_benchmark(name=...) to get a collection of tasks.
2. Sample a ruleset: Use benchmark.sample_ruleset(key) to select a specific task configuration.
3. Create an environment: Use xminigrid.make(env_id) to instantiate an environment.
4. Configure parameters: Replace the default env_params with your sampled ruleset.
5. Apply wrappers: Use GymAutoResetWrapper for automatic resets and RGBImgObservationWrapper if you need RGB image observations (note: this significantly impacts performance).
6. Step and Reset: Use env.reset and env.step within jax.jit for high performance.

import jax
import xminigrid
from xminigrid.wrappers import GymAutoResetWrapper
from xminigrid.experimental.img_obs import RGBImgObservationWrapper

key = jax.random.key(0)
reset_key, ruleset_key = jax.random.split(key)

# List available benchmarks via xminigrid.registered_benchmarks()
benchmark = xminigrid.load_benchmark(name="trivial-21k")
# Sample a ruleset
ruleset = benchmark.sample_ruleset(ruleset_key)

# List available environments via xminigrid.registered_environments()
env, env_params = xminigrid.make("XLand-MiniGrid-R9-25x25")
env_params = env_params.replace(ruleset=ruleset)

# Apply wrappers
env = GymAutoResetWrapper(env)
env = RGBImgObservationWrapper(env)

# JIT-compatible step and reset
timestep = jax.jit(env.reset)(env_params, reset_key)
timestep = jax.jit(env.step)(env_params, timestep, action=0)

# Optionally render the state
env.render(env_params, timestep)

To avoid the overhead of manual rule/goal composition and to facilitate reliable comparisons, XLand-MiniGrid provides pre-sampled benchmarks. These benchmarks contain up to three million unique tasks, where each task is represented as a tree of production rules leading to a goal.

Benchmarks are hosted on HuggingFace and are automatically downloaded and cached to ~/.xland_minigrid (or the path specified by the XLAND_MINIGRID_DATA environment variable) upon first use.

Common operations:

• load_benchmark(name=...): Loads a specific benchmark.
• sample_ruleset(key): Samples a random ruleset from the benchmark.
• get_ruleset(ruleset_id=...): Retrieves a specific ruleset by ID.
• shuffle(key=...).split(prop=...): Splits the benchmark into training and testing sets.

import jax.random
import xminigrid
from xminigrid.benchmarks import Benchmark

# Downloading to path specified by XLAND_MINIGRID_DATA, ~/.xland_minigrid by default
benchmark: Benchmark = xminigrid.load_benchmark(name="trivial-21k")

# Users can sample or get specific rulesets
benchmark.sample_ruleset(jax.random.key(0))
benchmark.get_ruleset(ruleset_id=benchmark.num_rulesets() - 1)

# Split them for train & test
train, test = benchmark.shuffle(key=jax.random.key(0)).split(prop=0.8)

The complexity of XLand-MiniGrid comes from dynamic Rules and Goals.

• Rules: Functions that change the environment state deterministically based on conditions.
• Goals: Functions that test conditions without changing the state.

Encoding for JAX Compatibility

To remain JIT-compatible and efficient, rules and goals are represented as vectors (encodings). The first element is always an ID, followed by arguments. This allows:

1. Dynamic rulesets (changing arguments between resets).
2. Compact storage for millions of tasks.
3. Efficient decoding using jax.lax.switch inside JIT loops.

Working with Encodings

Instead of calling class methods directly, use check_goal and check_rule with encodings.

from xminigrid.core.goals import AgentNearGoal, check_goal
from xminigrid.core.rules import AgentNearRule, check_rule

# Create and encode a goal
gol = AgentNearGoal(tile=TILES_REGISTRY[Tiles.BALL, Colors.YELLOW])
goal_encoding = gol.encode()

# Check goal using encoding (efficient for JIT)
is_goal_met = check_goal(goal_encoding, grid, agent, action, position)

# Create and encode a rule
rule = AgentNearRule(tile=..., prod_tile=...)
rule_encoding = rule.encode()

# Check rule using encoding
new_grid, new_agent = check_rule(rule_encoding[None, ...], grid, agent, action, position)

from xminigrid.core.goals import AgentNearGoal, check_goal
from xminigrid.core.rules import AgentNearRule, check_rule

# Goal usage
gol = AgentNearGoal(tile=TILES_REGISTRY[Tiles.BALL, Colors.YELLOW])
goal_encoding = gol.encode()
is_goal_met = check_goal(goal_encoding, grid, agent, action, position)

# Rule usage
rule = AgentNearRule(tile=..., prod_tile=...)
rule_encoding = rule.encode()
new_grid, new_agent = check_rule(rule_encoding[None, ...], grid, agent, action, position)

The Benchmark class is a JAX-compatible PyTreeNode that holds the dataset. It provides several methods for sampling, splitting, and filtering data, which are all JIT-compatible.

• get_ruleset(ruleset_id): Retrieves a specific RuleSet by its ID.
• sample_ruleset(key): Samples a random RuleSet using a JAX random key.
• shuffle(key): Returns a new Benchmark instance with the rulesets shuffled.
• split(prop): Splits the benchmark into two parts based on a proportion prop (0.0 to 1.0).
• filter_split(fn): Splits the benchmark into two based on a predicate function fn(single_goal, single_rules) -> bool.

import jax
from xminigrid.benchmarks import load_benchmark

benchmark = load_benchmark("trivial-21k")
key = jax.random.PRNGKey(0)

# Sample a random ruleset
random_ruleset = benchmark.sample_ruleset(key)

# Split the benchmark into 80% train and 20% test
train_bench, test_bench = benchmark.split(0.8)

# Filter benchmark based on a condition
def my_filter(goal, rules):
    return goal.shape[0] > 5 # Example condition

bench_a, bench_b = benchmark.filter_split(my_filter)

The xminigrid environments are designed to be stateless and fully JIT-compatible. The environment instance itself does not hold state; instead, all state information is passed through and returned via EnvParams and TimeStep containers.

Core Components

• EnvParams: A PyTreeNode containing configuration parameters (e.g., grid size, rulesets). Use env.default_params(**kwargs) to create them.
• TimeStep: A PyTreeNode containing the current state (hidden environment state like grid, agent, goal), step_type, reward, discount, and observation.

Main Methods

• reset(params: EnvParams, key: jax.Array) -> TimeStep
• step(params: EnvParams, timestep: TimeStep, action: int) -> TimeStep
• render(params: EnvParams, timestep: TimeStep) -> jax.Array (returns RGB images)
• observation_shape(params: EnvParams) -> tuple[int, int, int]
• num_actions(params: EnvParams) -> int

import xminigrid
import jax

key = jax.random.key(0)
key, reset_key = jax.random.split(key)

# Create environment and params
env, env_params = xminigrid.make("MiniGrid-Empty-8x8")

# JIT-compatible reset and step
timestep = jax.jit(env.reset)(env_params, reset_key)
timestep = jax.jit(env.step)(env_params, timestep, action=0)

XLand-MiniGrid uses JAX-compatible PyTree structures to represent the environment state and transitions. This allows for fully JIT-compatible training and simulation.

• State[EnvCarryT]: Represents the core environment state at a specific moment. It contains a key (for randomness), step_num, and the following fields:

  • grid: The GridState (a jax.Array).
  • agent: An AgentState object containing position, direction, and pocket.
  • goal_encoding: Encoding of the current goal.
  • rule_encoding: Encoding of the current rules.
  • carry: An EnvCarryT object (user-defined or environment-specific carry state).

• TimeStep[EnvCarryT]: Represents a single transition in a trajectory. It contains:

  • state: The State at that time step.
  • step_type: A StepType indicating if this is the FIRST, MID, or LAST step of an episode.
  • reward: The scalar reward received.
  • discount: The discount factor.
  • observation: The observation, which can be a jax.Array or a dict[str, jax.Array].

# Example conceptual structure of a TimeStep
# (Note: Actual instantiation depends on the specific environment API)
# timestep.state.agent.position
# timestep.reward
# timestep.observation