Jumanji

https://instadeepai.github.io/jumanji/
Admin
Jumanji is a diverse suite of scalable reinforcement learning environments written in JAX, featuring...
Tokens:98,257
Snippets:696
Trust Score:9.7
Update:6 months ago
Show doc for...
Tokens:
Raw
### Minesweeper Environment Usage Example

Source: https://instadeepai.github.io/jumanji/api/environments/minesweeper

Example of how to reset, step, and render the Minesweeper environment.

```APIDOC
## Minesweeper Environment Usage

### Description
This snippet demonstrates the basic usage of the Minesweeper environment, including resetting, taking actions, and rendering the state.

### Method
`reset`, `step`, `render`

### Endpoint
N/A (Environment Interaction)

### Parameters
None

### Request Example
```python
import jax
from jumanji.environments import Minesweeper

# Instantiate the environment
env = Minesweeper()

# Initialize a PRNG key for JAX
key = jax.random.PRNGKey(0)

# Reset the environment to get the initial state and timestep
state, timestep = jax.jit(env.reset)(key)

# Render the initial state
env.render(state)

# Generate a valid action (example: exploring a square)
action = env.action_spec.generate_value() 

# Take a step in the environment
state, timestep = jax.jit(env.step)(state, action)

# Render the state after the step
env.render(state)
```

### Response
None

### Response Example
None
```

--------------------------------

### Initialize and Step Level-Based Foraging

Source: https://instadeepai.github.io/jumanji/api/environments/lbf

Basic usage example showing how to reset the environment, perform a step, and render the state using JAX.

```python
from jumanji.environments import LevelBasedForaging
env = LevelBasedForaging()
key = jax.random.key(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### Instantiate and Use Snake Environment

Source: https://instadeepai.github.io/jumanji/api/environments/snake

Demonstrates how to instantiate the Snake environment, reset it, take a step, and render the state. Requires JAX to be installed.

```python
from jumanji.environments import Snake
env = Snake()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### Initialize and Use Sudoku Environment

Source: https://instadeepai.github.io/jumanji/api/environments/sudoku

Demonstrates how to initialize the Sudoku environment, reset it to get an initial state and timestep, and perform a step with a sample action. Includes rendering the state before and after the step.

```python
from jumanji.environments import Sudoku
env = Sudoku()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### reset()

Source: https://instadeepai.github.io/jumanji/api/wrappers

Starts a new sequence and returns the first TimeStep of the environment.

```APIDOC
## reset()

### Description
Starts a new sequence and returns the first `TimeStep` of this sequence.

### Response
#### Success Response (200)
- **TimeStep** (namedtuple) - A `TimeStep` containing:
    - **step_type**: A `StepType` of `FIRST`.
    - **reward**: `None`.
    - **discount**: `None`.
    - **observation**: A NumPy array, or a nested dict, list or tuple of arrays conforming to `observation_spec`.
```

--------------------------------

### MMST Environment Setup and Usage

Source: https://instadeepai.github.io/jumanji/api/environments/mmst

This snippet demonstrates how to initialize, reset, and step through the MMST environment using Jumanji and JAX.

```APIDOC
## MMST Environment Setup and Usage

### Description
This section provides a basic example of how to instantiate, reset, and interact with the MMST environment. It covers the essential steps for starting an episode and performing a single step.

### Method
```python
from jumanji.environments import MMST
import jax

# Initialize the environment
env = MMST()

# Create a JAX random key
key = jax.random.PRNGKey(0)

# Reset the environment to get the initial state and timestep
state, timestep = jax.jit(env.reset)(key)

# Render the initial state (optional)
env.render(state)

# Generate a valid action (replace with actual action selection logic)
action = env.action_spec.generate_value()

# Take a step in the environment
state, timestep = jax.jit(env.step)(state, action)

# Render the state after the step (optional)
env.render(state)
```

### Parameters

- **generator** (Optional[Generator]) - Generator whose `__call__` instantiates an environment instance. Implemented options are [`SplitRandomGenerator`]. Defaults to `SplitRandomGenerator(num_nodes=36, num_edges=72, max_degree=5, num_agents=3, num_nodes_per_agent=4, max_step=time_limit)`.
- **reward_fn** (Optional[RewardFn]) - Class of type `RewardFn`, whose `__call__` is used as a reward function. Implemented options are [`DenseRewardFn`]. Defaults to `DenseRewardFn(reward_values=(10.0, -1.0, -1.0))`.
- **time_limit** (int) - The number of steps allowed before an episode terminates. Defaults to 70.
- **viewer** (Optional[Viewer[State]]) - Viewer used for rendering. Defaults to `MMSTViewer`.

### Request Example
```json
{
  "example": "No specific request body for environment initialization, parameters are passed during instantiation."
}
```

### Response
#### Success Response (200)
- **state** (State) - The current state of the environment.
- **timestep** (TimeStep) - The timestep returned by the environment after a reset or step.

