### Quickstart Flat Buffer

Source: https://github.com/instadeepai/flashbax/blob/main/README.md

Minimal example demonstrating the usage of pure functions for a flat buffer, compatible with JAX transformations.

```python
import jax
import jax.numpy as jnp
import flashbax as fbx
```

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

### Install Flashbax and Dependencies

Source: https://github.com/instadeepai/flashbax/blob/main/examples/gym_dqn_example.ipynb

Clone the Flashbax repository and install it with example dependencies.

```bash
!git clone https://github.com/instadeepai/flashbax.git
```

```bash
!pip install ./flashbax[examples]
```

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

### Install Flashbax, Matrax, and Jumanji

Source: https://github.com/instadeepai/flashbax/blob/main/examples/matrax_iql_example.ipynb

Installs the necessary libraries for the example. This command should be run in a terminal or notebook environment.

```bash
%%capture
!pip install flashbax matrax jumanji
```

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

### Experiment Setup Function

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_dqn_example.ipynb

Initializes the experiment environment, network, optimizer, parameters, and buffer. Sets up the epsilon schedule for exploration.

```python
def set_up_experiment(
    env,
    batch_size,
    step_size,
    seed,
    buffer_size,
    epsilon_init,
    epsilon_final,
    epsilon_steps,
):
    """Sets up the experiment."""
    cores_count = len(jax.devices())  # get available TPU cores.
    network = get_network_fn(env.action_spec.num_values)  # define network.
    optim = optax.adam(step_size)  # define optimiser.

    rng, rng_e, rng_p = random.split(random.PRNGKey(seed), num=3)  # prng keys.
    _, timestep = env.reset(rng_e)
    obs = timestep.observation.grid
    dummy_obs = jnp.expand_dims(obs, 0)  # dummy for net init.
    params = network.init(rng_p, dummy_obs)  # initialise params.
    opt_state = optim.init(params)  # initialise optimiser stats.
    buffer_fn = fbx.make_flat_buffer(
        max_length=buffer_size,
        min_length=batch_size,
        sample_batch_size=batch_size,
        add_sequences=True,
        add_batch_size=None,
    )
    buffer_state = buffer_fn.init(
        TimeStep(
            observation=obs,
            action=jnp.zeros((), dtype=jnp.int32),
            reward=jnp.zeros(()),
            discount=jnp.zeros(()),
        )
    )  # initialise buffer state.
    epsilon_schedule_fn = optax.linear_schedule(
        epsilon_init, epsilon_final, epsilon_steps
    )  # define epsilon schedule.
    return (
        cores_count,
        network,
        optim,
        params,
        opt_state,
        buffer_fn,
        buffer_state,
        rng,
        epsilon_schedule_fn,
    )
```

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

### Install Flashbax

Source: https://github.com/instadeepai/flashbax/blob/main/README.md

Install the library using pip.

```bash
pip install flashbax
```

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

### Initialize Q-Network and Parameters

Source: https://github.com/instadeepai/flashbax/blob/main/examples/matrax_iql_example.ipynb

Initializes the Q-network and its parameters using a dummy observation and a JAX random key. This setup is necessary before training or inference.

```python
SEED = 42

# Instantiate Q-network
q_network = QNetwork(NUM_ACTIONS)

# Create a single dummy observation (i.e., batch size is 1)
# We add num_agents to num_obs to account for the addition of one-hot-encoded agent IDs
dummy_obs = jnp.zeros((1, NUM_OBS + NUM_AGENTS), jnp.float32)

# Generate random key for initialising params
key = jax.random.PRNGKey(SEED)
key, subkey = jax.random.split(key)

# Initialise Q-network
q_network_params = q_network.init(subkey, dummy_obs)
```

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

### Experiment Setup Function

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_prioritised_dqn_example.ipynb

This function initializes the experiment by setting up the environment, network, optimizer, PRNG keys, and buffer. It prepares all necessary components for training.

```python
def set_up_experiment(
    env,
    batch_size,
    step_size,
    seed,
    buffer_size,
    epsilon_init,
    epsilon_final,
    epsilon_steps,
    priority_exponent,
):
    """Sets up the experiment."""
    cores_count = len(jax.devices())  # get available TPU cores.
    network = get_network_fn(env.action_spec.num_values)  # define network.
    optim = optax.adam(step_size)  # define optimiser.

    rng, rng_e, rng_p = random.split(random.PRNGKey(seed), num=3)  # prng keys.
    _, timestep = env.reset(rng_e)
    obs = timestep.observation.grid
    dummy_obs = jnp.expand_dims(obs, 0)  # dummy for net init.
    params = network.init(rng_p, dummy_obs)  # initialise params.
    opt_state = optim.init(params)  # initialise optimiser stats.
    buffer_fn = fbx.make_prioritised_flat_buffer(
        max_length=buffer_size,
        min_length=batch_size,
        sample_batch_size=batch_size,
        add_sequences=True,
        add_batch_size=None,
        priority_exponent=priority_exponent,
    )
    buffer_state = buffer_fn.init(
        TimeStep(
            observation=obs,
            action=jnp.zeros((), dtype=jnp.int32),
            reward=jnp.zeros(()),
            discount=jnp.zeros(()),
        )
    )  # initialise buffer state.
    epsilon_schedule_fn = optax.linear_schedule(
        epsilon_init, epsilon_final, epsilon_steps
    )  # define epsilon schedule.
    return (
        cores_count,
        network,
        optim,
        params,
        opt_state,
        buffer_fn,
        buffer_state,
        rng,
        epsilon_schedule_fn,
    )
```

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

