### Setup PPO Training Function in Python

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/brax_minatar.ipynb

Configures and initializes the training process for Proximal Policy Optimization (PPO). This function sets up environment wrappers, learning rate schedules, and the training state, preparing for the main training loop.

```python
def make_train(config):
    config["NUM_UPDATES"] = (
        config["TOTAL_TIMESTEPS"] // config["NUM_STEPS"] // config["NUM_ENVS"]
    )
    config["MINIBATCH_SIZE"] = (
        config["NUM_ENVS"] * config["NUM_STEPS"] // config["NUM_MINIBATCHES"]
    )
    env, env_params = BraxGymnaxWrapper(config["ENV_NAME"]), None
    env = LogWrapper(env)
    env = ClipAction(env)
    env = VecEnv(env)
    if config["NORMALIZE_ENV"]:
        env = NormalizeVecObservation(env)
        env = NormalizeVecReward(env, config["GAMMA"])

    def linear_schedule(count):
        frac = (
            1.0
            - (count // (config["NUM_MINIBATCHES"] * config["UPDATE_EPOCHS"])) 
            / config["NUM_UPDATES"]
        )
        return config["LR"] * frac

    def train(rng):
        # INIT NETWORK
        network = ActorCritic(
            env.action_space(env_params).shape[0], activation=config["ACTIVATION"]
        )
        rng, _rng = jax.random.split(rng)
        init_x = jnp.zeros(env.observation_space(env_params).shape)
        network_params = network.init(_rng, init_x)
        if config["ANNEAL_LR"]:
            tx = optax.chain(
                optax.clip_by_global_norm(config["MAX_GRAD_NORM"]),
                optax.adam(learning_rate=linear_schedule, eps=1e-5),
            )
        else:
            tx = optax.chain(
                optax.clip_by_global_norm(config["MAX_GRAD_NORM"]),
                optax.adam(config["LR"], eps=1e-5),
            )
        train_state = TrainState.create(
            apply_fn=network.apply,
            params=network_params,
            tx=tx,
        )

        # INIT ENV
        rng, _rng = jax.random.split(rng)
        reset_rng = jax.random.split(_rng, config["NUM_ENVS"])
        obsv, env_state = env.reset(reset_rng, env_params)

        # TRAIN LOOP
        def _update_step(runner_state, unused):
            # COLLECT TRAJECTORIES
            def _env_step(runner_state, unused):
                train_state, env_state, last_obs, rng = runner_state

                # SELECT ACTION
                rng, _rng = jax.random.split(rng)
                pi, value = network.apply(train_state.params, last_obs)
                action = pi.sample(seed=_rng)
                log_prob = pi.log_prob(action)

                # STEP ENV
                rng, _rng = jax.random.split(rng)
                rng_step = jax.random.split(_rng, config["NUM_ENVS"])
                obsv, env_state, reward, done, info = env.step(rng_step, env_state, action, env_params)
                transition = Transition(
                    done, action, value, reward, log_prob, last_obs, info
                )
                runner_state = (train_state, env_state, obsv, rng)
                return runner_state, transition

            runner_state, traj_batch = jax.lax.scan(
                _env_step, runner_state, None, config["NUM_STEPS"]
            )

            # CALCULATE ADVANTAGE
            train_state, env_state, last_obs, rng = runner_state
            return runner_state # Placeholder for actual update logic
```

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

### DPO Brax Training Setup and Loop in Python

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/brax_minatar.ipynb

Sets up and runs the Discovered Policy Optimisation (DPO) training loop for Brax environments. It includes environment wrapping, network initialization, optimizer setup, and the core training update step using JAX's `lax.scan` for efficiency.