#### Response Example
```json
{
  "example": "The response includes a state object and a timestep object, which contain detailed information about the environment's status and the results of the last action."
}
```
```

--------------------------------

### Instantiate and Use RubiksCube Environment

Source: https://instadeepai.github.io/jumanji/api/environments/rubiks_cube

Instantiate the RubiksCube environment, reset it to get an initial state and timestep, and then perform a step with a generated action. Rendering is shown before and after the step.

```python
from jumanji.environments import RubiksCube
env = RubiksCube()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### Action Mask Example

Source: https://instadeepai.github.io/jumanji/environments/tetris

Indicates valid placement actions for different rotations of the current tetromino.

```python
[
    [ True,  False,  True,  True, False, False],
    [ True,  True, False,  False,  False, False],
    [ True,  False,  True,  True, False, False],
    [ True,  True, False,  False,  False, False],
]
```

--------------------------------

### Initialize and Interact with Minesweeper Environment

Source: https://instadeepai.github.io/jumanji/api/environments/minesweeper

Instantiate the Minesweeper environment, reset it to get the initial state and timestep, render the board, generate a valid action, and then step the environment with that action. This demonstrates a basic interaction loop.

```python
from jumanji.environments import Minesweeper
env = Minesweeper()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### Connector Environment Observation Grid Example

Source: https://instadeepai.github.io/jumanji/environments/connector

A sample observation grid for the Connector environment on a 6x6 grid, showing the spatial arrangement of agent paths, positions, and targets. Zeroes represent empty cells.

```plaintext
[[ 2  0  3  0  0  0]
 [ 1  0  4  4  4  0]
 [ 1  0  5  9  0  0]
 [ 1  0  0  0  0  0]
 [ 0  0  0  8  0  0]
 [ 0  0  6  7  7  7]]
```

--------------------------------

### Initialize and Step the FlatPack Environment

Source: https://instadeepai.github.io/jumanji/api/environments/flat_pack

Demonstrates basic environment instantiation, resetting, and stepping using JAX JIT compilation.

```python
from jumanji.environments import FlatPack
env = FlatPack()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### Install Jumanji from PyPI

Source: https://instadeepai.github.io/jumanji

Install the latest release of Jumanji using pip. Ensure you have the correct JAX version installed based on your hardware.

```bash
pip install -U jumanji
```

--------------------------------

### Initialize and Step the Knapsack Environment

Source: https://instadeepai.github.io/jumanji/api/environments/knapsack

Demonstrates how to instantiate the environment, reset it, and perform a step using JAX JIT compilation.

```python
from jumanji.environments import Knapsack
env = Knapsack()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### Install Jumanji from GitHub

Source: https://instadeepai.github.io/jumanji

Install the latest development version of Jumanji directly from its GitHub repository using pip.

```bash
pip install git+https://github.com/instadeepai/jumanji.git
```

--------------------------------

### Instantiate and Use JobShop Environment

Source: https://instadeepai.github.io/jumanji/api/environments/job_shop

Demonstrates how to instantiate the JobShop environment, reset it, render the initial state, take a step with a generated action, and render the resulting state.

```python
from jumanji.environments import JobShop
env = JobShop()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### Initialize and Use Connector Environment

Source: https://instadeepai.github.io/jumanji/api/environments/connector

Demonstrates how to initialize the Connector environment, reset it, render the initial state, take an action, and step through the environment. Requires JAX.

```python
from jumanji.environments import Connector
env = Connector()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### 2048 Game Board Observation Example

Source: https://instadeepai.github.io/jumanji/environments/game_2048

Example of a random observation of the 2048 game board represented as a JAX array of exponents.

```text
1
2
3
4
```

```text
[[ 2 0 1 4]
 [ 5 3 0 2]
 [ 0 2 3 2]
 [ 1 2 0 0]]
```

--------------------------------

### Sliding Tile Puzzle Observation Example

Source: https://instadeepai.github.io/jumanji/environments/sliding_tile_puzzle

An example of a random observation for the Sliding Tile Puzzle game, showing the puzzle grid, empty tile position, and action mask.

```python
[[ 1 2 3 4]
 [ 5 6 7 8]
 [ 9 10 0 12]
 [ 13 14 15 11]]
```

--------------------------------

### Initialize and step the SearchAndRescue environment

Source: https://instadeepai.github.io/jumanji/api/environments/search_and_rescue

Demonstrates how to instantiate the environment, reset it, and perform a step using JAX JIT compilation.

```python
from jumanji.environments import SearchAndRescue

env = SearchAndRescue()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### Instantiate and Use Cleaner Environment

Source: https://instadeepai.github.io/jumanji/api/environments/cleaner

Demonstrates how to instantiate the Cleaner environment, reset it, render the initial state, take a step with a sample action, and render the resulting state.