### Setup Fake Devices for JAX

Source: https://github.com/instadeepai/flashbax/blob/main/examples/quickstart_trajectory_buffer.ipynb

Sets up mock CPU devices for JAX, often used with `jax.pmap` for distributed computation.

```python
import chex
import jax
import jax.numpy as jnp

# Setup fake devices - we use this later with `jax.pmap`.
DEVICE_COUNT_MOCK = 2
chex.set_n_cpu_devices(DEVICE_COUNT_MOCK)
```

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

### Initialize Replay Buffer State

Source: https://github.com/instadeepai/flashbax/blob/main/examples/matrax_iql_example.ipynb

Initializes the replay buffer state with a dummy transition. This is a setup step before interacting with the buffer.

```python
transition = {
    "obs": jnp.zeros(
        (NUM_AGENTS, NUM_AGENTS), dtype="float32"  # second dim won't always be num_agents
    ),
    "action": jnp.zeros(NUM_AGENTS, dtype="int32"),
    "reward": jnp.zeros(NUM_AGENTS, dtype="float32"),
    "next_obs": jnp.zeros((NUM_AGENTS, NUM_AGENTS), dtype="float32"),
    "done": 0.0,
}  # store in dictionary

q_learning_memory_state = q_learning_memory.init(transition)
```

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

### Learner Function Setup

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_dqn_example.ipynb

Sets up the main learner function, which orchestrates environment interaction, data collection, and parameter updates for the DQN agent.

```python
def get_learner_fn(
    env,
    forward_pass,
    buffer_fn,
    opt_update,
    rollout_len,
    agent_discount,
    iterations,
    target_period,
    epsilon_schedule_fn,
    sgd_steps_per_rollout,
):
    """Returns a learner function that can be used to update parameters."""

    def rollout_fn(params_state, outer_rng, env_state, env_timestep):
        """Collects a trajectory from the environment."""
        def step_fn(env_data, rng):
            """Steps the environment and collects transition data."""
            (env_state, env_timestep, params_state) = env_data
            obs_tm1 = env_timestep.observation.grid  
            d_tm1 = env_timestep.discount  
            q_values_tm1 = forward_pass(params_state.online, jnp.expand_dims(obs_tm1, 0))  
            a_tm1_dist = distrax.EpsilonGreedy(preferences=q_values_tm1[0], epsilon=epsilon_schedule_fn(params_state.update_count))
            a_tm1 = a_tm1_dist.sample(seed=rng) 
            new_env_state, new_env_timestep = env.step(env_state, a_tm1)  #
            r_t = new_env_timestep.reward 
            return (
                new_env_state,
                new_env_timestep,
                params_state,
            ), TimeStep(  # return env state and transition data.
                observation=obs_tm1, action=a_tm1, discount=d_tm1, reward=r_t
            ) # We line up the observation with its discount, not the discount of the next observation as is usually seen.
              # This is so that we know in a transition that discount[1] is the discount of the next observation. 
              # e.g. indexing is v[t] = reward[t] + discount[t+1]*value[t+1]
              # Switching to Sutton and Barto's notation, we would have v[t] = reward[t+1] + discount[t+1]*value[t+1]
              # To do this, we would add r_tm1 = env_timestep.reward to the TimeStep dataclass, not the new_env_timestep.reward

        step_rngs = jax.random.split(outer_rng, rollout_len)
        (env_state, env_timestep, params_state), rollout = lax.scan(
            step_fn, (env_state, env_timestep, params_state), step_rngs
        )  # trajectory.

        return rollout, env_state, env_timestep

    def loss_fn(params, target_params, batch):
        """Computes the loss for a single batch."""
        # For ease of reading
        o_tm1 = batch.first.observation
        a_tm1 = batch.first.action.astype(jnp.int32)
        r_t = batch.first.reward
        d_t = agent_discount * batch.second.discount
        o_t = batch.second.observation
        # Compute Q-values for current and next states.
        q_tm1 = forward_pass(params, o_tm1) 
        q_t = forward_pass(target_params, o_t)  
        q_t_select = forward_pass(params, o_t)
        # Compute the TD-error.
        td_error = jax.vmap(
            rlax.double_q_learning
        )(  # compute multi-step temporal diff error.
            q_tm1=q_tm1,  # predictions.
            a_tm1=a_tm1,  # actions.
            r_t=r_t,  # rewards.
            discount_t=d_t,  # discount.
            q_t_value=q_t,  # target values.
            q_t_selector=q_t_select,  # selector values.
        )
        return jnp.mean(jnp.square(td_error))

    def update_fn(
        params_state: Params, buffer_state, opt_state, rng, env_state, env_timestep
    ):
        """Updates the parameters of the agent."""
        rng, rollout_rng, update_rng = random.split(rng, 3)

        data_rollout, new_env_state, new_env_timestep = rollout_fn(
            params_state, rollout_rng, env_state, env_timestep
        )  # collect trajectory from environment. This could be one step, or many steps.

        buffer_state = buffer_fn.add(buffer_state, data_rollout)  # store trajectory in buffer. 
        
        def sgd_step(carry, rng):
            """Performs a single SGD step."""
            params_state, opt_state, buffer_state = carry
            
            batch = buffer_fn.sample(buffer_state, rng).experience  # sample batch from buffer.

```

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

### Install Flashbax if not present

Source: https://github.com/instadeepai/flashbax/blob/main/examples/vault_demonstration.ipynb

Ensures Flashbax is installed. This snippet should be run first.

