### Install RLax

Source: https://context7.com/google-deepmind/rlax/llms.txt

Installs the RLax library using pip. Can install from PyPI or directly from the source repository.

```bash
pip install rlax
# or latest from source:
pip install git+https://github.com/google-deepmind/rlax.git
```

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

### Install RLax Development Version

Source: https://github.com/google-deepmind/rlax/blob/main/README.md

Install the latest development version of RLax directly from its GitHub repository.

```sh
pip install git+https://github.com/deepmind/rlax.git
```

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

### Install RLax from PyPI

Source: https://github.com/google-deepmind/rlax/blob/main/README.md

Use this command to install the latest released version of RLax from the Python Package Index.

```sh
pip install rlax
```

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

### Complete Example: Double DQN Agent

Source: https://context7.com/google-deepmind/rlax/llms.txt

A full working implementation using `rlax.double_q_learning`, `rlax.epsilon_greedy`, and `rlax.l2_loss` to train a DQN on BSuite's Catch environment.

```python
import collections
import jax
import jax.numpy as jnp
import haiku as hk
from haiku import nets
import optax
import rlax
```

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

### Actor Step Function with Epsilon-Greedy Exploration

Source: https://context7.com/google-deepmind/rlax/llms.txt

Performs a single step of the actor, applying the online network to get Q-values and sampling an action using epsilon-greedy exploration. Epsilon is annealed over time.

```python
@jax.jit
def actor_step(params, obs, step_count, key):
    q = network.apply(params.online, obs[None])[0]
    # Anneal epsilon from 1.0 -> 0.01 over 1000 steps
    epsilon = jnp.maximum(0.01, 1.0 - step_count / 1000.0)
    return rlax.epsilon_greedy(epsilon).sample(key, q)
```

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

### rlax.sample_start_indices

Source: https://github.com/google-deepmind/rlax/blob/main/docs/api.md

Samples a specified number of starting indices for a batch of trajectories. Ensures that all sampled indices are within the valid range.

```APIDOC
## rlax.sample_start_indices(rng_key: Array, batch_size: int, num_start_indices: int, max_valid_start_idx: int) -> Array | ndarray | bool | number

### Description
Sampling batch_size x num_start_indices starting indices.

### Parameters
#### Path Parameters
- **rng_key** (Array) - A pseudo random number generator’s key.
- **batch_size** (int) - The size of the batch of trajectories to index in.
- **num_start_indices** (int) - How many starting points per trajectory in the batch.
- **max_valid_start_idx** (int) - Maximum valid time index for all starting points.

### Returns
An array of starting points with shape [B, num_start_indices].
```

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

### General Off Policy Returns From Action Values

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Provides methods for calculating general off-policy returns starting from action values.

```APIDOC
## General Off Policy Returns From Action Values

### Description
Provides general off-policy returns from action values.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### QV Max

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements QV Max, a variant of QV learning.

```APIDOC
## QV Max

### Description
Implements QV Max.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### QV Learning

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements QV Learning, an algorithm that learns both value and Q-values.

```APIDOC
## QV Learning

### Description
Implements QV Learning.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### Q-Lambda

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements Q-Lambda, an algorithm that combines Q-learning with eligibility traces.

```APIDOC
## Q-Lambda

### Description
Implements Q-Lambda.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### Q Learning

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements the standard Q Learning algorithm, a model-free reinforcement learning method.

```APIDOC
## Q Learning

### Description
Implements Q Learning.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### Compute Double Q-learning TD Error with rlax.double_q_learning

Source: https://context7.com/google-deepmind/rlax/llms.txt

Computes the Double Q-learning TD error, which reduces overestimation bias by decoupling action selection and value estimation. The example demonstrates a single step calculation and its use within a batched loss function.

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

q_tm1        = jnp.array([0.1, 0.5, 0.3, 0.2])  # online Q(s,·)
a_tm1        = jnp.int32(1)
r_t          = jnp.float32(1.0)
discount_t   = jnp.float32(0.99)
q_t_value    = jnp.array([0.2, 0.7, 0.6, 0.4])  # target network Q(s',·)
q_t_selector = jnp.array([0.3, 0.8, 0.5, 0.4])  # online network Q(s',·) for selection

td_error = rlax.double_q_learning(
    q_tm1, a_tm1, r_t, discount_t, q_t_value, q_t_selector)
# online selects action 1 (argmax of q_t_selector)
# target evaluates: 1.0 + 0.99*q_t_value[1] = 1.0 + 0.99*0.7 = 1.693
# error = 1.693 - 0.5 = 1.193
print(td_error)  # 1.193

# Used in the simple DQN example:
batched_loss = jax.vmap(rlax.double_q_learning)
td_error = batched_loss(q_tm1, a_tm1, r_t, discount_t, q_t_val, q_t_select)
loss = jnp.mean(rlax.l2_loss(td_error))
```

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