```python
from jumanji.environments import Cleaner
env = Cleaner()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### GET observation_spec

Source: https://instadeepai.github.io/jumanji/api/environments/graph_coloring

Retrieve the observation specification for the environment.

```APIDOC
## GET observation_spec

### Description
Returns the observation spec for the environment.

### Response
- **adj_matrix** (BoundedArray) - Boolean array of shape (num_nodes, num_nodes) representing the adjacency matrix.
- **action_mask** (BoundedArray) - Boolean array of shape (num_nodes,) representing valid actions.
- **colors** (BoundedArray) - Int32 array of shape (num_nodes,) representing assigned colors.
- **current_node_index** (BoundedArray) - Int32 array of shape () representing the current node index.
```

--------------------------------

### GET action_spec

Source: https://instadeepai.github.io/jumanji/api/environments/graph_coloring

Retrieve the action specification for the environment.

```APIDOC
## GET action_spec

### Description
Returns the specification of the action for the `GraphColoring` environment.

### Response
- **action_spec** (DiscreteArray) - A `specs.DiscreteArray` object defining valid actions.
```

--------------------------------

### Instantiate and Use GraphColoring Environment

Source: https://instadeepai.github.io/jumanji/api/environments/graph_coloring

Demonstrates how to instantiate the GraphColoring environment, reset it, render the initial state, take a step with a generated action, and render the resulting state.

```python
from jumanji.environments import GraphColoring
env = GraphColoring()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### GET observation_spec

Source: https://instadeepai.github.io/jumanji/api/environments/cvrp

Returns the observation specification for the CVRP environment.

```APIDOC
## GET observation_spec

### Description
Returns the observation spec for the environment.

### Method
GET

### Response
- **coordinates** (BoundedArray) - float, shape (num_nodes + 1, 2)
- **demands** (BoundedArray) - float, shape (num_nodes + 1,)
- **unvisited_nodes** (BoundedArray) - bool, shape (num_nodes + 1,)
- **position** (DiscreteArray) - num_values = num_nodes + 1, shape ()
- **trajectory** (BoundedArray) - int32, shape (2 * num_nodes,)
- **capacity** (BoundedArray) - float, shape ()
- **action_mask** (BoundedArray) - bool, shape (num_nodes + 1,)
```

--------------------------------

### Initialize and Step CVRP Environment

Source: https://instadeepai.github.io/jumanji/api/environments/cvrp

Demonstrates how to instantiate the CVRP environment, reset it, and perform a step using JAX.

```python
from jumanji.environments import CVRP
env = CVRP()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### Initialize and Step PacMan Environment

Source: https://instadeepai.github.io/jumanji/api/environments/pac_man

Demonstrates how to instantiate the environment, reset it, and perform a step using JAX JIT compilation.

```python
from jumanji.environments import pac_man
env = PacMan()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### GET action_spec

Source: https://instadeepai.github.io/jumanji/api/environments/cvrp

Returns the action specification for the CVRP environment.

```APIDOC
## GET action_spec

### Description
Returns the action spec for the environment.

### Method
GET

### Response
- **action_spec** (DiscreteArray) - a specs.DiscreteArray spec.
```

--------------------------------

### GET observation_spec

Source: https://instadeepai.github.io/jumanji/api/environments/cleaner

Retrieves the specification of the observation for the Cleaner environment.

```APIDOC
## GET observation_spec

### Description
Returns the specification of the observation, including grid, agent locations, action mask, and step count.

### Response
#### Success Response (200)
- **grid** (BoundedArray) - int8 array of shape (num_rows, num_cols), values 0-2.
- **agent_locations_spec** (BoundedArray) - int32 array of shape (num_agents, 2).
- **action_mask** (BoundedArray) - bool array of shape (num_agent, 4).
- **step_count** (BoundedArray) - int32 array of shape ().
```

--------------------------------

### GET action_spec

Source: https://instadeepai.github.io/jumanji/api/environments/rubiks_cube

Retrieve the action specification for the RubiksCube environment.

```APIDOC
## GET action_spec

### Description
Returns the action spec. An action is composed of 3 elements that range in: 6 faces, each with cube_size//2 possible depths, and 3 possible directions.

### Response
- **action_spec** (MultiDiscreteArray) - The action specification object.
```

--------------------------------

### Initialize and Step the Game2048 Environment

Source: https://instadeepai.github.io/jumanji/api/environments/game_2048

Demonstrates how to instantiate the environment, reset it using a JAX PRNG key, and perform a step with a generated action.

```python
from jumanji.environments import Game2048
env = Game2048()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### get_finished_agents

Source: https://instadeepai.github.io/jumanji/api/environments/mmst

Get the done flags for each agent based on the environment state.

```APIDOC
## get_finished_agents

### Description
Get the done flags for each agent.

### Parameters
- **state** (State) - Required - The environment state.