```python
%%capture
try:
    import flashbax as fbx
except ModuleNotFoundError:
    print('installing flashbax')
    %pip install -q flashbax
    import flashbax as fbx
```

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

### Environment Setup Function

Source: https://github.com/instadeepai/flashbax/blob/main/examples/gym_dqn_example.ipynb

Defines a factory function to create gym environments with AutoResetWrapper and RecordEpisodeStatistics. Ensures environments are seeded for reproducibility.

```python
def make_env(env_id, seed):
    def thunk():
        
        env = gym.make(env_id)
        # We use an auto reset wrapper to automatically reset the environment
        # when the episode is done since we are using vectorized environments
        # and we want all the environments to always be active.
        # Additionally, we use the auto reset wrapper since we are adding transitions 
        # sequentially and we want to maintain the order in which we are adding the transitions.
        env = gym.wrappers.AutoResetWrapper(env)
        env = gym.wrappers.RecordEpisodeStatistics(env)
        env.action_space.seed(seed)

        return env

    return thunk
```

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

### Initialize Environment and Print Details

Source: https://github.com/instadeepai/flashbax/blob/main/examples/matrax_iql_example.ipynb

Initializes a multi-agent environment and prints the number of agents and actions available. This is a common starting point for setting up experiments.

```python
eval_env = matrax.make(ENV_NAME)

NUM_ACTIONS = env.num_actions
NUM_AGENTS = env.num_agents
NUM_OBS = NUM_AGENTS  # in matrax, observations have shape (num_agents, num_agents)

print("Number of agents:", env.num_agents)
print("Number of actions:", env.num_actions)
```

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

### Setup Training and Evaluation Environments

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_dqn_example.ipynb

Creates a Jumanji environment and wraps it with AutoResetWrapper for automatic resetting during training. Evaluation environments do not automatically reset.

```python
# We separate the environment into a training and evaluation environment. Training resets the environment automatically, while evaluation does not.
env = jumanji.make("Snake-v1", num_rows=6, num_cols=6)
training_env = AutoResetWrapper(env)
```

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

### Buffer wrap-around example

Source: https://github.com/instadeepai/flashbax/blob/main/examples/vault_demonstration.ipynb

Illustrates the state of the buffer and vault after the buffer has wrapped around, showing that the vault continues to store data correctly.

```python
Buffer state:
[[[6. 6.]
  [2. 2.]
  [3. 3.]
  [4. 4.]
  [5. 5.]]]

Vault state:
[[[1. 1.]
  [2. 2.]
  [3. 3.]
  [4. 4.]
  [5. 5.]
  [6. 6.]]]
```

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

### DQN Learner Function Setup

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_prioritised_dqn_example.ipynb

Sets up the core learner function for the DQN agent. This function orchestrates data collection (rollout) and parameter updates using a defined loss function.

```python
def get_learner_fn(
    env,
    forward_pass,
    buffer_fn,
    opt_update,
    rollout_len,
    agent_discount,
    iterations,
    target_period,
    epsilon_schedule_fn,
    sgd_steps_per_rollout,
    importance_sampling_exponent,
):
    """Returns a learner function that can be used to update parameters."""

    def rollout(params_state, outer_rng, env_state, env_timestep):
        """Collects a trajectory from the environment."""
        def step_fn(env_data, rng):
            """Steps the environment and collects transition data."""
            (env_state, env_timestep, params_state) = env_data
            obs_tm1 = env_timestep.observation.grid  
            d_tm1 = env_timestep.discount  
            q_values_tm1 = forward_pass(params_state.online, jnp.expand_dims(obs_tm1, 0))  
            a_tm1_dist = distrax.EpsilonGreedy(preferences=q_values_tm1[0], epsilon=epsilon_schedule_fn(params_state.update_count))
            a_tm1 = a_tm1_dist.sample(seed=rng) 
            new_env_state, new_env_timestep = env.step(env_state, a_tm1)  #
            r_t = new_env_timestep.reward 
            return (
                new_env_state,
                new_env_timestep,
                params_state,
            ), TimeStep(  # return env state and transition data.
                observation=obs_tm1, action=a_tm1, discount=d_tm1, reward=r_t
            ) # We line up the observation with its discount, not the discount of the next observation as is usually seen.
              # This is so that we know in a transition that discount[1] is the discount of the next observation. 
              # e.g. indexing is v[t] = reward[t] + discount[t+1]*value[t+1]
              # Switching to Sutton and Barto's notation, we would have v[t] = reward[t+1] + discount[t+1]*value[t+1]
              # To do this, we would add r_tm1 = env_timestep.reward to the TimeStep dataclass, not the new_env_timestep.reward

        step_rngs = jax.random.split(outer_rng, rollout_len)
        (env_state, env_timestep, params_state), rollout = lax.scan(
            step_fn, (env_state, env_timestep, params_state), step_rngs
        )  # trajectory.

        return rollout, env_state, env_timestep

    def loss_fn(params, target_params, batch):
        """Computes the loss for a single batch."""
        # For ease of reading
        o_tm1 = batch.experience.first.observation
        a_tm1 = batch.experience.first.action.astype(jnp.int32)
        r_t = batch.experience.first.reward
        d_t = agent_discount * batch.experience.second.discount
        o_t = batch.experience.second.observation
        # Compute Q-values for current and next states.
        q_tm1 = forward_pass(params, o_tm1) 
        q_t = forward_pass(target_params, o_t)  
        q_t_select = forward_pass(params, o_t)
        # Compute the TD-error.
        td_error = jax.vmap(
            rlax.double_q_learning
        )(  # compute multi-step temporal diff error.
            q_tm1=q_tm1,  # predictions.
            a_tm1=a_tm1,  # actions.
            r_t=r_t,  # rewards.
            discount_t=d_t,  # discount.
            q_t_value=q_t,  # target values.
            q_t_selector=q_t_select,  # selector values.
        )

        batch_loss = rlax.l2_loss(td_error)  # compute L2 loss.

        # Get the importance weights.
        importance_weights = (1. / batch.probabilities).astype(jnp.float32)
        importance_weights **= importance_sampling_exponent
        importance_weights /= jnp.max(importance_weights)

        # Reweight.
        loss = jnp.mean(importance_weights * batch_loss)  # compute L2 loss.
        new_priorities = jnp.abs(td_error) + 1e-7
        return loss, new_priorities

    def update_fn(
        params_state: Params, buffer_state, opt_state, rng, env_state, env_timestep
    ):
        """Updates the parameters of the agent."""
        rng, rollout_rng, update_rng = random.split(rng, 3)

        data_rollout, new_env_state, new_env_timestep = rollout(

```

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