### Compute Q-learning TD Error with rlax.q_learning

Source: https://context7.com/google-deepmind/rlax/llms.txt

Computes the Q-learning TD error using the max-Q bootstrap target: `δ = r + γ·max_a Q(s',a) - Q(s,a)`. This is the core update for off-policy value-based agents. The example shows a single step calculation and how to use `jax.vmap` with `rlax.l2_loss` for a training loop.

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

# 4-action environment, single step
q_tm1      = jnp.array([0.1, 0.5, 0.3, 0.2])  # Q(s, ·)
a_tm1      = jnp.int32(1)                        # action taken (index 1)
r_t        = jnp.float32(1.0)                    # reward received
discount_t = jnp.float32(0.99)
q_t        = jnp.array([0.2, 0.8, 0.6, 0.4])   # Q(s', ·)

td_error = rlax.q_learning(q_tm1, a_tm1, r_t, discount_t, q_t)
# target = 1.0 + 0.99 * max(q_t) = 1.0 + 0.99*0.8 = 1.792
# error  = 1.792 - q_tm1[1] = 1.792 - 0.5 = 1.292
print(td_error)  # 1.292

# In a training loop, use vmap + l2_loss:
batched_loss_fn = jax.vmap(rlax.q_learning)
td_errors = batched_loss_fn(q_tm1_batch, a_tm1_batch, r_t_batch,
                            discount_t_batch, q_t_batch)
loss = jnp.mean(rlax.l2_loss(td_errors))
```

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

### Compute GAE for PPO

Source: https://context7.com/google-deepmind/rlax/llms.txt

Computes Generalized Advantage Estimation (GAE) for PPO and actor-critic agents. Advantages are calculated using a lambda parameter for weighting.

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

T = 5
r_t        = jnp.array([1.0, 0.0, 1.0, -1.0, 0.5])
discount_t = jnp.array([0.99, 0.99, 0.99, 0.99, 0.0])
lambda_    = jnp.float32(0.95)
values     = jnp.array([0.5, 0.8, 0.6, 0.4, 0.2, 0.0])  # length T+1

advantages = rlax.truncated_generalized_advantage_estimation(
    r_t, discount_t, lambda_, values)
print(advantages.shape)  # (5,)
# Use advantages in policy gradient loss:
logits = jnp.array([[0.1, 0.5, 0.4]]*T)
actions = jnp.array([1, 0, 2, 1, 0])
weights = jnp.ones(T)
pg_loss = rlax.policy_gradient_loss(logits, actions, advantages, weights)
```

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

### Define Network and Optimizer

Source: https://context7.com/google-deepmind/rlax/llms.txt

Sets up a simple MLP network and an Adam optimizer. The network is wrapped with `hk.without_apply_rng` and `hk.transform` for use with JAX.

```python
Params      = collections.namedtuple("Params", "online target")
LearnerState = collections.namedtuple("LearnerState", "count opt_state")

def build_network(num_actions):
    def q(obs):
        net = hk.Sequential([hk.Flatten(), nets.MLP([64, num_actions])])
        return net(obs)
    return hk.without_apply_rng(hk.transform(q))

network   = build_network(num_actions=3)
optimizer = optax.adam(1e-3)
```

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

### Epsilon-Greedy Action Selection

Source: https://context7.com/google-deepmind/rlax/llms.txt

Selects actions using an epsilon-greedy strategy for exploration during training or a pure greedy strategy during evaluation. Epsilon can be annealed over training steps.

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

key = jax.random.PRNGKey(42)
q_values = jnp.array([0.1, 0.8, 0.3, 0.5])  # Q(s, ·)

# Training: epsilon-greedy with ε=0.1 (10% random, 90% greedy)
train_action = rlax.epsilon_greedy(epsilon=0.1).sample(key, q_values)
print(train_action)  # integer action index