### Response
- **finished_agents** (Array) - Array of boolean flags in the shape (number of agents, ).
```

--------------------------------

### Game2048 Initialization and Usage

Source: https://instadeepai.github.io/jumanji/api/environments/game_2048

Demonstrates how to initialize the Game2048 environment, reset the state, render the board, and take a step.

```APIDOC
## Game2048 Initialization and Usage

### Description
This snippet shows the basic usage of the Game2048 environment, including initialization, resetting the game, rendering the current state, and performing a game step.

### Method
N/A (Illustrative Example)

### Endpoint
N/A (Environment Interaction)

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Request Example
```python
from jumanji.environments import Game2048
import jax

key = jax.random.PRNGKey(0)
env = Game2048()
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value() # Example action generation
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

### Response
#### Success Response (200)
N/A (Environment interaction does not return HTTP responses)

#### Response Example
N/A
```

--------------------------------

### I-Tetromino Shape Example

Source: https://instadeepai.github.io/jumanji/environments/tetris

Represents the I-shaped tetromino as a 4x4 JAX array.

```python
[
    [1, 0, 0, 0],
    [1, 0, 0, 0],
    [1, 0, 0, 0],
    [1, 0, 0, 0],
]
```

--------------------------------

### Initialize and Step the TSP Environment

Source: https://instadeepai.github.io/jumanji/api/environments/tsp

Demonstrates how to instantiate the TSP environment, reset it, and perform a step using JAX JIT compilation.

```python
from jumanji.environments import TSP
env = TSP()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### GET observation_spec

Source: https://instadeepai.github.io/jumanji/api/environments/sudoku

Returns the observation specification containing the board and action_mask arrays.

```APIDOC
## GET observation_spec

### Description
Returns the observation spec containing the board and action_mask arrays.

### Response
#### Success Response (200)
- **Spec[Observation]** (Spec) - Spec containing all the specifications for all the Observation fields: - board: BoundedArray (jnp.int8) of shape (9,9). - action_mask: BoundedArray (bool) of shape (9,9,9).
```

--------------------------------

### Initialize and Step Tetris Environment

Source: https://instadeepai.github.io/jumanji/api/environments/tetris

Demonstrates how to instantiate the Tetris environment, reset it, and perform a step using JAX JIT compilation.

```python
from jumanji.environments import Tetris
env = Tetris()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### Get dm_env Observation Spec

Source: https://instadeepai.github.io/jumanji/api/wrappers

Retrieves the observation specification converted to the dm_env format.

```python
def observation_spec(self) -> dm_env.specs.Array:
    """Returns the dm_env observation spec."""
    return specs.jumanji_specs_to_dm_env_specs(self._env.observation_spec)
```

--------------------------------

### Initialize and Use MMST Environment

Source: https://instadeepai.github.io/jumanji/api/environments/mmst

Demonstrates how to initialize the MMST environment, reset it, render the initial state, take a step with a sample action, and render the resulting state. Ensure JAX is imported for JIT compilation and environment operations.

```python
from jumanji.environments import MMST
env = MMST()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

--------------------------------

### Get dm_env Action Spec

Source: https://instadeepai.github.io/jumanji/api/wrappers

Retrieves the action specification converted to the dm_env format.

```python
def action_spec(self) -> dm_env.specs.Array:
    """Returns the dm_env action spec."""
    return specs.jumanji_specs_to_dm_env_specs(self._env.action_spec)
```

--------------------------------

### Get Sokoban environment representation

Source: https://instadeepai.github.io/jumanji/api/environments/sokoban

Returns a string representation of the Sokoban environment configuration.

```python
def __repr__(self) -> str:
    """
    Returns a printable representation of the Sokoban environment.

    Returns:
        str: A string representation of the Sokoban environment.
    """
    return "\n".join(
        [
            "Bokoban environment:",
            f" - num_rows: {self.num_rows}",
            f" - num_cols: {self.num_cols}",
            f" - time_limit: {self.time_limit}",
            f" - generator: {self.generator}",
        ]
    )
```

--------------------------------

### PacMan Environment Initialization

Source: https://instadeepai.github.io/jumanji/api/environments/pac_man

Instantiates the PacMan environment with optional generator, viewer, and time limit configurations.

```APIDOC
## PacMan Environment Initialization

### Description
Instantiates a new PacMan environment instance for JAX-based reinforcement learning.

### Parameters
#### Request Body
- **generator** (Optional[Generator]) - Optional - Generator whose __call__ instantiates an environment instance. Implemented options are [AsciiGenerator].
- **time_limit** (Optional[int]) - Optional - The maximum number of environment steps before the episode terminates. Defaults to 1000.
- **viewer** (Optional[Viewer[State]]) - Optional - Viewer used for rendering. Defaults to PacManViewer.

### Request Example
{
  "time_limit": 1000
}
```

--------------------------------

### Wrapper.step(state, action)

Source: https://instadeepai.github.io/jumanji/api/wrappers

Runs one timestep of the environment's dynamics.

```APIDOC
## step(state, action)