### Import Reinforcement Learning Libraries

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_ppo_example.ipynb

Imports a comprehensive set of libraries required for reinforcement learning, including JAX, Haiku, Optax, RLax, and Jumanji. This setup is essential for building and training RL agents.

```python
from matplotlib import pyplot as plt
import seaborn as sns
import numpy as np
import jax
from jax import lax
from jax import random
from jax import numpy as jnp
import jax.numpy as jnp
import haiku as hk
import optax
import rlax
import timeit
import distrax
import chex
from jumanji.wrappers import AutoResetWrapper
import jumanji

print("The following devices are available: ", jax.devices())
```

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

### Q-Learning Action Selection (Training/Evaluation)

Source: https://github.com/instadeepai/flashbax/blob/main/examples/matrax_iql_example.ipynb

Selects an action for Q-learning, switching between epsilon-greedy (for training) and greedy (for evaluation) strategies. It applies the Q-network to get Q-values and then uses the appropriate selection method.

```python
def q_learning_select_action(key, params, num_timesteps, obs, evaluation=False):
    """A function to perform greedy or epsilon-greedy action selection,
    based on whether we are evaluating or training a policy.
    """

    obs = jnp.expand_dims(obs, axis=0)  # add a batch dim as the leading axis

    q_values = q_network.apply(params.online, obs)[0]  # remove batch dim

    action = select_epsilon_greedy_action(key, q_values, num_timesteps)
    greedy_action = select_greedy_action(q_values)

    action = jax.lax.select(evaluation, greedy_action, action)

    return action
```

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

### Inspect Sampled Batch Keys and Shape (Repeat)

Source: https://github.com/instadeepai/flashbax/blob/main/examples/quickstart_trajectory_buffer.ipynb

Prints the keys of the experience dictionary within the sampled batch and the shape of the 'reward' tensor. This is a repeat of a previous example for emphasis.

```python
print(batch.experience.keys()) # prints dict_keys(['obs', 'reward'])
# prints (4, 8) = (sample_batch_size, sample_sequence_length, *fake_timestep['reward'].shape)
print(batch.experience['reward'].shape)
```

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

### Broadcast Parameters to Device Shape in JAX

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_prioritised_dqn_example.ipynb

Broadcasts JAX parameters, optimizer state, and buffer state to match the device shape (cores_count, num_envs). This is essential for distributed training setups. It also initializes the update count to zeros.

```python
def broadcast_to_device_shape(
    cores_count, num_envs, params, opt_state, buffer_state, rng
):
    """Broadcasts parameters to device shape."""
    broadcast = lambda x: jnp.broadcast_to(x, (cores_count, num_envs) + x.shape)
    params = jax.tree.map(broadcast, params)  # broadcast to cores and batch.
    opt_state = jax.tree.map(broadcast, opt_state)  # broadcast to cores and batch
    buffer_state = jax.tree.map(broadcast, buffer_state)  # broadcast to cores and batch
    params_state = Params(
        online=params,
        target=params,
        update_count=jnp.zeros(shape=(cores_count, num_envs)),
    )
    rng, step_rngs = get_rng_keys(cores_count, num_envs, rng)
    return params_state, opt_state, buffer_state, step_rngs, rng
```

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

### Initialize and Use Buffers

Source: https://github.com/instadeepai/flashbax/blob/main/README.md

Demonstrates how to instantiate various buffer types and perform basic operations like initialization, adding data, and sampling.

```python
import flashbax as fbx

buffer = fbx.make_trajectory_buffer(...)
# OR
buffer = fbx.make_prioritised_trajectory_buffer(...)
# OR
buffer = fbx.make_flat_buffer(...)
# OR
buffer = fbx.make_prioritised_flat_buffer(...)
# OR
buffer = fbx.make_item_buffer(...)
# OR
buffer = fbx.make_trajectory_queue(...)

# Initialise
state = buffer.init(example_timestep)
# Add Data
state = buffer.add(state, example_data)
# Sample Data
batch = buffer.sample(state, rng_key)
```

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

### Set Up Experiment Parameters

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_dqn_example.ipynb

Initializes experiment parameters including environment, network, optimizer, and buffer configurations.