```python
import jax
import jax.numpy as jnp
import flax.linen as nn
import numpy as np
import optax
from flax.linen.initializers import constant, orthogonal
from typing import Sequence, NamedTuple, Any
from flax.training.train_state import TrainState
import distrax

# Assuming ActorCritic and Transition are defined as above
# Assuming BraxGymnaxWrapper, LogWrapper, ClipAction, VecEnv, NormalizeVecObservation, NormalizeVecReward are imported

class Transition(NamedTuple):
    done: jnp.ndarray
    action: jnp.ndarray
    value: jnp.ndarray
    reward: jnp.ndarray
    log_prob: jnp.ndarray
    obs: jnp.ndarray
    info: jnp.ndarray

def make_train(config):
    config["NUM_UPDATES"] = (
        config["TOTAL_TIMESTEPS"] // config["NUM_STEPS"] // config["NUM_ENVS"]
    )
    config["MINIBATCH_SIZE"] = (
        config["NUM_ENVS"] * config["NUM_STEPS"] // config["NUM_MINIBATCHES"]
    )
    env, env_params = BraxGymnaxWrapper(config["ENV_NAME"]), None
    env = LogWrapper(env)
    env = ClipAction(env)
    env = VecEnv(env)
    if config["NORMALIZE_ENV"]:
        env = NormalizeVecObservation(env)
        env = NormalizeVecReward(env, config["GAMMA"])

    def linear_schedule(count):
        frac = (
            1.0
            - (count // (config["NUM_MINIBATCHES"] * config["UPDATE_EPOCHS"])) 
            / config["NUM_UPDATES"]
        )
        return config["LR"] * frac

    def train(rng):
        # INIT NETWORK
        network = ActorCritic(
            env.action_space(env_params).shape[0], activation=config["ACTIVATION"]
        )
        rng, _rng = jax.random.split(rng)
        init_x = jnp.zeros(env.observation_space(env_params).shape)
        network_params = network.init(_rng, init_x)
        if config["ANNEAL_LR"]:
            tx = optax.chain(
                optax.clip_by_global_norm(config["MAX_GRAD_NORM"]),
                optax.adam(learning_rate=linear_schedule, eps=1e-5),
            )
        else:
            tx = optax.chain(
                optax.clip_by_global_norm(config["MAX_GRAD_NORM"]),
                optax.adam(config["LR"], eps=1e-5),
            )
        train_state = TrainState.create(
            apply_fn=network.apply,
            params=network_params,
            tx=tx,
        )

        # INIT ENV
        rng, _rng = jax.random.split(rng)
        reset_rng = jax.random.split(_rng, config["NUM_ENVS"])
        obsv, env_state = env.reset(reset_rng, env_params)

        # TRAIN LOOP
        def _update_step(runner_state, unused):
            # COLLECT TRAJECTORIES
            def _env_step(runner_state, unused):
                train_state, env_state, last_obs, rng = runner_state

                # SELECT ACTION
                rng, _rng = jax.random.split(rng)
                pi, value = network.apply(train_state.params, last_obs)
                action = pi.sample(seed=_rng)
                log_prob = pi.log_prob(action)

                # STEP ENV
                rng, _rng = jax.random.split(rng)
                rng_step = jax.random.split(_rng, config["NUM_ENVS"])
                obsv, env_state, reward, done, info = env.step(rng_step, env_state, action, env_params)
                transition = Transition(
                    done, action, value, reward, log_prob, last_obs, info
                )
                runner_state = (train_state, env_state, obsv, rng)
                return runner_state, transition

            runner_state, traj_batch = jax.lax.scan(

```

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

### Training Function Setup in Python

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/walkthrough.ipynb

Sets up the training configuration and the main training function using JAX. This includes calculating necessary parameters like NUM_UPDATES and MINIBATCH_SIZE, initializing the environment and the ActorCritic network, and defining the optimizer with optional learning rate annealing.

```python
import jax
import jax.numpy as jnp
import flax.linen as nn
import numpy as np
import optax
from flax.linen.initializers import constant, orthogonal
from typing import Sequence, NamedTuple, Any
from flax.training.train_state import TrainState
import distrax
import gymnax
from gymnax.wrappers.purerl import FlattenObservationWrapper, LogWrapper

def make_train(config):
    config["NUM_UPDATES"] = (
        config["TOTAL_TIMESTEPS"] // config["NUM_STEPS"] // config["NUM_ENVS"]
    )
    config["MINIBATCH_SIZE"] = (
        config["NUM_ENVS"] * config["NUM_STEPS"] // config["NUM_MINIBATCHES"]
    )
    env, env_params = gymnax.make(config["ENV_NAME"])
    env = FlattenObservationWrapper(env)
    env = LogWrapper(env)

    def linear_schedule(count):
        frac = 1.0 - (count // (config["NUM_MINIBATCHES"] * config["UPDATE_EPOCHS"])) / config["NUM_UPDATES"]
        return config["LR"] * frac

    def train(rng):

        # INIT NETWORK
        network = ActorCritic(env.action_space(env_params).n, activation=config["ACTIVATION"])
        rng, _rng = jax.random.split(rng)
        init_x = jnp.zeros(env.observation_space(env_params).shape)
        network_params = network.init(_rng, init_x)
        if config["ANNEAL_LR"]:
            tx = optax.chain(
                optax.clip_by_global_norm(config["MAX_GRAD_NORM"]),
                optax.adam(learning_rate=linear_schedule, eps=1e-5),
            )
        else:
            tx = optax.chain(optax.clip_by_global_norm(config["MAX_GRAD_NORM"]), optax.adam(config["LR"], eps=1e-5))
        train_state = TrainState.create(
            apply_fn=network.apply,
            params=network_params,
            tx=tx,
        )

        # INIT ENV
        rng, _rng = jax.random.split(rng)
        reset_rng = jax.random.split(_rng, config["NUM_ENVS"])
        obsv, env_state = jax.vmap(env.reset, in_axes=(0, None))(reset_rng, env_params)

        # TRAIN LOOP
        def _update_step(runner_state, unused):
            # COLLECT TRAJECTORIES
            def _env_step(runner_state, unused):
                train_state, env_state, last_obs, rng = runner_state

                # SELECT ACTION
                rng, _rng = jax.random.split(rng)
                pi, value = network.apply(train_state.params, last_obs)
                action = pi.sample(seed=_rng)
                log_prob = pi.log_prob(action)

                # STEP ENV
                rng, _rng = jax.random.split(rng)
                rng_step = jax.random.split(_rng, config["NUM_ENVS"])
                obsv, env_state, reward, done, info = jax.vmap(env.step, in_axes=(0,0,0,None))(
                    rng_step, env_state, action, env_params
                )
                transition = Transition(
                    done, action, value, reward, log_prob, last_obs, info
                )
                runner_state = (train_state, env_state, obsv, rng)
                return runner_state, transition

            runner_state, traj_batch = jax.lax.scan(
                _env_step, runner_state, None, config["NUM_STEPS"]
            )

            # CALCULATE ADVANTAGE
            train_state, env_state, last_obs, rng = runner_state
            _, last_val = network.apply(train_state.params, last_obs)


```

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

