Jumanji Documentation

The SlidingTilePuzzle is a JAX JIT-able implementation of the classic sliding tile puzzle game. Players move tiles into an empty space to reach a target configuration. It is available via the Jumanji registry.

Registered Versions:

• SlidingTilePuzzle-v0: A version with a 5x5 grid size.

The Maze environment is a JAX JIT-able 2D maze problem. The agent (green) must navigate a size-configurable 2D matrix of free space and walls to reach a single target cell (red).

Key characteristics:

• Sparse Reward: The agent receives a reward of 0 at every step and a reward of 1 only upon reaching the target.
• Movement: The agent can move Up, Right, Down, or Left. If a move is blocked by a wall, the agent remains in its current position.
• Generation: Mazes are randomly generated using a recursive division function. By default, the maze, agent position, and target position are regenerated on every reset().

The Sokoban environment is a JAX implementation of the classic box-pushing puzzle. The agent's goal is to move all boxes onto their designated target locations. This implementation follows the rules from the DeepMind 'Imagination Augmented Agents' paper and uses levels from the Boxoban dataset.

Key characteristics:

• Goal: Place all 4 boxes on their targets.
• Difficulty Splits: Available levels include 'unfiltered', 'medium', and 'hard'.
• Registered Version: Sokoban-v0 (uses unfiltered training levels).

The GraphColoring environment is a JAX JIT-able, episodic, single-agent environment designed for combinatorial optimization. The objective is to assign a color to each vertex of a graph such that no two adjacent vertices share the same color, while minimizing the total number of unique colors used.

Key characteristics:

• Problem Type: Combinatorial optimization (Graph Coloring).
• Implementation: JAX JIT-able.
• Goal: Find a valid coloring using the minimum number of colors possible.

The Multi Minimum Spanning Tree (MMST) environment involves a random connected graph with groups of nodes (of the same type) that must be connected. The objective is to connect all nodes within each group using the shortest path possible, without reusing utility nodes (nodes that do not belong to any group).

Key Characteristics:

• Goal: Connect all nodes of the same type together.
• Termination: An episode ends when all groups of nodes are connected or the maximum number of steps is reached.
• Agent Model: While it can be viewed as a multi-agent problem, this implementation treats it as a single agent that outputs multiple actions per step (one for each group/agent).

Registered Versions:

• MMST-v0: 3 agents, 36 nodes, 72 edges, 4 nodes to connect per agent, and a step limit of 70.

Jumanji provides a set of wrappers in the jumanji.wrappers module to modify environment behavior, transform observations, or adapt Jumanji environments to other interfaces (like Gymnasium).

Key wrappers include:

• Wrapper: The base class for all wrappers.
• JumanjiToDMEnvWrapper: Adapts a Jumanji environment to the DeepMind Control Suite (DMControl) interface.
• JumanjiToGymWrapper: Adapts a Jumanji environment to the Gymnasium interface.
• MultiToSingleWrapper: Transforms environments with multiple observation/action spaces into a single space.
• VmapWrapper: Enables vectorization (vmap) of the environment.
• AutoResetWrapper: Automatically resets the environment when a terminal state is reached.
• jumanji_to_gym_obs: A utility to transform Jumanji observations into a format compatible with Gymnasium.

Jumanji records two distinct types of evaluation during the training process:

• Stochastic evaluation: Uses the same policy (sampling from the distribution) used during the training phase.
• Greedy evaluation: Uses the argmax over the action logits to select the most likely action.

The observation returned at each step is a tuple containing three JAX arrays:

1. grid: A jax.array (int32) of shape (grid_size, grid_size). This 2D matrix represents the positions, targets, and paths of all agents.
2. action_mask: A jax.array (bool) of shape (num_agents, 5). This indicates the valid actions available to each agent.
3. step_count: A jax.array (int32) of shape () representing the number of steps taken since the last reset.

Grid Encoding

Agents are encoded in groups of 3 integers. For any agent $i$ (where $i$ starts at 0), the components are mapped as follows:

• Path: 1 + (3 * i)
• Position: 2 + (3 * i)
• Target: 3 + (3 * i)

Example Encoding (Agent 1, 2, and 3):

• Agent 1: Path=1, Position=2, Target=3
• Agent 2: Path=4, Position=5, Target=6
• Agent 3: Path=7, Position=8, Target=9

The Search & Rescue environment is a multi-agent simulation where a group of agents searches a 2D square space (with wrapped boundaries) for multiple targets.

Core Mechanics:

• Goal: Agents aim to locate targets within a fixed number of steps.
• Update Sequence:
  1. Agent velocities and positions are updated.
  2. Target positions are updated.
  3. Targets within an agent's detection range and view cone are marked as 'found'.
  4. Agents receive rewards for finding previously undetected targets.
  5. Local views (observations) are generated.
• Customization: You can extend the environment by implementing the following interfaces:
  • ObservationFn: To customize agent observations.
  • RewardFn: To customize reward logic.
  • TargetDynamics: To model different target movement scenarios.

The reward is dense: a reward of -1 is given at each time step until termination.

Termination Scenarios:

1. Finished schedule: All operations for all jobs have been processed.
2. Illegal action: The agent selects an action that violates the action_mask.
3. Simultaneously idle: All machines are inactive at the same time.

Penalties: If the agent selects an invalid action or all machines become simultaneously idle, a large penalty is applied: -num_jobs * max_num_ops * max_op_duration. This represents an upper bound on the possible makespan.