```python
(cores_count, network, optim, params, opt_state, buffer_fn, buffer_state, rng, epsilon_schedule_fn) = set_up_experiment(
    env=training_env,
    batch_size=BATCH_SIZE,
    step_size=LEARNING_RATE,
    seed=SEED,
    buffer_size=BUFFER_SIZE,
    epsilon_steps=EPSILON_STEPS,
    epsilon_init=EPSILON_INIT,
    epsilon_final=EPSILON_FINAL,
)
```

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

### JAX and Network Initialization

Source: https://github.com/instadeepai/flashbax/blob/main/examples/gym_dqn_example.ipynb

Sets up JAX devices, random seeds, the Q-network, and the Adam optimizer. Initializes network parameters and the training state.

```python
with jax.default_device(jax.devices('cpu')[0]):

    # Specify the random seeds we will use for reproducibility
    random.seed(seed)
    np.random.seed(seed)
    key = jax.random.PRNGKey(seed)
    key, q_key = jax.random.split(key, 2)

    # Set up the network and optimiser
    q_network = get_network_fn(num_actions)
    optim = optax.adam(learning_rate=learning_rate)

    # Get an initial observation from the environment to initialize the network 
    dummy_obs, _ = envs.reset(seed=seed)
    if num_envs>1:
        dummy_obs = dummy_obs[0]
    # Initialize the network parameters
    params = q_network.init(q_key, dummy_obs, None)
    # Initialize the optimiser state
    opt_state=optim.init(params)
    # Initialize the initial train state
    q_state = TrainState(
        params=params,
        target_params=params,
        opt_state=opt_state
    )
```

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

### Initialize Experiment and Device State

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_ppo_example.ipynb

Utilities for setting up the environment, network, optimizer, and broadcasting states across JAX devices.

```python
def set_up_experiment(env, rollout_len, step_size, seed, buffer_size):
    """Sets up the experiment and returns the necessary information."""

    cores_count = len(jax.devices())  # get available TPU cores.
    network = get_network_fn(env.action_spec.num_values)  # define network.
    optim = optax.adam(step_size)  # define optimiser.

    rng, rng_e, rng_p = random.split(random.PRNGKey(seed), num=3)  # prng keys.
    _, timestep = env.reset(rng_e)
    obs = timestep.observation.grid 
    dummy_obs = jnp.expand_dims(obs , 0) # dummy for net init.
    params = network.init(rng_p, dummy_obs)  # initialise params.
    opt_state = optim.init(params)  # initialise optimiser stats.
    buffer_fn = fbx.make_trajectory_queue(
        max_length_time_axis=buffer_size,
        add_batch_size=1,
        add_sequence_length=rollout_len,
        sample_sequence_length=rollout_len
    )
    dummy_transition = TimeStep(observation=obs,
                                action=jnp.zeros((), dtype=jnp.int32),
                                reward=timestep.reward,
                                discount=timestep.discount,
                                behaviour_action_log_prob=jnp.zeros(()),
                                behaviour_value=jnp.zeros(()),
                                )
    buffer_state = buffer_fn.init(dummy_transition) # initialise buffer.
    return cores_count, network, optim, params, opt_state, buffer_fn, buffer_state, rng


def get_rng_keys(cores_count, num_envs, rng):
    """Returns a list of rng keys for each environment."""
    rng, *step_rngs = jax.random.split(rng, cores_count * num_envs + 1)
    reshape = lambda x: x.reshape((cores_count, num_envs) + x.shape[1:])
    step_rngs = reshape(jnp.stack(step_rngs))  # add dimension to pmap over.
    return rng, step_rngs

def broadcast_to_device_shape(cores_count, num_envs, params, opt_state, buffer_state, rng):
    """Broadcasts the parameters to the shape of the device."""
    broadcast = lambda x: jnp.broadcast_to(x, (cores_count, num_envs) + x.shape)
    params = jax.tree.map(broadcast, params)  # broadcast to cores and batch.
    opt_state = jax.tree.map(broadcast, opt_state)  # broadcast to cores and batch
    buffer_state = jax.tree.map(broadcast, buffer_state)  # broadcast to cores and batch
    params_state = Params(online=params, update_count=jnp.zeros(shape=(cores_count, num_envs)))
    rng, step_rngs = get_rng_keys(cores_count, num_envs, rng)
    return params_state, opt_state, buffer_state, step_rngs, rng
```

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

### Initialize Environment and Parameters

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_dqn_example.ipynb

Initializes environment parameters, RNG keys, and returns the initial states for parameters, optimizer, buffer, and RNGs.

```python
params_state = Params(
        online=params,
        target=params,
        update_count=jnp.zeros(shape=(cores_count, num_envs)),
    )
    rng, step_rngs = get_rng_keys(cores_count, num_envs, rng)
    return params_state, opt_state, buffer_state, step_rngs, rng
```

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

### Set up Adam Optimizer for Q-Learning

Source: https://github.com/instadeepai/flashbax/blob/main/examples/matrax_iql_example.ipynb

Initializes the Adam optimizer with a specified learning rate and creates the initial optimizer state for the Q-network parameters.

```python
LEARNING_RATE = 3e-4

# Initialise Q-network optimiser
OPTIMISER = optax.adam(learning_rate=LEARNING_RATE)

Q_LEARN_OPTIM_STATE = OPTIMISER.init(q_network_params)  # initial optim state

# Create Learn State
q_learning_learn_state = QLearnState(
    0, Q_LEARN_OPTIM_STATE
)  # count set to zero initially
```

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

### Get Network Function

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_dqn_example.ipynb

Creates and returns a Haiku network function for processing observations and outputting action logits. Includes convolutional and linear layers.