### Vectorized Multi-Seed PPO Training with jax.vmap in JAX

Source: https://context7.com/luchris429/purejaxrl/llms.txt

This example showcases how to leverage `jax.vmap` to train hundreds of agents in parallel across different random seeds. It demonstrates the efficiency gains of PureJaxRL for large-scale experiments, allowing for statistically significant results in a fraction of the time. The code vectorizes the JIT-compiled training function and measures the total training time.

```python
import jax
import jax.numpy as jnp
import time
from purejaxrl.ppo import make_train

config = {
    "LR": 2.5e-4,
    "NUM_ENVS": 4,
    "NUM_STEPS": 128,
    "TOTAL_TIMESTEPS": 5e5,
    "UPDATE_EPOCHS": 4,
    "NUM_MINIBATCHES": 4,
    "GAMMA": 0.99,
    "GAE_LAMBDA": 0.95,
    "CLIP_EPS": 0.2,
    "ENT_COEF": 0.01,
    "VF_COEF": 0.5,
    "MAX_GRAD_NORM": 0.5,
    "ACTIVATION": "tanh",
    "ENV_NAME": "CartPole-v1",
    "ANNEAL_LR": True,
}

# Create 256 different random seeds
rng = jax.random.PRNGKey(42)
 rngs = jax.random.split(rng, 256)

# Vectorize and JIT compile training across seeds
train_vjit = jax.jit(jax.vmap(make_train(config)))

# Train 256 agents in parallel
t0 = time.time()
outs = jax.block_until_ready(train_vjit(rngs))
print(f"Trained 256 seeds in {time.time() - t0:.2f}s")

# Results shape: (256, num_updates, num_envs)
mean_returns = outs["metrics"]["returned_episode_returns"].mean(axis=(0, 2))
print(f"Mean return across seeds: {mean_returns[-1]:.2f}")
```

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

### Setup JaxRL Fitness Function for Meta-Evolution

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/walkthrough.ipynb

Configures the JaxRL training function and defines a single rollout function to evaluate the performance of meta-parameters. This function calculates the mean of returned episode returns and is vectorized for efficient evaluation across multiple rollouts.

```python
config["NUM_STEPS"] = 64
train_fn = make_train(config)
def single_rollout(rng_input, meta_params):
    out = train_fn(meta_params, rng_input)
    return out["metrics"]['returned_episode_returns'].mean()
vmap_rollout = jax.vmap(single_rollout, in_axes=(0, None))
rollout = jax.jit(jax.vmap(vmap_rollout, in_axes=(None, 0)))
```

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

### Hyperparameter Search with Double vmap

Source: https://context7.com/luchris429/purejaxrl/llms.txt

This technique leverages double `vmap` in JAX to efficiently search across hyperparameter values while simultaneously running multiple random seeds for each configuration. This enables rapid and statistically robust hyperparameter optimization. The example demonstrates searching over entropy coefficients for PPO. It requires JAX and PureJaxRL's PPO module, taking a base configuration, a list of hyperparameters to search, and PRNGKeys for seeds.

```python
import jax
import jax.numpy as jnp
from purejaxrl.ppo import make_train

config = {
    "LR": 2.5e-4,
    "NUM_ENVS": 4,
    "NUM_STEPS": 128,
    "TOTAL_TIMESTEPS": 5e5,
    "UPDATE_EPOCHS": 4,
    "NUM_MINIBATCHES": 4,
    "GAMMA": 0.99,
    "GAE_LAMBDA": 0.95,
    "CLIP_EPS": 0.2,
    "VF_COEF": 0.5,
    "MAX_GRAD_NORM": 0.5,
    "ACTIVATION": "tanh",
    "ENV_NAME": "CartPole-v1",
    "ANNEAL_LR": True,
}

# Modify make_train to accept ent_coef as parameter
def make_train_with_ent_coef(config):
    # ... (same as make_train but with ent_coef as function argument)
    pass

# Search over entropy coefficients
ent_coef_search = jnp.array([0.0, 0.001, 0.005, 0.01, 0.05, 0.1])

rng = jax.random.PRNGKey(42)
rngs = jax.random.split(rng, 8)  # 8 seeds per hyperparameter

# Double vmap: outer over hyperparams, inner over seeds
train_vvjit = jax.jit(
    jax.vmap(
        jax.vmap(make_train_with_ent_coef(config), in_axes=(None, 0)),
        in_axes=(0, None)
    )
)

# Results shape: (num_hyperparams, num_seeds, num_updates, num_envs)
outs = jax.block_until_ready(train_vvjit(ent_coef_search, rngs))
```

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

### PPO Training Pipeline with make_train in JAX

Source: https://context7.com/luchris429/purejaxrl/llms.txt

This snippet demonstrates how to set up and run a Proximal Policy Optimization (PPO) training pipeline using PureJaxRL's `make_train` function. It configures training parameters, JIT-compiles the training function for performance, and executes a single training run. The output includes final model parameters and episode returns.