# Evaluation: pure greedy
eval_action = rlax.greedy().sample(key, q_values)
print(eval_action)  # 1 (argmax)
```

```python
# Annealing epsilon over training (common DQN pattern):
import optax
epsilon_schedule = optax.polynomial_schedule(
    init_value=1.0, end_value=0.01,
    transition_steps=10000, power=1.0)
step = jnp.int32(5000)
epsilon = epsilon_schedule(step)
action = rlax.epsilon_greedy(epsilon).sample(key, q_values)
```

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

### Quantile Q Learning

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements Quantile Q Learning, a distributional reinforcement learning algorithm.

```APIDOC
## Quantile Q Learning

### Description
Implements Quantile Q Learning.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### rlax.retrace_continuous

Source: https://github.com/google-deepmind/rlax/blob/main/docs/api.md

Retrace continuous implementation for off-policy reinforcement learning.

```APIDOC
## rlax.retrace_continuous

### Description
Retrace continuous.

### Parameters
* **q_tm1** (Array | ndarray | bool | number) - Q-values at times [0, …, K - 1].
* **q_t** (Array | ndarray | bool | number) - Q-values evaluated at actions collected using behavior policy at times [1, …, K - 1].
* **v_t** (Array | ndarray | bool | number) - Value estimates of the target policy at times [1, …, K].
* **r_t** (Array | ndarray | bool | number) - reward at times [1, …, K].
* **discount_t** (Array | ndarray | bool | number) - discount at times [1, …, K].
* **log_rhos** (Array | ndarray | bool | number) - Log importance weight pi_target/pi_behavior evaluated at actions collected using behavior policy [1, …, K - 1].
* **lambda_** (Array | ndarray | bool | number | float) - scalar or a vector of mixing parameter lambda.
* **stop_target_gradients** (bool) - bool indicating whether or not to apply stop gradient to targets. Defaults to True.

### Returns
(Array | ndarray | bool | number) - Retrace error.
```

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

### rlax.n_step_bootstrapped_returns

Source: https://context7.com/google-deepmind/rlax/llms.txt

Computes strided n-step bootstrapped returns, accumulating n rewards before bootstrapping from the value function.

```APIDOC
## rlax.n_step_bootstrapped_returns — n-step returns

Computes strided n-step bootstrapped returns, accumulating n rewards before bootstrapping from the value function.

### Parameters
- **r_t** (numeric array): Rewards received at time t.
- **discount_t** (numeric array): Discount factor at time t.
- **v_t** (numeric array): Value function estimates at time t.
- **n** (int): The number of steps for n-step returns.

### Returns
- **returns** (numeric array): The computed n-step bootstrapped returns.
```

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

### Initialize and Update Pop-Art State

Source: https://context7.com/google-deepmind/rlax/llms.txt

Implements Pop-Art for adaptive reward normalization, crucial for environments with varying reward scales. Initializes state and provides an update function for network parameters.

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

# Setup: 2 output heads (tasks), step_size=1e-4
init_state_fn, popart_update_fn = rlax.popart(
    num_outputs=2, step_size=1e-4, scale_lb=1e-4, scale_ub=1e6)

# Initialize PopArt state
state = init_state_fn()
print(state.shift)        # [0.0, 0.0] — running mean
print(state.scale)        # [1.0, 1.0] — running std

# Dummy linear layer params (weights [hidden, 2], bias [2])
key = jax.random.PRNGKey(0)
params = {
    "w": jax.random.normal(key, (16, 2)),
    "b": jnp.zeros(2),
}

# Targets for batch: 8 samples, task indices
targets = jnp.array([10.0, 20.0, 15.0, 25.0, 12.0, 18.0, 22.0, 8.0])
indices = jnp.array([0, 1, 0, 1, 0, 1, 0, 1])  # which task head

new_state, new_params = popart_update_fn(state, params, targets, indices)

# Normalize/unnormalize values:
normalized   = rlax.normalize(new_state, targets, indices)
unnormalized = rlax.unnormalize(new_state, normalized, indices)
```

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

### rlax.quantile_q_learning

Source: https://github.com/google-deepmind/rlax/blob/main/docs/api.md

Implements Q-learning for quantile-valued Q distributions. This function is part of distributional reinforcement learning approaches.

```APIDOC
## rlax.quantile_q_learning

### Description
Implements Q-learning for quantile-valued Q distributions.