```python
def get_network_fn(num_outputs: int):
    def network_fn(obs: chex.Array) -> chex.Array:
        """Outputs action logits."""
        network = hk.Sequential(
            [
            hk.Conv2D(32, kernel_shape=2, stride=1),
            jax.nn.relu,
            hk.Conv2D(32, kernel_shape=2, stride=1),
            jax.nn.relu,
            hk.Flatten(),
            hk.Linear(256),
            jax.nn.relu,
            hk.Linear(128),
            jax.nn.relu,
            hk.Linear(num_outputs),
            ]
        )
        return network(obs)

    return hk.without_apply_rng(hk.transform(network_fn))
```

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

### Get Learner Function and Pmap

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_dqn_example.ipynb

Retrieves the learner function, which encapsulates the training logic, and then uses `jax.pmap` to replicate it across multiple cores for parallel training.

```python
learn = get_learner_fn(
    env=training_env,
    forward_pass=network.apply,
    buffer_fn=buffer_fn,
    opt_update=optim.update,
    rollout_len=ROLLOUT_LEN,
    agent_discount=AGENT_DISCOUNT,
    iterations=TRAINING_ITERS,
    target_period=TARGET_PERIOD,
    epsilon_schedule_fn=epsilon_schedule_fn,
    sgd_steps_per_rollout=SGD_STEPS_PER_ROLLOUT,
)
learn = jax.pmap(learn, axis_name="i")  # replicate over multiple cores.
```

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

### Initialize Experiment and Evaluation Functions

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_prioritised_dqn_example.ipynb

Sets up the experiment environment and defines JAX-jitted functions for evaluating agent performance over episodes.

```python
(
    cores_count,
    network,
    optim,
    params,
    opt_state,
    buffer_fn,
    buffer_state,
    rng,
    epsilon_schedule_fn,
) = set_up_experiment(
    env=training_env,
    batch_size=BATCH_SIZE,
    step_size=LEARNING_RATE,
    seed=SEED,
    buffer_size=BUFFER_SIZE,
    epsilon_steps=EPSILON_STEPS,
    epsilon_init=EPSILON_INIT,
    epsilon_final=EPSILON_FINAL,
    priority_exponent=PRIORTIY_EXPONENT,
)

@jax.jit
def eval_one_episode(params, rng):
    """Evaluates one episode."""
    state, timestep = env.reset(rng)
    
    def step(val):
        params, state, timestep, tot_r, rng, done = val
        rng, key_step = jax.random.split(rng)
        obs = timestep.observation.grid
        q_values = network.apply(params, obs[jnp.newaxis,])
        a_t = jnp.argmax(q_values, axis=-1)[0]
        state, timestep = env.step(state, a_t)
        tot_r += timestep.reward
        return (params, state, timestep, tot_r, rng, timestep.last())
       
    params, state, timestep, tot_r, rng, done = jax.lax.while_loop(lambda val : val[5] == False, step, (params, state, timestep, 0, rng, False))

    return params, tot_r

@jax.jit
def eval(params, rng):
    """Evaluates multiple episodes."""
    rngs = random.split(rng, NUM_EVAL_EPISODES)
    params = jax.tree.map(lambda x: x[0][0], params)
    _, tot_r = jax.lax.scan(eval_one_episode, params, rngs)
    return tot_r.mean()
```

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

### Get RNG Keys Function

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_prioritised_dqn_example.ipynb

This function is designed to return a batch of random number generator keys, likely for use across multiple environments or parallel processes.

```python
def get_rng_keys(cores_count, num_envs, rng):
    """Returns a batch of random number generator keys."""
```

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

### Create and Use a Flat Buffer

Source: https://context7.com/instadeepai/flashbax/llms.txt

Initialize a flat buffer for sampling individual transitions and demonstrate adding data and sampling batches.

```python
import jax
import jax.numpy as jnp
import flashbax as fbx

# Create a flat buffer with max 10000 transitions, min 1000 before sampling
buffer = fbx.make_flat_buffer(
    max_length=10000,        # Maximum number of timesteps
    min_length=1000,         # Minimum timesteps before sampling allowed
    sample_batch_size=32,    # Batch size when sampling
    add_sequences=False,     # Adding single timesteps (not sequences)
    add_batch_size=None,     # Adding one timestep at a time
)

# Define experience structure with an example timestep
fake_timestep = {
    "obs": jnp.zeros((84, 84, 4), dtype=jnp.float32),
    "action": jnp.array(0, dtype=jnp.int32),
    "reward": jnp.array(0.0, dtype=jnp.float32),
    "done": jnp.array(False, dtype=jnp.bool_),
}

# Initialize buffer state
state = buffer.init(fake_timestep)

# Add timesteps to the buffer
for i in range(1500):
    timestep = {
        "obs": jnp.ones((84, 84, 4)) * i,
        "action": jnp.array(i % 4),
        "reward": jnp.array(float(i)),
        "done": jnp.array(i % 100 == 0),
    }
    state = buffer.add(state, timestep)

# Check if we can sample
print(f"Can sample: {buffer.can_sample(state)}")  # True

# Sample a batch of transitions
rng_key = jax.random.PRNGKey(0)
batch = buffer.sample(state, rng_key)

# Access transition pairs (s_t, s_{t+1})
print(f"Current obs shape: {batch.experience.first['obs'].shape}")   # (32, 84, 84, 4)
print(f"Next obs shape: {batch.experience.second['obs'].shape}")     # (32, 84, 84, 4)
print(f"Actions shape: {batch.experience.first['action'].shape}")    # (32,)
print(f"Rewards shape: {batch.experience.first['reward'].shape}")    # (32,)

```

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