### Description
Run one timestep of the environment's dynamics.

### Method
step

### Parameters
#### Path Parameters
- **state** (State) - Required - State object containing the dynamics of the environment.
- **action** (Array) - Required - Array containing the action to take.

### Returns
- **state** (State) - State object corresponding to the next state of the environment.
- **timestep** (TimeStep[Observation]) - TimeStep object corresponding the timestep returned by the environment.

### Source Code
```python
def step(self, state: State, action: chex.Array) -> Tuple[State, TimeStep[Observation]]:
    """Run one timestep of the environment's dynamics.

    Args:
        state: State object containing the dynamics of the environment.
        action: Array containing the action to take.

    Returns:
        state: State object corresponding to the next state of the environment,
        timestep: TimeStep object corresponding the timestep returned by the environment,
    """
    return self._env.step(state, action)
```
```

--------------------------------

### Grid Observation Example

Source: https://instadeepai.github.io/jumanji/environments/tetris

Represents the current state of the game grid as a 2D JAX array.

```python
[
    [0, 0, 0, 0, 0, 1],
    [0, 0, 0, 0, 1, 1],
    [0, 0, 0, 0, 1, 1],
    [0, 1, 0, 0, 1, 1],
    [0, 1, 1, 1, 0, 1],
    [0, 1, 0, 1, 1, 1],
    [1, 1, 0, 1, 1, 1],
]
```

--------------------------------

### Cleaner Environment Initialization

Source: https://instadeepai.github.io/jumanji/api/environments/cleaner

Shows the constructor for the Cleaner environment, detailing parameters like generator, time limit, timestep penalty, and viewer.

```python
def __init__(
    self,
    generator: Optional[Generator] = None,
    time_limit: Optional[int] = None,
    penalty_per_timestep: float = 0.5,
    viewer: Optional[Viewer[State]] = None,
) -> None:
    """Instantiates a `Cleaner` environment.

    Args:
        generator: `Generator` whose `__call__` instantiates an environment instance.
            Implemented options are [`RandomGenerator`]. Defaults to `RandomGenerator` with
            `num_rows=10`, `num_cols=10` and `num_agents=3`.
        time_limit: max number of steps in an episode. Defaults to `num_rows * num_cols`.
        penalty_per_timestep: the penalty returned at each timestep in the reward.
        viewer: `Viewer` used for rendering. Defaults to `CleanerViewer` with "human" render
            mode.
    """
    self.generator = generator or RandomGenerator(num_rows=10, num_cols=10, num_agents=3)
    self.num_agents = self.generator.num_agents
    self.num_rows = self.generator.num_rows
    self.num_cols = self.generator.num_cols
    self.grid_shape = (self.num_rows, self.num_cols)
    self.time_limit = time_limit or (self.num_rows * self.num_cols)
    super().__init__()
    self.penalty_per_timestep = penalty_per_timestep

    # Create viewer used for rendering
    self._viewer = viewer or CleanerViewer("Cleaner", render_mode="human")
```

--------------------------------

### Define Agent Positions

Source: https://instadeepai.github.io/jumanji/environments/mmst

Example array representing the current node positions of the agents in the environment.

```python
jnp.array([1, 3])
```

--------------------------------

### Sokoban Environment Initialization

Source: https://instadeepai.github.io/jumanji/api/environments/sokoban

Instantiates a new Sokoban environment with a specified generator, time limit, and viewer.

```APIDOC
## Constructor: Sokoban

### Description
Instantiates a `Sokoban` environment with a specific generator, time limit, and viewer.

### Parameters
- **generator** (Optional[Generator]) - Optional - Generator whose `__call__` instantiates an environment instance. Defaults to `HuggingFaceDeepMindGenerator`.
- **time_limit** (int) - Optional - Max steps for the environment. Defaults to 120.
- **viewer** (Optional[Viewer]) - Optional - 'Viewer' object used to render the environment.
```

--------------------------------

### GET observation_spec

Source: https://instadeepai.github.io/jumanji/api/environments/pac_man

Retrieves the observation specifications for the PacMan environment, detailing the structure of the game state.

```APIDOC
## GET observation_spec

### Description
Returns the specifications for the observation fields of the PacMan environment.

### Method
GET

### Response
#### Success Response (200)
- **player_locations** (BoundedArray) - Tree of int32 player positions.
- **grid** (BoundedArray) - Integer representation of the maze with walls.
- **ghost_locations** (jax array) - Integer array of ghost positions.
- **power_up_locations** (jax array) - Integer array of power-pellet locations.
- **pellet_locations** (jax array) - Integer array of pellet locations.
- **action_mask** (jax array) - Boolean array defining current valid actions.
- **frightened_state_time** (int) - Time remaining in scatter mode.
- **score** (int) - Total score obtained by the player.
```

--------------------------------

### restart function

Source: https://instadeepai.github.io/jumanji/api/types

Creates a TimeStep object with step_type set to StepType.FIRST, indicating the start of an episode.

```APIDOC
## restart(observation, extras=None, shape=(), dtype=float)