See “Distributional Reinforcement Learning with Quantile Regression” by
Dabney et al. ([https://arxiv.org/abs/1710.10044](https://arxiv.org/abs/1710.10044)).

### Parameters
* **dist_q_tm1** – Q distribution at time t-1.
* **tau_q_tm1** – Q distribution probability thresholds.
* **a_tm1** – action index at time t-1.
* **r_t** – reward at time t.
* **discount_t** – discount at time t.
* **dist_q_t_selector** – Q distribution at time t for selecting greedy action in
  target policy. This is separate from dist_q_t as in Double Q-Learning, but
  can be computed with the target network and a separate set of samples.
* **dist_q_t** – target Q distribution at time t.
* **huber_param** – Huber loss parameter, defaults to 0 (no Huber loss).
* **stop_target_gradients** – bool indicating whether or not to apply stop gradient
  to targets.

### Returns
Quantile regression Q learning loss.
```

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

### SARSA Lambda

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements SARSA Lambda, which combines SARSA with eligibility traces.

```APIDOC
## SARSA Lambda

### Description
Implements SARSA Lambda.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### rlax.epsilon_greedy / rlax.greedy

Source: https://context7.com/google-deepmind/rlax/llms.txt

Provides epsilon-greedy and greedy distributions for selecting discrete actions from Q-value estimates during training and evaluation.

```APIDOC
## rlax.epsilon_greedy / rlax.greedy — Action selection distributions

### Description
Provides epsilon-greedy and greedy distributions for selecting discrete actions from Q-value estimates during training and evaluation.

### Usage
```python
import jax
import jax.numpy as jnp
import rlax

key = jax.random.PRNGKey(42)
q_values = jnp.array([0.1, 0.8, 0.3, 0.5])  # Q(s, \cdot)

# Training: epsilon-greedy with \epsilon=0.1 (10% random, 90% greedy)
train_action = rlax.epsilon_greedy(epsilon=0.1).sample(key, q_values)
print(train_action)  # integer action index

# Evaluation: pure greedy
eval_action = rlax.greedy().sample(key, q_values)
print(eval_action)  # 1 (argmax)

# Annealing epsilon over training (common DQN pattern):
import optax
epsilon_schedule = optax.polynomial_schedule(
    init_value=1.0, end_value=0.01,
    transition_steps=10000, power=1.0)
step = jnp.int32(5000)
epsilon = epsilon_schedule(step)
action = rlax.epsilon_greedy(epsilon).sample(key, q_values)
```
```

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

### N Step Bootstrapped Returns

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Calculates N-step bootstrapped returns, which combine Monte Carlo and TD learning.

```APIDOC
## N Step Bootstrapped Returns

### Description
Calculates N-step bootstrapped returns.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### Transformed N Step Q Learning

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements transformed N-step Q Learning, applying transformations to N-step Q learning updates.

```APIDOC
## Transformed N Step Q Learning

### Description
Implements transformed N-step Q Learning.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### rlax.n_step_bootstrapped_returns

Source: https://github.com/google-deepmind/rlax/blob/main/docs/api.md

Computes strided n-step bootstrapped return targets over a sequence. The returns are computed according to the equation Gₜ = rₜ₊₁ + γₜ₊₁ [(1 - λₜ₊₁) vₜ₊₁ + λₜ₊₁ Gₜ₊₁] iterated n times.

```APIDOC
## rlax.n_step_bootstrapped_returns

### Description
Computes strided n-step bootstrapped return targets over a sequence.

### Parameters
* **r_t** (Array | ndarray | bool | number) - rewards at times [1, …, T].
* **discount_t** (Array | ndarray | bool | number) - discounts at times [1, …, T].
* **v_t** (Array | ndarray | bool | number) - state or state-action values to bootstrap from at time [1, …., T].
* **n** (int) - number of steps over which to accumulate reward before bootstrapping.
* **lambda_t** (Array | ndarray | bool | number | float | int, optional) - lambdas at times [1, …, T]. Shape is [], or [T-1]. Defaults to 1.0.
* **stop_target_gradients** (bool, optional) - bool indicating whether or not to apply stop gradient to targets. Defaults to False.

### Returns
Array | ndarray | bool | number - estimated bootstrapped returns at times [0, …., T-1]
```

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

### rlax.compute_parametric_kl_penalty_and_dual_loss

Source: https://github.com/google-deepmind/rlax/blob/main/docs/api.md

Optimizes hard KL constraints between the current and previous policies using parametric methods.

```APIDOC
## rlax.compute_parametric_kl_penalty_and_dual_loss

### Description
Optimize hard KL constraints between the current and previous policies.

### Parameters
* **kl_constraints** (Sequence[Tuple[Array | ndarray | bool | number, LagrangePenalty]]) - Sequence of KL constraints.
* **projection_operator** (Callable[[Array | ndarray | bool | number], Array | ndarray | bool | number]) - Function to project values.
* **use_stop_gradient** (bool, optional) - Indicates whether or not to apply stop gradient. Defaults to True.

### Returns
Tuple containing the KL penalty and dual loss.
```

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

### General Off Policy Returns From Q and V

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Calculates general off-policy returns using both Q-values and V-values.

```APIDOC
## General Off Policy Returns From Q and V

### Description
Provides general off-policy returns from Q and V values.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### Persistent Q Learning

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements Persistent Q Learning, a variant of Q-learning that aims for more stable learning.

```APIDOC
## Persistent Q Learning

### Description
Implements Persistent Q Learning.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### Compute n-step Returns with rlax.n_step_bootstrapped_returns

Source: https://context7.com/google-deepmind/rlax/llms.txt

Calculates strided n-step bootstrapped returns. This function accumulates n rewards before bootstrapping from the value function, suitable for n-step learning methods.

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

r_t        = jnp.array([1.0, 0.5, 0.8, -0.5, 0.2, 1.0])
discount_t = jnp.array([0.99]*6)
v_t        = jnp.array([0.5, 0.8, 0.6, 0.4, 0.3, 0.0])
n = 3  # 3-step returns

returns = rlax.n_step_bootstrapped_returns(r_t, discount_t, v_t, n)
print(returns.shape)  # (6,)
# returns[0] = r[0] + γ*r[1] + γ²*r[2] + γ³*v[3]
```

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

### Categorical Q Learning

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements Categorical Q Learning, a fundamental algorithm for learning action-value functions.

```APIDOC
## Categorical Q Learning

### Description
Implements Categorical Q Learning.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### rlax.zero_policy_targets

Source: https://github.com/google-deepmind/rlax/blob/main/docs/api.md

Create policy targets with constant weights. Actions will be sampled from the distribution and associated with a constant weights_scale.

```APIDOC
## rlax.zero_policy_targets

### Description
Create policy targets with constant weights. The actions will be sampled from distribution and will all be associated to a constant weights_scale. If distribution is a (discrete or continuous) uniform probability distribution, distilling these targets will push the agent’s policy towards a uniform distribution. The strength of the penalty associated with a non-uniform policy depends on weights_scale (e.g. in the extreme case weights_scale==0, distillation loss is 0 for any policy).

### Parameters
* **distribution** (Distribution | Distribution) – a distrax or tfp distribution for sampling actions.
* **rng_key** (Array) – a JAX pseudo random number generator.
* **num_samples** (int) – number of action sampled in the PolicyTarget.
* **weights_scale** (float) – constant weight associated with each action. Defaults to 0.0.

### Returns
a PolicyTarget to learn from (PolicyTarget).
```

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

### Retrace Continuous

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements a continuous version of the Retrace algorithm.

```APIDOC
## Retrace Continuous

### Description
Implements Retrace Continuous.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### rlax.general_off_policy_returns_from_q_and_v

Source: https://github.com/google-deepmind/rlax/blob/main/docs/api.md

Calculates targets for various off-policy evaluation algorithms using Q and V values.

```APIDOC
## rlax.general_off_policy_returns_from_q_and_v

### Description
Calculates targets for various off-policy evaluation algorithms.

Given a window of experience of length K+1, generated by a behaviour policy
μ, for each time-step t we can estimate the return G_t from that step
onwards, under some target policy π, using the rewards in the trajectory,
the values under π of states and actions selected by μ, according to equation:

> Gₜ = rₜ₊₁ + γₜ₊₁ * (vₜ₊₁ - cₜ₊₁ * q(aₜ₊₁) + cₜ₊₁* Gₜ₊₁),

where, depending on the choice of c_t, the algorithm implements:

> Importance Sampling             c_t = π(x_t, a_t) / μ(x_t, a_t),
> Harutyunyan’s et al. Q(lambda)  c_t = λ,
> Precup’s et al. Tree-Backup     c_t = π(x_t, a_t),
> Munos’ et al. Retrace           c_t = λ min(1, π(x_t, a_t) / μ(x_t, a_t)).

See “Safe and Efficient Off-Policy Reinforcement Learning” by Munos et al.
([https://arxiv.org/abs/1606.02647](https://arxiv.org/abs/1606.02647)).

### Parameters
* **q_t** – Q-values under π of actions executed by μ at times [1, …, K - 1].
* **v_t** – Values under π at times [1, …, K].
* **r_t** – rewards at times [1, …, K].
* **discount_t** – discounts at times [1, …, K].
* **c_t** – weights at times [1, …, K - 1].
* **stop_target_gradients** – bool indicating whether or not to apply stop gradient
  to targets.

### Returns
Off-policy estimates of the generalized returns from states visited at times
[0, …, K - 1].
```

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

### rlax.constant_policy_targets

Source: https://github.com/google-deepmind/rlax/blob/main/docs/api.md

Create policy targets with constant weights. Actions will be sampled from the distribution and associated with a constant weights_scale.

```APIDOC
## rlax.constant_policy_targets

### Description
Create policy targets with constant weights. The actions will be sampled from distribution and will all be associated to a constant weights_scale. If distribution is a (discrete or continuous) uniform probability distribution, distilling these targets will push the agent’s policy towards a uniform distribution. The strength of the penalty associated with a non-uniform policy depends on weights_scale (e.g. in the extreme case weights_scale==0, distillation loss is 0 for any policy).

### Parameters
* **distribution** (Distribution | Distribution) – a distrax or tfp distribution for sampling actions.
* **rng_key** (Array) – a JAX pseudo random number generator.
* **num_samples** (int) – number of action sampled in the PolicyTarget.
* **weights_scale** (float) – constant weight associated with each action. Defaults to 1.0.

### Returns
a PolicyTarget to learn from (PolicyTarget).
```

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

### MPO Compute Weights and Temperature Loss

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Computes weights and temperature loss for the Maximum a Posteriori Policy Optimization (MPO) algorithm.

```APIDOC
## MPO Compute Weights and Temperature Loss

### Description
Computes weights and temperature loss for MPO.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### Softmax Action Sampling

Source: https://context7.com/google-deepmind/rlax/llms.txt

Samples actions using a temperature-scaled softmax distribution. Higher temperatures lead to more uniform sampling, while lower temperatures favor actions with higher logits.

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

key = jax.random.PRNGKey(0)
logits = jnp.array([1.0, 2.0, 0.5, 1.5])

# Default temperature=1.0
action = rlax.softmax(temperature=1.0).sample(key, logits)

# Higher temperature = more uniform; lower = more greedy
action_hot  = rlax.softmax(temperature=0.1).sample(key, logits)   # near-greedy
action_cold = rlax.softmax(temperature=10.0).sample(key, logits)  # near-uniform
```

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

### Transformed Q Lambda

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements transformed Q Lambda, applying transformations to the Q-Lambda algorithm.

```APIDOC
## Transformed Q Lambda

### Description
Implements transformed Q Lambda.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### Learner Step Function

Source: https://context7.com/google-deepmind/rlax/llms.txt

Performs a single step of the learner, computing gradients, updating optimizer state, and applying updates to the online network parameters. The target network is updated periodically.

```python
@jax.jit
def learner_step(params, data, learner_state):
    grads = jax.grad(loss_fn)(params.online, params.target, *data)
    updates, opt_state = optimizer.update(grads, learner_state.opt_state)
    online_params = optax.apply_updates(params.online, updates)
    # Sync target network every 50 steps
    target_params = rlax.periodic_update(
        online_params, params.target, learner_state.count, 50)
    return Params(online_params, target_params), \
           LearnerState(learner_state.count + 1, opt_state)
```

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

### rlax.popart

Source: https://context7.com/google-deepmind/rlax/llms.txt

Implements Pop-Art (Preserving Outputs Precisely, while Adaptively Rescaling Targets) for adaptive reward normalization.

```APIDOC
## rlax.popart

### Description
Implements Pop-Art (Preserving Outputs Precisely, while Adaptively Rescaling Targets), which normalizes value function targets online while preserving the network's output via inverse weight updates. This is crucial for multi-task RL or environments with highly varying reward scales.

### Parameters
- **num_outputs** (int) - The number of output heads (e.g., tasks).
- **step_size** (float, optional) - The step size for updating the Pop-Art state. Defaults to 1e-4.
- **scale_lb** (float, optional) - Lower bound for the scale. Defaults to 1e-4.
- **scale_ub** (float, optional) - Upper bound for the scale. Defaults to 1e6.

### Returns
- **init_state_fn** (function) - Function to initialize the Pop-Art state.
- **popart_update_fn** (function) - Function to perform a Pop-Art update step.
```

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

### Categorical L2 Project

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Provides functionality for Categorical L2 Projection, used in value function approximation.

```APIDOC
## Categorical L2 Project

### Description
Provides Categorical L2 Projection.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### rlax.clipped_surrogate_pg_loss

Source: https://context7.com/google-deepmind/rlax/llms.txt

Computes the PPO clipped objective L_clip = -E[min(r_t·Â_t, clip(r_t, 1-ε, 1+ε)·Â_t)], preventing excessively large policy updates.

```APIDOC
## rlax.clipped_surrogate_pg_loss — PPO clipped surrogate loss

Computes the PPO clipped objective `L_clip = -E[min(r_t·Â_t, clip(r_t, 1-ε, 1+ε)·Â_t)]`, preventing excessively large policy updates.

### Parameters

- **prob_ratios_t** (jax.numpy.ndarray) - Ratio of new policy probs to old policy probs: π_new(a|s) / π_old(a|s)
- **adv_t** (jax.numpy.ndarray) - 
- **epsilon** (float) - PPO clip parameter

### Request Example

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

T = 5
# Ratio of new policy probs to old policy probs: π_new(a|s) / π_old(a|s)
prob_ratios_t = jnp.array([1.1, 0.9, 1.3, 0.7, 1.05])
adv_t         = jnp.array([0.5, -0.3, 1.2, -0.1, 0.8])
epsilon       = 0.2  # PPO clip parameter

loss = rlax.clipped_surrogate_pg_loss(prob_ratios_t, adv_t, epsilon)
print(loss)  # scalar PPO clip loss

# In a full PPO update:
def ppo_loss(logits_new, logits_old, actions, advantages, epsilon=0.2):
    log_probs_new = jax.nn.log_softmax(logits_new)
    log_probs_old = jax.nn.log_softmax(logits_old)
    log_ratios = log_probs_new[jnp.arange(len(actions)), actions] \
               - log_probs_old[jnp.arange(len(actions)), actions]
    ratios = jnp.exp(log_ratios)
    return rlax.clipped_surrogate_pg_loss(ratios, advantages, epsilon)
```

### Response

#### Success Response (200)

- **loss** (scalar) - scalar PPO clip loss

#### Response Example

```python
# scalar loss value
```
```

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

### Quantile Expected Sarsa

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements Quantile Expected SARSA, a method for learning distributional value functions.

```APIDOC
## Quantile Expected Sarsa

### Description
Implements Quantile Expected Sarsa.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### Retrace

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements Retrace, an off-policy correction algorithm for estimating action values.

```APIDOC
## Retrace

### Description
Implements Retrace.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### Expected SARSA

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements Expected SARSA, an on-policy temporal difference learning algorithm.

```APIDOC
## Expected SARSA

### Description
Implements Expected SARSA.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### VTrace

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements VTrace, an off-policy correction algorithm for policy gradient methods.

```APIDOC
## VTrace

### Description
Implements VTrace.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### rlax.softmax

Source: https://context7.com/google-deepmind/rlax/llms.txt

Provides a temperature-scaled softmax distribution for sampling actions, useful for continuous exploration proportional to value estimates.

```APIDOC
## rlax.softmax — Softmax action sampling

### Description
Provides a temperature-scaled softmax distribution for sampling actions, useful for continuous exploration proportional to value estimates.

### Usage
```python
import jax
import jax.numpy as jnp
import rlax

key = jax.random.PRNGKey(0)
logits = jnp.array([1.0, 2.0, 0.5, 1.5])

# Default temperature=1.0
action = rlax.softmax(temperature=1.0).sample(key, logits)

# Higher temperature = more uniform; lower = more greedy
action_hot  = rlax.softmax(temperature=0.1).sample(key, logits)   # near-greedy
action_cold = rlax.softmax(temperature=10.0).sample(key, logits)  # near-uniform
```
```

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

### rlax.vmpo_compute_weights_and_temperature_loss

Source: https://github.com/google-deepmind/rlax/blob/main/docs/api.md

Computes the weights and temperature loss for V-MPO. This function is part of the V-MPO algorithm's E-step.

```APIDOC
## rlax.vmpo_compute_weights_and_temperature_loss

### Description
Computes the weights and temperature loss for V-MPO.

### Parameters
* **advantages** (Array | ndarray | bool | number) - Advantages for the E-step. Shape E*.
* **restarting_weights** (Array | ndarray | bool | number) - Restarting weights, 0 means that this step is the start of a new episode and we ignore losses at this step because the agent cannot influence these. Shape E*.
* **importance_weights** (Array | ndarray | bool | number) - Optional importance weights. Shape E*.
* **temperature_constraint** (LagrangePenalty) - Lagrange constraint for the E-step temperature optimization.
* **projection_operator** (Callable[[Array | ndarray | bool | number], Array | ndarray | bool | number]) - Function to project dual variables (temperature and kl constraint alphas) into the positive range.
* **top_k_fraction** (float) - Fraction of samples to use in the E-step.
* **axis_name** (str | None, optional) - Optional axis name for pmap or ‘vmap’. If None, computations are performed locally on each device.
* **use_stop_gradient** (bool, optional) - bool indicating whether or not to apply stop gradient. Defaults to True.

### Returns
Tuple[Array | ndarray | bool | number, Array | ndarray | bool | number, Array | ndarray | bool | number] - The temperature loss, normalized weights and number of samples used.
```

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

### Compute Multi-step λ-Return Targets with rlax.lambda_returns

Source: https://context7.com/google-deepmind/rlax/llms.txt

Computes bootstrapped λ-return targets recursively from the end of a trajectory. This is useful for calculating target values in methods that use eligibility traces.

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

r_t        = jnp.array([1.0, 0.5, -1.0, 0.0])
discount_t = jnp.array([0.99, 0.99, 0.99, 0.0])  # terminal at last step
v_t        = jnp.array([1.0, 0.8, 0.6, 0.0])
lambda_    = 0.9

returns = rlax.lambda_returns(r_t, discount_t, v_t, lambda_)
print(returns)  # shape (4,), G_t targets for each timestep
# returns[3] = 0.0 + 0.0*... = 0.0 (terminal)
# returns[2] = -1.0 + 0.99*(0.1*0.6 + 0.9*returns[3])
```

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

### Calculate Importance Sampling Ratios

Source: https://context7.com/google-deepmind/rlax/llms.txt

Computes importance sampling ratios for categorical policies. Used as importance weights in algorithms like V-trace or retrace.

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

pi_logits_t = jnp.array([[1.0, 2.0, 0.5], [0.5, 1.5, 1.0]])  # shape [T, A]
mu_logits_t = jnp.array([[0.8, 1.9, 0.6], [0.4, 1.4, 1.2]])
a_t         = jnp.array([1, 2])  # actions taken by behavior policy

rhos = rlax.categorical_importance_sampling_ratios(pi_logits_t, mu_logits_t, a_t)
print(rhos.shape)  # (2,)  — one ratio per timestep
```

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

### Truncated Generalized Advantage Estimation

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Implements Truncated Generalized Advantage Estimation (GAE), a method for estimating advantage functions.

```APIDOC
## Truncated Generalized Advantage Estimation

### Description
Implements Truncated Generalized Advantage Estimation.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```

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

### Compute Discounted Returns

Source: https://context7.com/google-deepmind/rlax/llms.txt

Computes the discounted cumulative return G_t = r_{t+1} + γ_{t+1} G_{t+1} from a complete trajectory. Used for full trajectory returns without bootstrapping.

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

r_t        = jnp.array([1.0, 0.5, -1.0, 2.0])
discount_t = jnp.array([0.99, 0.99, 0.99, 0.0])  # episode ends
v_t        = jnp.array([0.0]*4)  # no bootstrapping (terminal)

returns = rlax.discounted_returns(r_t, discount_t, v_t)
print(returns)
# returns[0] = 1.0 + 0.99*(0.5 + 0.99*(-1.0 + 0.99*2.0))
```

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

### Discounted Returns

Source: https://github.com/google-deepmind/rlax/blob/main/docs/index.md

Calculates discounted returns, a core concept in reinforcement learning for evaluating sequences of rewards.

```APIDOC
## Discounted Returns

### Description
Calculates discounted returns.

### Method
Not specified (likely a function call within the RLAX library).

### Endpoint
Not applicable (SDK/library function).

### Parameters
Not specified in the source.

### Request Example
Not applicable.

### Response
Not specified.
```