### Initialize Q-Learning Parameters

Source: https://github.com/instadeepai/flashbax/blob/main/examples/matrax_iql_example.ipynb

Initializes the QLearnParams object with online and target network weights.

```python
q_learning_params = QLearnParams(
    online=q_network_params, target=q_network_params
)  # target equal to online
```

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

### Execute Training Loop

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_ppo_example.ipynb

Initializes environment states and runs the training loop across multiple cores using jax.pmap.

```python
rng, *env_rngs = jax.random.split(rng, cores_count * NUM_ENVS + 1)
env_states, env_timesteps = jax.vmap(env.reset)(jnp.stack(env_rngs))  # init envs.
reshape = lambda x: x.reshape((cores_count, NUM_ENVS) + x.shape[1:])
env_states = jax.tree.map(reshape, env_states)  # add dimension to pmap over.
env_timesteps = jax.tree.map(reshape, env_timesteps)  # add dimension to pmap over.
params_state, opt_state, buffer_state, step_rngs, rng = broadcast_to_device_shape(cores_count, NUM_ENVS, params, opt_state, buffer_state, rng)


learn = get_learner_fn(env=training_env, forward_pass=network.apply, buffer_fn=buffer_fn, opt_update=optim.update, rollout_len=ROLLOUT_LEN, agent_discount=AGENT_DISCOUNT, iterations=TRAINING_ITERS, gae_lambda=GAE_LAMBDA, clip_epsilon=CLIP_EPSILON, sgd_epochs=SGD_EPOCHS)
learn = jax.pmap(learn, axis_name='i')  # replicate over multiple cores.

avg_reward = []
total_time = 0
for training_eval_iters in range(1, TRAINING_EVAL_ITERS+1):
    # Train
    start = timeit.default_timer()
    params_state, buffer_state, opt_state, step_rngs, env_states, env_timesteps = learn(params_state, buffer_state, opt_state, step_rngs, env_states, env_timesteps)
    params_state = jax.tree.map(lambda x: x.block_until_ready(), params_state) # wait for params to be ready so time is accurate.
    total_time += timeit.default_timer() - start
    # Eval
    rng, eval_rng = jax.random.split(rng, num=2)
    tot_r = eval(params_state.online, eval_rng)
    avg_reward.append(tot_r)
    if training_eval_iters % 2 == 0:
        print(f"Average Reward at iteration {training_eval_iters}: {tot_r}")
```

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

### Instantiate and Configure Replay Buffer

Source: https://github.com/instadeepai/flashbax/blob/main/examples/matrax_iql_example.ipynb

Instantiates a replay buffer (`fbx.make_item_buffer`) for Q-learning. It's configured with maximum and minimum lengths, sample batch size, and enabled batch addition for storing environment transitions.

```python
# Instantiate the replay buffer
BATCH_SIZE = 64
q_learning_memory = fbx.make_item_buffer(
    max_length=50_000, min_length=64, sample_batch_size=BATCH_SIZE, add_batches=True
)
```

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

### Initialize Environments and Broadcast Parameters

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_dqn_example.ipynb

Initializes multiple environments in parallel and broadcasts parameters and states to the device shape.

```python
rng, *env_rngs = jax.random.split(rng, cores_count * NUM_ENVS + 1)
env_states, env_timesteps = jax.vmap(env.reset)(jnp.stack(env_rngs))  # init envs.
reshape = lambda x: x.reshape((cores_count, NUM_ENVS) + x.shape[1:])
env_states = jax.tree.map(reshape, env_states)  # add dimension to pmap over.
env_timesteps = jax.tree.map(reshape, env_timesteps)  # add dimension to pmap over.
params_state, opt_state, buffer_state, step_rngs, rng = broadcast_to_device_shape(
    cores_count, NUM_ENVS, params, opt_state, buffer_state, rng
)
```

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

### Instantiate and use a Flashbax flat buffer

Source: https://github.com/instadeepai/flashbax/blob/main/README.md

Demonstrates the lifecycle of a flat buffer including initialization, adding data, and sampling. Note that sampling requires the buffer to meet the minimum transition length.

```python
# Instantiate the flat buffer NamedTuple using `make_flat_buffer` using a simple configuration.
# The returned `buffer` is simply a container for the pure functions needed for using a flat buffer.
buffer = fbx.make_flat_buffer(max_length=32, min_length=2, sample_batch_size=1)

# Initialise the buffer's state.
fake_timestep = {"obs": jnp.array([0, 0]), "reward": jnp.array(0.0)}
state = buffer.init(fake_timestep)

# Now we add data to the buffer.
state = buffer.add(state, {"obs": jnp.array([1, 2]), "reward": jnp.array(3.0)})
print(buffer.can_sample(state))  # False because min_length not reached yet.

state = buffer.add(state, {"obs": jnp.array([4, 5]), "reward": jnp.array(6.0)})
print(buffer.can_sample(state))  # Still False because we need 2 *transitions* (i.e. 3 timesteps).

state = buffer.add(state, {"obs": jnp.array([7, 8]), "reward": jnp.array(9.0)})
print(buffer.can_sample(state))  # True! We have 2 transitions (3 timesteps).

# Get a transition from the buffer.
rng_key = jax.random.PRNGKey(0)  # Source of randomness.
batch = buffer.sample(state, rng_key)  # Sample

# We have a transition! Prints: obs = [[4 5]], obs' = [[7 8]]
print(
    f"obs = {batch.experience.first['obs']}, obs' = {batch.experience.second['obs']}"
)
```

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

