On-Policy Rollout Functions

These helpers gather transitions from an environment using an on-policy agent. They fill the agent’s rollout buffer so that calling agent.learn() will update the policy from the collected data.

agilerl.rollouts.collect_rollouts(agent: PPO, env: Env | VectorEnv, n_steps: int | None = None, **kwargs: Any) → tuple[list[float], ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]], dict[str, Any]]

Collect rollouts for non-recurrent on-policy algorithms.

Parameters:

  • agent (RLAlgorithm) – The agent to collect rollouts for.

  • env (gym.Env | gym.vector.VectorEnv) – The environment to collect rollouts from.

  • n_steps (int | None) – The number of steps to collect rollouts for.

Returns:

The scores for the episodes completed in the rollouts, followed by the observation, done flag, scores, and info for the current step.

Return type:

tuple[list[float], npt.NDArray, npt.NDArray, npt.NDArray, dict[str, Any]]

agilerl.rollouts.collect_rollouts_recurrent(agent: PPO, env: Env | VectorEnv, n_steps: int | None = None, **kwargs: Any) → tuple[list[float], ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]], dict[str, Any]]

Collect rollouts for recurrent on-policy algorithms.

Parameters:

  • agent (RLAlgorithm) – The agent to collect rollouts for.

  • env (gym.Env | gym.vector.VectorEnv) – The environment to collect rollouts from.

  • n_steps (int | None) – The number of steps to collect rollouts for.

Returns:

The scores for the episodes completed in the rollouts, followed by the observation, done flag, scores, and info for the current step.

Return type:

tuple[list[float], npt.NDArray, npt.NDArray, npt.NDArray, dict[str, Any]]

Example

Using a non-recurrent PPO agent:

import gymnasium as gym
from agilerl.algorithms import PPO
from agilerl.rollouts import collect_rollouts

env = gym.make("CartPole-v1")
agent = PPO(env.observation_space, env.action_space)

collect_rollouts(agent, env, n_steps=agent.learn_step)
agent.learn()

For recurrent policies, use collect_rollouts_recurrent:

num_envs = 4
env = gym.vector.SyncVectorEnv([lambda: gym.make("CartPole-v1")] * num_envs)
agent = PPO(
    env.single_observation_space,
    env.single_action_space,
    recurrent=True,
    num_envs=num_envs,
)

collect_rollouts_recurrent(agent, env, n_steps=5)
agent.learn()