```python
import jax
import jax.numpy as jnp
from purejaxrl.ppo import make_train

# Configuration for PPO training
config = {
    "LR": 2.5e-4,
    "NUM_ENVS": 4,
    "NUM_STEPS": 128,
    "TOTAL_TIMESTEPS": 5e5,
    "UPDATE_EPOCHS": 4,
    "NUM_MINIBATCHES": 4,
    "GAMMA": 0.99,
    "GAE_LAMBDA": 0.95,
    "CLIP_EPS": 0.2,
    "ENT_COEF": 0.01,
    "VF_COEF": 0.5,
    "MAX_GRAD_NORM": 0.5,
    "ACTIVATION": "tanh",
    "ENV_NAME": "CartPole-v1",
    "ANNEAL_LR": True,
    "DEBUG": True,
}

# Create and JIT compile the training function
rng = jax.random.PRNGKey(42)
train_jit = jax.jit(make_train(config))

# Run training
out = train_jit(rng)

# Access results
final_params = out["runner_state"][0].params
episode_returns = out["metrics"]["returned_episode_returns"]
print(f"Final mean return: {episode_returns[-1].mean():.2f}")
```

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

### Configuration and Initialization for CartPole-v1 Training

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/walkthrough.ipynb

Sets up the configuration dictionary for the CartPole-v1 environment and initializes the JAX training process. It defines hyperparameters, environment name, and JAX random number generation keys. The `make_train` function is then JIT compiled for optimized execution.

```python
config = {
    "LR": 2.5e-4,
    "NUM_ENVS": 4,
    "NUM_STEPS": 128,
    "TOTAL_TIMESTEPS": 5e5,
    "UPDATE_EPOCHS": 4,
    "NUM_MINIBATCHES": 4,
    "GAMMA": 0.99,
    "GAE_LAMBDA": 0.95,
    "CLIP_EPS": 0.2,
    "ENT_COEF": 0.01,
    "VF_COEF": 0.5,
    "MAX_GRAD_NORM": 0.5,
    "ACTIVATION": "tanh",
    "ENV_NAME": "CartPole-v1",
    "ANNEAL_LR": True,
}

rng = jax.random.PRNGKey(42)
train_jit = jax.jit(make_train(config))
out = train_jit(rng)
```

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

### Prepare and Shuffle Minibatches for PPO Training

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/brax_minatar.ipynb

Prepares and shuffles the data into minibatches for PPO training. It asserts that the batch size matches the total number of steps and environments, then shuffles the data using a random permutation. The shuffled data is then reshaped into minibatches for processing.

```python
def _update_epoch(update_state, unused):
    def _update_minbatch(train_state, batch_info):
        traj_batch, advantages, targets = batch_info

        def _loss_fn(params, traj_batch, gae, targets):
            # RERUN NETWORK
            pi, value = network.apply(params, traj_batch.obs)
            log_prob = pi.log_prob(traj_batch.action)

            # CALCULATE VALUE LOSS
            value_pred_clipped = traj_batch.value + (
                value - traj_batch.value
            ).clip(-config["CLIP_EPS"], config["CLIP_EPS"])
            value_losses = jnp.square(value - targets)
            value_losses_clipped = jnp.square(value_pred_clipped - targets)
            value_loss = (
                0.5 * jnp.maximum(value_losses, value_losses_clipped).mean()
            )

            # CALCULATE ACTOR LOSS
            alpha = config["DPO_ALPHA"]
            beta = config["DPO_BETA"]
            log_diff = log_prob - traj_batch.log_prob
            ratio = jnp.exp(log_diff)
            gae = (gae - gae.mean()) / (gae.std() + 1e-8)
            is_pos = (gae >= 0.0).astype("float32")
            r1 = ratio - 1.0
            drift1 = nn.relu(r1 * gae - alpha * nn.tanh(r1 * gae / alpha))
            drift2 = nn.relu(
                log_diff * gae - beta * nn.tanh(log_diff * gae / beta)
            )
            drift = drift1 * is_pos + drift2 * (1 - is_pos)
            loss_actor = -(ratio * gae - drift).mean()
            entropy = pi.entropy().mean()

            total_loss = (
                loss_actor
                + config["VF_COEF"] * value_loss
                - config["ENT_COEF"] * entropy
            )
            return total_loss, (value_loss, loss_actor, entropy)

        grad_fn = jax.value_and_grad(_loss_fn, has_aux=True)
        total_loss, grads = grad_fn(
            train_state.params, traj_batch, advantages, targets
        )
        train_state = train_state.apply_gradients(grads=grads)
        return train_state, total_loss

    train_state, traj_batch, advantages, targets, rng = update_state
    rng, _rng = jax.random.split(rng)
    batch_size = config["MINIBATCH_SIZE"] * config["NUM_MINIBATCHES"]
    assert (
        batch_size == config["NUM_STEPS"] * config["NUM_ENVS"]
    ), "batch size must be equal to number of steps * number of envs"
    permutation = jax.random.permutation(_rng, batch_size)
    batch = (traj_batch, advantages, targets)
    batch = jax.tree_util.tree_map(
        lambda x: x.reshape((batch_size,) + x.shape[2:]),
        batch
    )
    shuffled_batch = jax.tree_util.tree_map(
        lambda x: jnp.take(x, permutation, axis=0),
        batch
    )
    minibatches = jax.tree_util.tree_map(
        lambda x: jnp.reshape(
            x, [config["NUM_MINIBATCHES"], -1] + list(x.shape[1:])
        ),
        shuffled_batch
    )
    train_state, total_loss = jax.lax.scan(

```

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

