DOCSEARCHbyAlgolia
Submit Docs

RLax

repository·main·Indexed 23 days ago

https://github.com/google-deepmind/rlax

A library of JAX-based building blocks for reinforcement learning agents. RLax provides implementations of RL-specific mathematical operations rather than complete algorithms, including value learning (Q-learning, SARSA, V-trace), policy optimization losses (PPO-style, MPO, DPG), exploration utilities (Gaussian and Dirichlet noise), and distributional RL components. It supports both on-policy and off-policy learning and is designed for just-in-time compilation via jax.jit for CPU, GPU, and TPU hardware.

Tokens
1.8K
Snippets
0
Records
13
Agent score
81%

What's inside rlax

  1. What RLax provides

    main

    RLax is a library of JAX-based building blocks for reinforcement learning agents. It does not provide complete algorithms, but rather implementations of RL-specific mathematical operations.

    Key capabilities include:

    • Values: State and action-values.
    • Bellman Equations: Values for non-linear generalizations of the Bellman equations.
    • Return Distributions: Distributional value functions.
    • General Value Functions: Values for cumulants other than the main reward.
    • Policies: Policy-gradients for both continuous and discrete action spaces.

    The library supports both on-policy and off-policy learning.

  2. Explore RLax API categories

    main

    RLax provides a collection of reinforcement learning primitives organized into several functional categories:

    • Value Learning: Functions for computing returns and TD errors (e.g., q_learning, vtrace, retrace, sarsa).
    • Policy Optimization: Loss functions for policy gradients and optimization (e.g., policy_gradient_loss, mpo_loss, clipped_surrogate_pg_loss).
    • Exploration: Noise injection and intrinsic reward utilities (e.g., add_gaussian_noise, episodic_memory_intrinsic_rewards).
    • Utilities: General RL helpers for gradients, trees, and broadcasting (e.g., clip_gradient, tree_map_zipped).
    • General Value Functions: Rewards for control tasks (e.g., pixel_control_rewards).
    • Model Learning: Tools for extracting training data from sequences.
    • Pop Art: Normalization techniques for value functions.
    • Transforms: Mathematical transforms like logit, sigmoid, and twohot mappings.
    • Losses: Standard RL losses (e.g., huber_loss, l2_loss).
    • Distributions: Sampling and divergence utilities for various distributions (e.g., categorical_sample, squashed_gaussian).

  3. Apply RLax functions to batches using vmap

    main
    RLax functions are defined for agents interacting with a single stream of experience. To support parallel data generation or replay buffers (batch processing), use the JAX vmap construct to apply these functions to batches of data.

  4. Understand RLax naming conventions for timesteps

    main

    RLax functions often operate on sequences of data (policies, actions, rewards, values) across consecutive timesteps. To clarify which timestep an input refers to, the library uses specific suffixes:

    • _tm1: Refers to the timestep before the current one (e.g., a_tm1 is the action selected in the source state).
    • _t: Refers to the current timestep (e.g., r_t is the reward collected in the destination state).

    Common examples include:

    • q_tm1: The action value in the source state of a transition.
    • a_tm1: The action that was selected in the source state.
    • r_t: The resulting rewards collected in the destination state.
    • discount_t: The discount associated with a transition.
    • q_t: The action values in the destination state.

  5. Install RLax

    main

    You can install the latest released version of RLax from PyPI or the latest development version from GitHub.

    To run the provided examples/, you will also need to install optax, haiku, and bsuite.

    All RLax code can be just-in-time compiled for different hardware (CPU, GPU, TPU) using jax.jit.

  6. Explore the RLax public API

    main

    RLax provides a collection of reinforcement learning building blocks implemented in JAX. The public API is exposed through the top-level rlax module.

    Important Note: Do not use symbols located in the _src subpackages directly, as they are internal implementation details and are not part of the stable public API. Always import from the top-level rlax module.

    The library includes modules for:

    • Distributions: Categorical and Gaussian distributions, epsilon-greedy, and softmax operations.
    • Losses: Expectile loss, L2 loss, and various policy gradient losses.
    • Value Learning: Q-learning, SARSA, TD-learning, and quantile regression.
    • Multistep Returns: Discounted returns, n-step returns, and GAE.
    • Exploration: Gaussian, Dirichlet, and Ornstein-Uhlenbeck noise.
    • PopArt: Normalization and unnormalization utilities.
    • VTrace: V-trace and leaky V-trace implementations.

  7. Use Value Learning primitives

    main

    RLax offers a wide range of functions for value-based reinforcement learning, including temporal difference (TD) learning, n-step returns, and distributional RL. Key functions include:

    • Standard TD/SARSA: td_learning, sarsa, expected_sarsa.
    • N-Step & Lambda: td_lambda, sarsa_lambda, q_lambda, n_step_bootstrapped_returns.
    • Off-Policy Correction: vtrace, leaky_vtrace, retrace, retrace_continuous.
    • Distributional RL: categorical_q_learning, quantile_q_learning.
    • Advanced Returns: discounted_returns, lambda_returns, general_off_policy_returns_from_q_and_v.

  8. Use RLax Transforms

    main

    RLax provides mathematical transforms often used in RL, such as mapping values to probabilities or handling discrete representations:

    • Standard Transforms: sigmoid, logit, power, identity.
    • Signed Transforms: signed_expm1, signed_hyperbolic, signed_logp1, signed_parabolic.
    • Two-Hot Encoding: transform_to_2hot and transform_from_2hot for representing continuous values in discrete bins.

  9. Use Exploration utilities

    main

    To encourage exploration in RL agents, RLax provides noise injection and intrinsic reward functions:

    • Noise Injection:
      • add_gaussian_noise
      • add_dirichlet_noise
      • add_ornstein_uhlenbeck_noise
    • Intrinsic Rewards: episodic_memory_intrinsic_rewards for memory-based exploration.
    • Similarity Search: knn_query for k-nearest neighbor lookups.

  10. Use RLax Utilities

    main

    RLax includes several utility functions for common RL engineering tasks:

    • Gradient & Optimization: clip_gradient, create_ema (Exponential Moving Average), periodic_update.
    • Tree Operations: tree_map_zipped, tree_select, tree_split_key, tree_split_leaves.
    • Tensor Manipulation: batched_index, lhs_broadcast, one_hot, transpose_first_axis_to_last, transpose_last_axis_to_first.
    • Masking: replace_masked, tree_replace_masked.
    • Other: AllSum class for summation utilities.

  11. Use Pop Art for value normalization

    main

    Pop Art (Preserving Outputs of Policy via Art) is a technique used to normalize value targets. RLax provides the following components:

    • popart: High-level interface.
    • art: Core logic.
    • normalize and unnormalize: For scaling values.
    • unnormalize_linear: For linear unnormalization.

  12. Use Policy Optimization losses

    main

    RLax provides loss functions for various policy optimization algorithms. Common functions include:

    • Policy Gradients: policy_gradient_loss, clipped_surrogate_pg_loss (for PPO-style updates).
    • Actor-Critic/DPG: dpg_loss (Deterministic Policy Gradient).
    • Maximum A Posteriori Policy Optimization (MPO): mpo_loss and mpo_compute_weights_and_temperature_loss.
    • Entropy & Regularization: entropy_loss, rm_loss, rpg_loss.
    • Distillation: sampled_policy_distillation_loss.