RLlib scaling guide

RLlib is a distributed and scalable RL library, based on Ray. An RLlib Algorithm uses Ray actors wherever parallelization of its sub-components can speed up sample and learning throughput.

Scalable axes in RLlib: All RLlib Algorithm classes provide three scaling axes:

• The number of EnvRunner actors in the EnvRunnerGroup, settable through config.env_runners(num_env_runners=n).

• The number of vectorized sub-environments on each EnvRunner actor, settable through config.env_runners(num_envs_per_env_runner=p).

• The number of Learner actors in the LearnerGroup, settable through config.learners(num_learners=m).

Scaling the number of EnvRunner actors

Control the degree of parallelism of the sampling machinery of the Algorithm by increasing the number of remote EnvRunner actors in the EnvRunnerGroup through the config:

from ray.rllib.algorithms.ppo import PPOConfig

config = (
    PPOConfig()
    # Use 4 EnvRunner actors (default is 2).
    .env_runners(num_env_runners=4)
)

To assign resources to each EnvRunner, use these config settings:

config.env_runners(
    num_cpus_per_env_runner=...,
    num_gpus_per_env_runner=...,
)

See this example of an EnvRunner and RL environment requiring a GPU resource.

The number of GPUs can be fractional, for example 0.5, to allocate only a fraction of a GPU per EnvRunner.

There’s always one “local” EnvRunner in the EnvRunnerGroup. To sample using only this local EnvRunner, set num_env_runners=0. This local EnvRunner sits directly in the main Algorithm process.

Hint

The Ray team might deprecate the local EnvRunner. It exists for historical reasons, and whether to keep it in the set is under debate.

Scaling the number of envs per EnvRunner actor

RLlib vectorizes RL environments on EnvRunner actors through gymnasium’s VectorEnv API. To create more than one environment copy per EnvRunner, set the following in your config:

from ray.rllib.algorithms.ppo import PPOConfig

config = (
    PPOConfig()
    # Use 10 sub-environments (vector) per EnvRunner.
    .env_runners(num_envs_per_env_runner=10)
)

Note

Unlike single-agent environments, RLlib can’t vectorize multi-agent setups yet. The Ray team is working on a solution that uses the gymnasium >= 1.x custom vectorization feature.

With more than one environment per EnvRunner, the RLModule runs inference on a batch of data and computes actions for all sub-environments in parallel.

By default, the individual sub-environments in a vector step and reset in sequence. Only the action computation of the RL environment loop runs in parallel, because observations can move through the model in a batch. However, gymnasium supports an asynchronous vectorization setting that gives each sub-environment its own Python process. The vector environment can then step or reset in parallel. Activate this asynchronous vectorization through:

import gymnasium as gym

config.env_runners(
    gym_env_vectorize_mode=gym.envs.registration.VectorizeMode.ASYNC,  # default is `SYNC`
)

This setting can significantly speed up sampling when combined with num_envs_per_env_runner > 1, especially when your RL environment’s stepping process is time-consuming.

See this example script that demonstrates a large speedup with async vectorization.

Scaling the number of Learner actors

Learning updates happen in the LearnerGroup, which manages either a single local Learner instance or any number of remote Learner actors.

Set the number of remote Learner actors through:

from ray.rllib.algorithms.ppo import PPOConfig

config = (
    PPOConfig()
    # Use 2 remote Learner actors (default is 0) for distributed data parallelism.
    # Choosing 0 creates a local Learner instance on the main Algorithm process.
    .learners(num_learners=2)
)

Typically, you use as many Learner actors as you have GPUs available for training. Set the number of GPUs per Learner to 1:

config.learners(num_gpus_per_learner=1)

Warning

For some algorithms, such as IMPALA and APPO, the performance of a single remote Learner actor with num_learners=1 compared to a single local Learner instance with num_learners=0 depends on whether a GPU is available. With exactly one GPU, run these two algorithms with num_learners=0, num_gpus_per_learner=1. With no GPU, set num_learners=1, num_gpus_per_learner=0. With more than one GPU, set num_learners=..., num_gpus_per_learner=1.

The number of GPUs can be fractional, for example 0.5, to allocate only a fraction of a GPU per Learner. For example, pack five Learner instances onto one GPU by setting num_learners=1, num_gpus_per_learner=0.2. See this fractional GPU example for details.

Note

If you specify num_gpus_per_learner > 0 and your machine doesn’t have enough GPUs, the experiment might stall until the Ray autoscaler brings up enough machines to fulfill the resource request. If your cluster has autoscaling turned off, this setting results in a seemingly hanging experiment run.

If you set num_gpus_per_learner=0, RLlib builds the RLModule instances on CPUs only, even if GPUs are available on the cluster.

Outlook: More RLlib elements that should scale

Other components and aspects of RLlib should also scale up.

For example, RLlib scales Learner actors only through “distributed data parallel” (DDP), so the model size is limited to whatever fits on a single GPU.

The Ray team is working on closing these gaps. Future areas of improvement include:

• Enable training large models, such as a “large language model” (LLM). The team is working on a “Reinforcement Learning from Human Feedback” (RLHF) prototype setup. The main problems to solve are the model-parallel and tensor-parallel distribution across multiple GPUs, and a reasonably fast transfer of weights between Ray actors.

• Enable training with thousands of multi-agent policies. A possible solution for this scaling problem is to split the MultiRLModule into manageable groups of individual policies across the EnvRunner and Learner actors.

• Enable vector envs for multi-agent.