### JAX Configuration and Training Initialization

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/brax_minatar.ipynb

This Python code snippet defines a configuration dictionary for the training process and initializes a JAX-compiled training function. The configuration includes hyperparameters such as learning rate, batch sizes, and environment settings. The `jax.jit` decorator is used to compile the `make_train` function for performance.

```python
config = {
    "LR": 3e-4,
    "NUM_ENVS": 2048,
    "NUM_STEPS": 10,
    "TOTAL_TIMESTEPS": 5e7,
    "UPDATE_EPOCHS": 4,
    "NUM_MINIBATCHES": 32,
    "GAMMA": 0.99,
    "GAE_LAMBDA": 0.95,
    "CLIP_EPS": 0.2,
    "DPO_ALPHA": 2.0,
    "DPO_BETA": 0.6,
    "ENT_COEF": 0.0,
    "VF_COEF": 0.5,
    "MAX_GRAD_NORM": 0.5,
    "ACTIVATION": "tanh",
    "ENV_NAME": "hopper",
    "ANNEAL_LR": False,
    "NORMALIZE_ENV": True,
    "DEBUG": True,
}
rng = jax.random.PRNGKey(30)
train_jit = jax.jit(make_train(config))
out = train_jit(rng)
```

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

### JAX Configuration and Training Initialization

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/brax_minatar.ipynb

Defines the configuration dictionary for the training process, including hyperparameters like learning rate, batch sizes, and environment settings. It then initializes a JAX PRNG key and JIT-compiles the training function for performance.

```python
config = {
    "LR": 3e-4,
    "NUM_ENVS": 2048,
    "NUM_STEPS": 10,
    "TOTAL_TIMESTEPS": 5e7,
    "UPDATE_EPOCHS": 4,
    "NUM_MINIBATCHES": 32,
    "GAMMA": 0.99,
    "GAE_LAMBDA": 0.95,
    "CLIP_EPS": 0.2,
    "ENT_COEF": 0.0,
    "VF_COEF": 0.5,
    "MAX_GRAD_NORM": 0.5,
    "ACTIVATION": "tanh",
    "ENV_NAME": "hopper",
    "ANNEAL_LR": False,
    "NORMALIZE_ENV": True,
    "DEBUG": True,
}
rng = jax.random.PRNGKey(30)
train_jit = jax.jit(make_train(config))
out = train_jit(rng)
```

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

### Initialize OpenES Strategy for Meta-Evolution

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/walkthrough.ipynb

Initializes the OpenES (Open Evolutionary Strategies) algorithm. This involves creating an instance of the MetaLossNetwork, initializing its parameters, and then setting up the OpenES strategy with specified population size, optimizer, and other relevant parameters. The initial state for the strategy is also prepared.

```python
meta_loss_network = MetaLossNetwork()
meta_params_pholder = meta_loss_network.init(jax.random.PRNGKey(0), jnp.zeros((1, 1)))
rng = jax.random.PRNGKey(42)
rng, rng_init = jax.random.split(rng)
strategy = OpenES(
    popsize=popsize,
    pholder_params=meta_params_pholder,
    opt_name="adam",
    centered_rank=True,
    maximize=True,
)
state = strategy.initialize(rng_init)
```

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

### Configure and Train with JAX

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/brax_minatar.ipynb

This code defines a configuration dictionary for a reinforcement learning experiment and then uses JAX to compile and execute a training function. It initializes the random number generator and runs the training process.

```python
config = {
    "LR": 5e-3,
    "NUM_ENVS": 64,
    "NUM_STEPS": 128,
    "TOTAL_TIMESTEPS": 1e7,
    "UPDATE_EPOCHS": 4,
    "NUM_MINIBATCHES": 8,
    "GAMMA": 0.99,
    "GAE_LAMBDA": 0.95,
    "CLIP_EPS": 0.2,
    "ENT_COEF": 0.01,
    "VF_COEF": 0.5,
    "MAX_GRAD_NORM": 0.5,
    "ACTIVATION": "relu",
    "ENV_NAME": "SpaceInvaders-MinAtar",
    "ANNEAL_LR": True,
}

rng = jax.random.PRNGKey(42)
train_jit = jax.jit(make_train(config))
out = train_jit(rng)
```

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

### Create BraxGymnaxWrapper for Environment Integration

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/brax_minatar.ipynb

This wrapper integrates Brax environments into the Gymnax framework. It initializes a Brax environment, applies episode and auto-reset wrappers, and provides methods for reset and step operations, mimicking the Gymnax interface. It exposes action and observation sizes and defines observation and action spaces.