### Description
Returns a TimeStep with step_type set to StepType.FIRST.

### Parameters
- **observation** (Observation) - Required - Array or tree of arrays.
- **extras** (Optional[Dict]) - Optional - Environment metric(s) or information returned by the environment but not observed by the agent. Defaults to None.
- **shape** (Union[int, Sequence[int]]) - Optional - Parameter to specify the shape of the rewards and discounts. Defaults to ().
- **dtype** (Union[dtype, type]) - Optional - Parameter to specify the data type of the rewards and discounts. Defaults to float.

### Response
- **TimeStep** - TimeStep identified as a reset.
```

--------------------------------

### GraphColoring Environment Initialization

Source: https://instadeepai.github.io/jumanji/api/environments/graph_coloring

Instantiate the GraphColoring environment with optional generator and viewer configurations.

```APIDOC
## Constructor: GraphColoring

### Description
Instantiate a `GraphColoring` environment for combinatorial optimization.

### Parameters
- **generator** (Optional[Generator]) - Optional - Callable to instantiate environment instances. Defaults to `RandomGenerator` (20 nodes, 0.8 edge probability).
- **viewer** (Optional[Viewer[State]]) - Optional - Environment viewer for rendering. Defaults to `GraphColoringViewer`.
```

--------------------------------

### Graph Adjacency Matrix Observation

Source: https://instadeepai.github.io/jumanji/environments/graph_coloring

Example of a 4x4 adjacency matrix represented as a JAX boolean array.

```python
[[False,  True, False,  True],
[ True, False,  True, False],
[False,  True, False,  True],
[ True, False,  True, False]]
```

--------------------------------

### GET action_spec

Source: https://instadeepai.github.io/jumanji/api/environments/pac_man

Retrieves the action specification for the PacMan environment, defining the 5 possible discrete actions.

```APIDOC
## GET action_spec

### Description
Returns the action specification for the environment. The action space consists of 5 discrete actions: [0, 1, 2, 3, 4] corresponding to [Up, Right, Down, Left, No-op].

### Method
GET

### Response
#### Success Response (200)
- **action_spec** (DiscreteArray) - A specs.DiscreteArray object defining the available actions.
```

--------------------------------

### Cleaner Environment Initialization

Source: https://instadeepai.github.io/jumanji/api/environments/cleaner

Instantiates and resets the Cleaner environment, then performs a step.

```APIDOC
## Cleaner Environment Initialization

### Description
This snippet demonstrates how to instantiate the Cleaner environment, reset it to an initial state, and perform a step with a sample action.

### Method
N/A (This is a usage example)

### Endpoint
N/A

### Parameters
#### Path Parameters
N/A

#### Query Parameters
N/A

#### Request Body
N/A

### Request Example
```python
from jumanji.environments import Cleaner
import jax

env = Cleaner()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

### Response
#### Success Response (200)
N/A (This is a usage example, not an API response)

#### Response Example
N/A
```

--------------------------------

### GET action_spec

Source: https://instadeepai.github.io/jumanji/api/environments/sudoku

Returns the action specification for the Sudoku environment, defining the structure of actions used to interact with the board.

```APIDOC
## GET action_spec

### Description
Returns the action spec. An action is composed of 3 integers: the row index, the column index and the value to be placed in the cell.

### Response
#### Success Response (200)
- **action_spec** (MultiDiscreteArray) - MultiDiscreteArray object.
```

--------------------------------

### TSP Environment Instantiation and Usage

Source: https://instadeepai.github.io/jumanji/api/environments/tsp

This snippet shows how to instantiate the TSP environment, reset it, render the initial state, take a step, and render the resulting state.

```APIDOC
## TSP Environment Instantiation and Usage

### Description
This code snippet demonstrates the basic usage of the TSP environment, including resetting the environment, rendering the state, taking an action, and observing the next state.

### Method
N/A (This is a usage example)

### Endpoint
N/A

### Request Example
```python
from jumanji.environments import TSP
env = TSP()
key = jax.random.PRNGKey(0)
state, timestep = jax.jit(env.reset)(key)
env.render(state)
action = env.action_spec.generate_value()
state, timestep = jax.jit(env.step)(state, action)
env.render(state)
```

### Response
N/A (This is a usage example)
```

--------------------------------

### Environment Reset API

Source: https://instadeepai.github.io/jumanji/api/environments/snake

Resets the environment to its initial state. This method is crucial for starting new episodes or simulations.