### Configure Training Parameters

Source: https://github.com/instadeepai/flashbax/blob/main/examples/gym_dqn_example.ipynb

Sets up hyperparameters for the DQN training process, including environment ID, buffer size, learning rate, and exploration parameters.

```python
# We specify our parameters 
env_id = "CartPole-v1"
seed = 42
num_envs = 1

total_timesteps = 50_000
learning_starts = 1_000
train_frequency = 5
target_network_frequency = 500
sample_batch_size = 128
buffer_size = 50_000
tau = 1.0
learning_rate = 1e-3
start_e = 1.0
end_e = 0.01
exploration_fraction = 0.5
gamma = 0.99
```

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

### Sample from Buffer

Source: https://github.com/instadeepai/flashbax/blob/main/examples/quickstart_flat_buffer.ipynb

Use the sample function to retrieve a batch of data from the buffer, which returns a TransitionSample object.

```python
rng_key = jax.random.PRNGKey(0)  # Setup source of randomness
batch = buffer.sample(state, rng_key)  # Sample a batch of data
```

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

### JIT Compilation and Buffer Donation

Source: https://context7.com/instadeepai/flashbax/llms.txt

Shows how to use jax.jit with donate_argnums to perform in-place buffer updates, preventing unnecessary memory copies during training.

```python
import jax
import jax.numpy as jnp
import flashbax as fbx
from functools import partial

# Create buffer
buffer = fbx.make_flat_buffer(max_length=10000, min_length=100, sample_batch_size=32)

example = {"obs": jnp.zeros((4,)), "action": jnp.array(0), "reward": jnp.array(0.0)}
state = buffer.init(example)

# Define JIT-compiled training step with buffer donation
@partial(jax.jit, donate_argnums=(1,))  # Donate buffer_state for in-place updates
def train_step(train_state, buffer_state, new_data, rng_key):
    # Add new experience
    buffer_state = buffer.add(buffer_state, new_data)

    # Sample if possible
    def do_sample(args):
        state, key = args
        batch = buffer.sample(state, key)
        # Process batch (e.g., compute loss, update params)
        return state, batch.experience.first['obs'].mean()

    def skip_sample(args):
        state, _ = args
        return state, 0.0

    buffer_state, loss = jax.lax.cond(
        buffer.can_sample(buffer_state),
        do_sample,
        skip_sample,
        (buffer_state, rng_key)
    )

    return train_state, buffer_state, loss

# Training loop
train_state = None  # Your training state here
rng_key = jax.random.PRNGKey(0)

for i in range(1000):
    rng_key, step_key = jax.random.split(rng_key)
    new_data = {
        "obs": jnp.ones((4,)) * i,
        "action": jnp.array(i % 4),
        "reward": jnp.array(float(i % 10)),
    }
    train_state, state, loss = train_step(train_state, state, new_data, step_key)
```

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

### Parallelize Buffer with jax.pmap

Source: https://github.com/instadeepai/flashbax/blob/main/examples/quickstart_flat_buffer.ipynb

Demonstrates initializing and using a buffer across multiple devices using pmap.

```python
# Define a function to create a fake batch of data
def get_fake_batch(fake_timestep: chex.ArrayTree, batch_size) -> chex.ArrayTree:
    return jax.tree.map(lambda x: jnp.stack([x + i for i in range(batch_size)]), fake_timestep)

add_batch_size = 8

# Re-instantiate the buffer
buffer = fbx.make_flat_buffer(max_length, min_length, sample_batch_size, add_sequences, add_batch_size)

# Initialize the buffer's state with a "device" dimension
fake_timestep_per_device = jax.tree.map(
    lambda x: jnp.stack([x + i for i in range(DEVICE_COUNT_MOCK)]), fake_timestep
)
state = jax.pmap(buffer.init)(fake_timestep_per_device)

# Fill the buffer above its minimum length
fake_batch = jax.pmap(get_fake_batch, static_broadcasted_argnums=1)(
    fake_timestep_per_device, add_batch_size
)
# Add two timesteps to form one transition pair
state = jax.pmap(buffer.add)(state, fake_batch)
state = jax.pmap(buffer.add)(state, fake_batch)
assert buffer.can_sample(state).all()

# Sample from the buffer
rng_key_per_device = jax.random.split(rng_key, DEVICE_COUNT_MOCK)
batch = jax.pmap(buffer.sample)(state, rng_key_per_device)
```

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

### Define Training and Evaluation Parameters

Source: https://github.com/instadeepai/flashbax/blob/main/examples/anakin_prioritised_dqn_example.ipynb

Sets up constants for training and evaluation, including batch size, learning rate, environment configurations, buffer size, and agent-specific hyperparameters like discount factor and exploration strategy.

```python
# Number of Training-Evaluation iterations
TRAINING_EVAL_ITERS = 10

# Training parameters
BATCH_SIZE = 32
LEARNING_RATE = 5e-4
SEED = 42
NUM_ENVS = 8
BUFFER_SIZE = 10_000
ROLLOUT_LEN = 512
SGD_STEPS_PER_ROLLOUT = 64
TRAINING_ITERS = 20
TARGET_PERIOD = 10
AGENT_DISCOUNT = 0.99
EPSILON_INIT = 1.0
EPSILON_FINAL = 0.1
EPSILON_STEPS = 10_000
PRIORTIY_EXPONENT = 0.6
IMPORTANCE_SAMPLING_EXPONENT = 0.6

# Evaluation parameters
NUM_EVAL_EPISODES = 50
```