```python
import jax
import jax.numpy as jnp
import chex
import numpy as np
from flax import struct
from functools import partial
from typing import Optional, Tuple, Union, Any
from gymnax.environments import environment, spaces
from gymnax.wrappers.purerl import GymnaxWrapper
from brax import envs
from brax.envs.wrappers.training import EpisodeWrapper, AutoResetWrapper

class BraxGymnaxWrapper:
    def __init__(self, env_name, backend="positional"):
        env = envs.get_environment(env_name=env_name, backend=backend)
        env = EpisodeWrapper(env, episode_length=1000, action_repeat=1)
        env = AutoResetWrapper(env)
        self._env = env
        self.action_size = env.action_size
        self.observation_size = (env.observation_size,)

    def reset(self, key, params=None):
        state = self._env.reset(key)
        return state.obs, state

    def step(self, key, state, action, params=None):
        next_state = self._env.step(state, action)
        return next_state.obs, next_state, next_state.reward, next_state.done > 0.5, {}

    def observation_space(self, params):
        return spaces.Box(
            low=-jnp.inf,
            high=jnp.inf,
            shape=(self._env.observation_size,),
        )

    def action_space(self, params):
        return spaces.Box(
            low=-1.0,
            high=1.0,
            shape=(self._env.action_size,),
        )
```

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

### Environment Wrappers for Gymnax and Brax

Source: https://context7.com/luchris429/purejaxrl/llms.txt

Demonstrates the use of various environment wrappers in PureJaxRL for preprocessing observations, logging metrics, and vectorizing environments. It covers both discrete action spaces (Gymnax) and continuous action spaces (Brax).

```python
import jax
import gymnax
from purejaxrl.wrappers import (
    FlattenObservationWrapper,
    LogWrapper,
    BraxGymnaxWrapper,
    VecEnv,
    NormalizeVecObservation,
    NormalizeVecReward,
    ClipAction,
)

# Discrete action environment with Gymnax
env, env_params = gymnax.make("CartPole-v1")
env = FlattenObservationWrapper(env)  # Flatten observations
env = LogWrapper(env)                  # Track episode returns/lengths

# Reset and step
rng = jax.random.PRNGKey(0)
obs, state = env.reset(rng, env_params)
action = env.action_space(env_params).sample(rng)
obs, state, reward, done, info = env.step(rng, state, action, env_params)

# Access logged metrics
print(f"Episode return: {info['returned_episode_returns']}")
print(f"Episode length: {info['returned_episode_lengths']}")

# Continuous action environment with Brax
env = BraxGymnaxWrapper("hopper", backend="positional")
env = LogWrapper(env)
env = ClipAction(env)
env = VecEnv(env)  # Vectorize for batch processing
env = NormalizeVecObservation(env)
env = NormalizeVecReward(env, gamma=0.99)

# Reset vectorized environment
num_envs = 2048
reset_rngs = jax.random.split(rng, num_envs)
obs, state = env.reset(reset_rngs, None)
print(f"Observation shape: {obs.shape}")  # (2048, obs_dim)
```

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

### DQN Training Pipeline with Flashbax

Source: https://context7.com/luchris429/purejaxrl/llms.txt

Implements a Deep Q-Network training pipeline featuring experience replay with Flashbax, epsilon-greedy exploration, and target network updates. It requires JAX and PureJaxRL's DQN module. The input is a configuration dictionary and a PRNGKey, and the output includes training metrics like returns.

```python
import jax
from purejaxrl.dqn import make_train

config = {
    "NUM_ENVS": 10,
    "BUFFER_SIZE": 10000,
    "BUFFER_BATCH_SIZE": 128,
    "TOTAL_TIMESTEPS": 5e5,
    "EPSILON_START": 1.0,
    "EPSILON_FINISH": 0.05,
    "EPSILON_ANNEAL_TIME": 25e4,
    "TARGET_UPDATE_INTERVAL": 500,
    "LR": 2.5e-4,
    "LEARNING_STARTS": 10000,
    "TRAINING_INTERVAL": 10,
    "LR_LINEAR_DECAY": False,
    "GAMMA": 0.99,
    "TAU": 1.0,
    "ENV_NAME": "CartPole-v1",
    "SEED": 0,
    "NUM_SEEDS": 1,
    "WANDB_MODE": "disabled",
}

# Train DQN agent
rng = jax.random.PRNGKey(config["SEED"])
rngs = jax.random.split(rng, config["NUM_SEEDS"])
train_vjit = jax.jit(jax.vmap(make_train(config)))
outs = jax.block_until_ready(train_vjit(rngs))

# Access metrics
returns = outs["metrics"]["returns"]
print(f"Final return: {returns[-1].mean():.2f}")
```

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

### Evolve Meta-Parameters using OpenES

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/walkthrough.ipynb

This Python snippet demonstrates the initialization of the OpenES evolutionary strategy from the evosax library. It sets hyperparameters such as population size, number of generations, and number of rollouts, preparing for an evolutionary optimization process on a task like CartPole-v1.

```python
from evosax import OpenES

# Hyperparameters
popsize = 128
num_generations = 512
num_rollouts = 2
```

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

### Meta-Learning with Evosax and PureJaxRL

Source: https://context7.com/luchris429/purejaxrl/llms.txt

Illustrates how to integrate PureJaxRL with Evosax for meta-learning RL components through evolution. This includes defining a meta-learned loss network, setting up fitness evaluation, and running an evolution loop.