```APIDOC
## reset(key)

### Description
Resets the environment.

### Method
POST

### Endpoint
/websites/instadeepai_github_io_jumanji/reset

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
- **key** (PRNGKey) - Required - random key used to sample the snake and fruit positions.

### Request Example
```json
{
  "key": "some_random_key_string"
}
```

### Response
#### Success Response (200)
- **state** (State) - `State` object corresponding to the new state of the environment.
- **timestep** (TimeStep[Observation]) - `TimeStep` object corresponding to the first timestep returned by the environment.

#### Response Example
```json
{
  "state": {
    "board_shape": [10, 10],
    "snake_head": [3, 7],
    "fruit_position": [5, 2],
    "snake_body": [[3, 7], [3, 6]],
    "length": 2,
    "step_count": 0
  },
  "timestep": {
    "observation": {
      "board": [[0, 0, ...], ...],
      "snake_head": [3, 7],
      "fruit_position": [5, 2]
    },
    "reward": 0.0,
    "discount": 1.0,
    "step_type": "FIRST"
  }
}
```
```

--------------------------------

### Example Searcher Observation Array

Source: https://instadeepai.github.io/jumanji/environments/search_and_rescue

Represents a segmented view of distances along rays from an agent to the nearest neighbor or target.

```python
[-1.0, -1.0, 0.5, -1.0, -1.0]
```

--------------------------------

### Knapsack Environment Initialization

Source: https://instadeepai.github.io/jumanji/api/environments/knapsack

Instantiates the Knapsack environment with optional custom generators, reward functions, and viewers.

```APIDOC
## POST /environments/knapsack/init

### Description
Instantiates a new Knapsack environment instance.

### Parameters
#### Request Body
- **generator** (Optional[Generator]) - Optional - Generator whose __call__ instantiates an environment instance. Defaults to RandomGenerator (50 items, 12.5 budget).
- **reward_fn** (Optional[RewardFn]) - Optional - Reward function (DenseReward or SparseReward). Defaults to DenseReward.
- **viewer** (Optional[Viewer[State]]) - Optional - Viewer used for rendering. Defaults to KnapsackViewer.
```

--------------------------------

### Action Mask Example

Source: https://instadeepai.github.io/jumanji/environments/graph_coloring

Boolean array indicating valid color choices for a node given 4 available colors.

```python
[True, False, True, False]
```

--------------------------------

### Wrapper Class Initialization

Source: https://instadeepai.github.io/jumanji/api/wrappers

Initializes the Wrapper class, taking an existing environment as input and setting up the necessary superclass initializations.

```APIDOC
## Wrapper(env)

### Description
Wraps the environment to allow modular transformations.

### Method
__init__

### Parameters
#### Request Body
- **env** (Environment[State, ActionSpec, Observation]) - Required - The environment to wrap.

### Source Code
```python
def __init__(self, env: Environment[State, ActionSpec, Observation]):
    self._env = env
    super().__init__()
```
```

--------------------------------

### Connector Environment Initialization

Source: https://instadeepai.github.io/jumanji/api/environments/connector

Shows the constructor for the Connector environment, detailing parameters like generator, reward function, time limit, and viewer. Defaults are provided for convenience.

```python
def __init__(
    self,
    generator: Optional[Generator] = None,
    reward_fn: Optional[RewardFn] = None,
    time_limit: int = 50,
    viewer: Optional[Viewer[State]] = None,
) -> None:
    """Create the `Connector` environment.

    Args:
        generator: `Generator` whose `__call__` instantiates an environment instance.
            Implemented options are [`UniformRandomGenerator`, `RandomWalkGenerator`].
            Defaults to `RandomWalkGenerator` with `grid_size=10` and `num_agents=10`.
        reward_fn: class of type `RewardFn`, whose `__call__` is used as a reward function.
            Implemented options are [`DenseRewardFn`]. Defaults to `DenseRewardFn`.
        time_limit: the number of steps allowed before an episode terminates. Defaults to 50.
        viewer: `Viewer` used for rendering. Defaults to `ConnectorViewer` with "human" render
            mode.
    """
    self._generator = generator or RandomWalkGenerator(grid_size=10, num_agents=10)
    self._rewarder = reward_fn or DenseRewardFn()
    self.time_limit = time_limit
    self.num_agents = self._generator.num_agents
    self.grid_size = self._generator.grid_size
    super().__init__()
    self._agent_ids = jnp.arange(self.num_agents)
    self._viewer = viewer or ConnectorViewer("Connector", self.num_agents, render_mode="human")
```

--------------------------------

### Sokoban Environment Representation

Source: https://instadeepai.github.io/jumanji/api/environments/sokoban

Provides a method to get a printable string representation of the Sokoban environment's current state.

```APIDOC
## __repr__()

### Description
Returns a printable representation of the Sokoban environment.

### Method
GET (Implicit)

### Endpoint
/

### Returns
- **str** (str) - A string representation of the Sokoban environment.
```

--------------------------------

### GraphColoring Environment Initialization

Source: https://instadeepai.github.io/jumanji/api/environments/graph_coloring

Shows the constructor for the GraphColoring environment, including optional generator and viewer arguments. Defaults are provided for graph generation and rendering.