```python
import jax
import jax.numpy as jnp
import flax.linen as nn
from evosax import OpenES
from purejaxrl.ppo import make_train

# Define a meta-learned loss network
class MetaLossNetwork(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(64)(x)
        x = nn.tanh(x)
        x = nn.Dense(64)(x)
        x = nn.tanh(x)
        x = nn.Dense(1)(x)
        return x

# Hyperparameters
popsize = 128
num_generations = 100
num_rollouts = 2

config = {
    "LR": 2.5e-4,
    "NUM_ENVS": 4,
    "NUM_STEPS": 64,
    "TOTAL_TIMESTEPS": 5e5,
    "UPDATE_EPOCHS": 4,
    "NUM_MINIBATCHES": 4,
    "GAMMA": 0.99,
    "GAE_LAMBDA": 0.95,
    "CLIP_EPS": 0.2,
    "ENT_COEF": 0.01,
    "VF_COEF": 0.5,
    "MAX_GRAD_NORM": 0.5,
    "ACTIVATION": "tanh",
    "ENV_NAME": "CartPole-v1",
    "ANNEAL_LR": True,
}

# Setup fitness evaluation
def single_rollout(rng_input, meta_params):
    out = train_fn(meta_params, rng_input)
    return out["metrics"]["returned_episode_returns"].mean()

vmap_rollout = jax.vmap(single_rollout, in_axes=(0, None))
rollout = jax.jit(jax.vmap(vmap_rollout, in_axes=(None, 0)))

# Initialize evolution strategy
meta_loss_network = MetaLossNetwork()
meta_params_pholder = meta_loss_network.init(jax.random.PRNGKey(0), jnp.zeros((1, 1)))
rng = jax.random.PRNGKey(42)
strategy = OpenES(
    popsize=popsize,
    pholder_params=meta_params_pholder,
    opt_name="adam",
    centered_rank=True,
    maximize=True,
)
state = strategy.initialize(rng)

# Run evolution loop
for gen in range(num_generations):
    rng, rng_ask, rng_eval = jax.random.split(rng, 3)
    x, state = strategy.ask(rng_ask, state)
    batch_rng = jax.random.split(rng_eval, num_rollouts)
    fitness = rollout(batch_rng, x).mean(axis=1)
    state = strategy.tell(x, fitness, state)
    print(f"Gen {gen}: Mean={fitness.mean():.2f}, Best={state.best_fitness:.2f}")
```

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

### Reinforcement Learning Training Function in JAX

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/brax_minatar.ipynb

Sets up and executes a reinforcement learning training loop using JAX. It configures the environment, network, optimizer, and handles the training steps, including environment interaction and gradient updates. Dependencies include JAX, Optax, and Gymnax.

```python
import jax
import jax.numpy as jnp
import flax.linen as nn
import numpy as np
import optax
from flax.linen.initializers import constant, orthogonal
from typing import Sequence, NamedTuple, Any
from flax.training.train_state import TrainState
import distrax
import gymnax
from gymnax.wrappers.purerl import LogWrapper, FlattenObservationWrapper

class ActorCritic(nn.Module):
    action_dim: Sequence[int]
    activation: str = "tanh"

    @nn.compact
    def __call__(self, x):
        if self.activation == "relu":
            activation = nn.relu
        else:
            activation = nn.tanh
        actor_mean = nn.Dense(
            64, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0)
        )(x)
        actor_mean = activation(actor_mean)
        actor_mean = nn.Dense(
            64, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0)
        )(actor_mean)
        actor_mean = activation(actor_mean)
        actor_mean = nn.Dense(
            self.action_dim, kernel_init=orthogonal(0.01), bias_init=constant(0.0)
        )(actor_mean)
        pi = distrax.Categorical(logits=actor_mean)

        critic = nn.Dense(
            64, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0)
        )(x)
        critic = activation(critic)
        critic = nn.Dense(
            64, kernel_init=orthogonal(np.sqrt(2)), bias_init=constant(0.0)
        )(critic)
        critic = activation(critic)
        critic = nn.Dense(1, kernel_init=orthogonal(1.0), bias_init=constant(0.0))(
            critic
        )

        return pi, jnp.squeeze(critic, axis=-1)


class Transition(NamedTuple):
    done: jnp.ndarray
    action: jnp.ndarray
    value: jnp.ndarray
    reward: jnp.ndarray
    log_prob: jnp.ndarray
    obs: jnp.ndarray
    info: jnp.ndarray

def make_train(config):
    config["NUM_UPDATES"] = (
        config["TOTAL_TIMESTEPS"] // config["NUM_STEPS"] // config["NUM_ENVS"]
    )
    config["MINIBATCH_SIZE"] = (
        config["NUM_ENVS"] * config["NUM_STEPS"] // config["NUM_MINIBATCHES"]
    )
    env, env_params = gymnax.make(config["ENV_NAME"])
    env = FlattenObservationWrapper(env)
    env = LogWrapper(env)

    def linear_schedule(count):
        frac = 1.0 - (count // (config["NUM_MINIBATCHES"] * config["UPDATE_EPOCHS"])) / config["NUM_UPDATES"]
        return config["LR"] * frac

    def train(rng):

        # INIT NETWORK
        network = ActorCritic(env.action_space(env_params).n, activation=config["ACTIVATION"])
        rng, _rng = jax.random.split(rng)
        init_x = jnp.zeros(env.observation_space(env_params).shape)
        network_params = network.init(_rng, init_x)
        if config["ANNEAL_LR"]:
            tx = optax.chain(
                optax.clip_by_global_norm(config["MAX_GRAD_NORM"]),
                optax.adam(learning_rate=linear_schedule, eps=1e-5),
            )
        else:
            tx = optax.chain(optax.clip_by_global_norm(config["MAX_GRAD_NORM"]), optax.adam(config["LR"], eps=1e-5))
        train_state = TrainState.create(
            apply_fn=network.apply,
            params=network_params,
            tx=tx,
        )

        # INIT ENV
        rng, _rng = jax.random.split(rng)
        reset_rng = jax.random.split(_rng, config["NUM_ENVS"])
        obsv, env_state = jax.vmap(env.reset, in_axes=(0, None))(reset_rng, env_params)

        # TRAIN LOOP
        def _update_step(runner_state, unused):
            # COLLECT TRAJECTORIES
            def _env_step(runner_state, unused):
                train_state, env_state, last_obs, rng = runner_state

                # SELECT ACTION
                rng, _rng = jax.random.split(rng)
                pi, value = network.apply(train_state.params, last_obs)
                action = pi.sample(seed=_rng)
                log_prob = pi.log_prob(action)

                # STEP ENV
                rng, _rng = jax.random.split(rng)
                rng_step = jax.random.split(_rng, config["NUM_ENVS"])
                obsv, env_state, reward, done, info = jax.vmap(env.step, in_axes=(0,0,0,None))(
                    rng_step, env_state, action, env_params
                )
                transition = Transition(
                    done, action, value, reward, log_prob, last_obs, info
                )
                runner_state = (train_state, env_state, obsv, rng)
                return runner_state, transition

            runner_state, traj_batch = jax.lax.scan(
                _env_step, runner_state, None, config["NUM_STEPS"]
            )

            # CALCULATE ADVANTAGE
            return traj_batch

        return train_state, obsv, env_state, rng
```

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