```python
def __init__(
    self,
    generator: Optional[Generator] = None,
    viewer: Optional[Viewer[State]] = None,
):
    """Instantiate a `GraphColoring` environment.

    Args:
        generator: callable to instantiate environment instances.
            Defaults to `RandomGenerator` which generates graphs with
            20 `num_nodes` and `edge_probability` equal to 0.8.
        viewer: environment viewer for rendering.
            Defaults to `GraphColoringViewer`.
    """
    self.generator = generator or RandomGenerator(num_nodes=20, edge_probability=0.8)
    self.num_nodes = self.generator.num_nodes
    super().__init__()

    # Create viewer used for rendering
    self._env_viewer = viewer or GraphColoringViewer(name="GraphColoring")
```

--------------------------------

### Maze Environment Initialization

Source: https://instadeepai.github.io/jumanji/api/environments/maze

Instantiates the Maze environment with optional generator, time limit, and viewer configurations.

```APIDOC
## Maze Environment Initialization

### Description
Instantiates a new Maze environment instance.

### Parameters
#### Request Body
- **generator** (Optional[Generator]) - Optional - Generator whose __call__ instantiates an environment instance. Defaults to RandomGenerator(num_rows=10, num_cols=10).
- **time_limit** (Optional[int]) - Optional - Maximum number of environment steps before the episode terminates. Defaults to num_rows * num_cols.
- **viewer** (Optional[Viewer[State]]) - Optional - Viewer used for rendering. Defaults to MazeEnvViewer.
```

--------------------------------

### Reset Environment Sequence

Source: https://instadeepai.github.io/jumanji/api/wrappers

Starts a new sequence in the environment and returns the initial TimeStep. This is used to re-initialize the environment for a new episode.

```python
def reset(self) -> dm_env.TimeStep:
    """Starts a new sequence and returns the first `TimeStep` of this sequence.

    Returns:
        A `TimeStep` namedtuple containing:
            - step_type: A `StepType` of `FIRST`.
            - reward: `None`, indicating the reward is undefined.
            - discount: `None`, indicating the discount is undefined.
            - observation: A NumPy array, or a nested dict, list or tuple of arrays.
                Scalar values that can be cast to NumPy arrays (e.g. Python floats)
                are also valid in place of a scalar array. Must conform to the
                specification returned by `observation_spec`.
    """
    reset_key, self._key = jax.random.split(self._key)
    self._state, timestep = self._jitted_reset(reset_key)
    return dm_env.restart(observation=timestep.observation)
```

--------------------------------

### Cleaner Environment Constructor

Source: https://instadeepai.github.io/jumanji/api/environments/cleaner

Details the parameters for initializing the Cleaner environment.

```APIDOC
## Cleaner Environment Constructor

### Description
Initializes the Cleaner environment with customizable parameters for generation, time limit, timestep penalty, and viewer.

### Method
N/A (This is a constructor)

### Endpoint
N/A

### Parameters
#### Path Parameters
N/A

#### Query Parameters
N/A

#### Request Body
N/A

### Request Example
```python
# Example instantiation with custom parameters (if needed)
# from jumanji.environments.routing.cleaner.generator import RandomGenerator
# from jumanji.viewer import CleanerViewer
# generator = RandomGenerator(num_rows=5, num_cols=5, num_agents=2)
# viewer = CleanerViewer('CustomCleaner', render_mode='rgb_array')
# env = Cleaner(generator=generator, time_limit=50, penalty_per_timestep=0.1, viewer=viewer)
```

### Response
#### Success Response (200)
N/A

#### Response Example
N/A

### Parameters
#### Path Parameters
N/A

#### Query Parameters
N/A

#### Request Body
N/A

### Parameters
Name | Type | Description | Default
---|---|---|---
`generator` | `Optional[Generator]` | `Generator` whose `__call__` instantiates an environment instance. Implemented options are [`RandomGenerator`]. Defaults to `RandomGenerator` with `num_rows=10`, `num_cols=10` and `num_agents=3`. | `None`
`time_limit` | `Optional[int]` | max number of steps in an episode. Defaults to `num_rows * num_cols`. | `None`
`penalty_per_timestep` | `float` | the penalty returned at each timestep in the reward. | `0.5`
`viewer` | `Optional[Viewer[State]]` | `Viewer` used for rendering. Defaults to `CleanerViewer` with "human" render mode. | `None`
```

--------------------------------

### Define MMST Node Types Observation

Source: https://instadeepai.github.io/jumanji/environments/mmst

Example representation of node types for a 12-node graph with 2 agents, formatted as a JAX array.

```python
node_types = jnp.array( [ [1, 0, -1, 2, -1, 3, 1, -1, 3, 1, -1, -1], [3, 2, -1, 0, -1, 1, 3, -1, 1, 3, -1, -1], ], dtype=jnp.int32, )
```
Tokens:
Raw