### PPO for Continuous Action Spaces with Brax

Source: https://context7.com/luchris429/purejaxrl/llms.txt

This PPO implementation is designed for continuous action spaces and integrates with Brax physics environments. It employs a Gaussian policy with a learned standard deviation and supports optional normalization for observations and rewards. Dependencies include JAX and PureJaxRL's PPO-Continuous-Action module. The training function accepts a configuration and a PRNGKey, outputting training metrics such as episode returns.

```python
import jax
from purejaxrl.ppo_continuous_action import make_train

config = {
    "LR": 3e-4,
    "NUM_ENVS": 2048,
    "NUM_STEPS": 10,
    "TOTAL_TIMESTEPS": 5e7,
    "UPDATE_EPOCHS": 4,
    "NUM_MINIBATCHES": 32,
    "GAMMA": 0.99,
    "GAE_LAMBDA": 0.95,
    "CLIP_EPS": 0.2,
    "ENT_COEF": 0.0,
    "VF_COEF": 0.5,
    "MAX_GRAD_NORM": 0.5,
    "ACTIVATION": "tanh",
    "ENV_NAME": "hopper",
    "ANNEAL_LR": False,
    "NORMALIZE_ENV": True,
    "DEBUG": True,
}

rng = jax.random.PRNGKey(30)
train_jit = jax.jit(make_train(config))
out = train_jit(rng)

episode_returns = out["metrics"]["returned_episode_returns"]
print(f"Final return: {episode_returns[-1].mean():.2f}")
```

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

### Vectorized Training and Plotting with JAX

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/walkthrough.ipynb

Demonstrates running multiple training instances in parallel using `jax.vmap` for increased efficiency. It times the vectorized execution and plots the results for each of the 256 parallel runs, showing the mean episode returns over update steps.

```python
rng = jax.random.PRNGKey(42)
rngs = jax.random.split(rng, 256)
train_vjit = jax.jit(jax.vmap(make_train(config)))
t0 = time.time()
outs = jax.block_until_ready(train_vjit(rngs))
print(f"time: {time.time() - t0:.2f} s")

for i in range(256):
    plt.plot(outs["metrics"]["returned_episode_returns"][i].mean(-1).reshape(-1))
plt.xlabel("Update Step")
plt.ylabel("Return")
plt.show()
```

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

### Distributed Training with JAX vmap

Source: https://github.com/luchris429/purejaxrl/blob/main/examples/walkthrough.ipynb

This Python code demonstrates distributed training using JAX's `vmap` for parallel execution. It sets up an entropy coefficient search and multiple random number generator keys, then applies `jax.jit` and nested `jax.vmap` to create a highly optimized training function. The execution time is measured and printed.

```python
ent_coef_search = jnp.array([0.0, 0.001, 0.0005, 0.01, 0.05, 0.1, 0.5])

rng = jax.random.PRNGKey(42)
rngs = jax.random.split(rng, 8)

train_vvjit = jax.jit(
    jax.vmap(
        jax.vmap(
            make_train(config), in_axes=(None, 0)
        ), in_axes=(0, None)
    )
)
t0 = time.time()
outs = jax.block_until_ready(train_vvjit(ent_coef_search, rngs))
print(f"time: {time.time() - t0:.2f} s")
```