2️⃣ Rollouts

Rollout Buffer

The way that our implementation of VPG will work is simple: we perform a rollout across num_envs many environments in parallel, and store the trajectories for each. We then learn from that set of rollouts, and then discard it afterwards. One rollout, one learning step. This means we are always learning on-policy: we only ever learn from data that the current model actually generated. We will use a rollout buffer to store the trajectories.

Exercise - implement Rollout Buffer

Difficulty: 🔴🔴🔴⚪⚪
Importance: 🔵🔵🔵⚪⚪
You should spend up to 20 minutes on this exercise.

The Rollout class will store a set of num_envs many trajectories. We do not shuffle up anything, or break up an episode into little experiences as we did for DQN. The smallest datapoint is one full trajectory:

$$\tau = s_0 \; a_0 \; r_0 \; s_1 \; a_1 \; r_1 \ldots s_T \; a_T \; r_T$$

The following methods need to be completed:

  • add_step - adds information gathered from timestep $t$ to the rollout buffer
  • get_batches - returns a list of RolloutTensors objects, each containing batch_size many trajectories.

We store the tensors for each step as separate lists, and then stack once at the end to get the final tensors with the .get function. This ends up being cheaper as it avoids spinning up indexed-write kernels per step.

Hint

Use t.split to write get_batches.

RolloutTensors = namedtuple("RolloutTensors", ["obs", "actions", "logprobs", "rewards", "dones"])


class Rollout:
    _obs: list[Float[Tensor, " num_envs *obs_shape"]]
    _actions: list[Int[Tensor, " num_envs"]]
    _logprobs: list[Float[Tensor, " num_envs"]]
    _rewards: list[Float[Tensor, " num_envs"]]
    _dones: list[Bool[Tensor, " num_envs"]]
    timestep: int

    def __init__(self, max_steps: int):
        """
        Args:
            max_steps: maximum number of steps to rollout per environment (the per-step tensors carry the
                       number of environments and the observation/action shapes themselves)
        """

        self.MAX_SIZE = max_steps

        # Per-step we append tensor references to Python lists (free) and t.stack() once at the
        # end, instead of 5 indexed-write kernels per step into a preallocated buffer. Each stored
        # tensor is freshly produced per step (the env returns a new state tensor each step), so
        # holding references is safe. This removes ~2500 tiny kernel launches per full rollout.
        self._obs, self._actions, self._logprobs, self._rewards, self._dones = [], [], [], [], []
        self.timestep = 0

    def add_step(
        self,
        obs: Float[Tensor, " num_envs *obs_shape"],
        actions: Int[Tensor, " num_envs"],
        logprobs: Float[Tensor, " num_envs"],
        rewards: Float[Tensor, " num_envs"],
        dones: Bool[Tensor, " num_envs"],
    ):
        """
        Stores one timestep of the rollout (the same transition for all `num_envs` environments at once)
        and advances `self.timestep`.

        Args:
            obs:      (num_envs, *obs_shape)  the observation (o_t) the agent saw *before* acting
            actions:  (num_envs,)             the action (a_t) it took in each env after seeing o_t
            logprobs: (num_envs,)             log pi(a_t | o_t) of the action that was *taken* 
                                              Recorded at collection time, so
                                              these are the behaviour policy's logprobs (used for importance weights)
            rewards:  (num_envs,)             the reward (r_{t+1}) received for taking a_t
            dones:    (num_envs,)  bool       True if that env's episode ended on this step (d_{t+1} = terminated OR truncated),
                                              i.e. the env auto-reset and the *next* stored obs starts a new episode

        Raises ValueError if the rollout already holds `max_steps` timesteps.
        """

        if self.timestep >= self.MAX_SIZE:
            raise ValueError("Rollout is full, cannot add more steps")

        raise NotImplementedError()

    def reset(self):
        self._obs.clear(); self._actions.clear(); self._logprobs.clear()
        self._rewards.clear(); self._dones.clear()
        self.timestep = 0

    def get(self) -> tuple[Tensor, ...]:
        """
        Stack the per-step lists from (num_envs, *tensor_shape) into (num_envs, timestep, *tensor_shape) tensors. Rollouts can stop early
        (see gen_rollout), so the time dimension is however many steps were actually collected.
        """
        assert self.timestep > 0, "Rollout is empty"
        return RolloutTensors(
            t.stack(self._obs, dim=1),
            t.stack(self._actions, dim=1),
            t.stack(self._logprobs, dim=1),
            t.stack(self._rewards, dim=1).float(),
            t.stack(self._dones, dim=1),
        )

    def get_batches(self, batch_size: int) -> list[RolloutTensors]:
        """
        Splits the rollout buffer into batches of size `batch_size`, and returns a list of
        `RolloutTensors` objects, each containing `batch_size` many trajectories.
        """

        raise NotImplementedError()


tests.test_rollout(Rollout)
Solution 
RolloutTensors = namedtuple("RolloutTensors", ["obs", "actions", "logprobs", "rewards", "dones"])


class Rollout:
    _obs: list[Float[Tensor, " num_envs *obs_shape"]]
    _actions: list[Int[Tensor, " num_envs"]]
    _logprobs: list[Float[Tensor, " num_envs"]]
    _rewards: list[Float[Tensor, " num_envs"]]
    _dones: list[Bool[Tensor, " num_envs"]]
    timestep: int

    def __init__(self, max_steps: int):
        """
        Args:
            max_steps: maximum number of steps to rollout per environment (the per-step tensors carry the
                       number of environments and the observation/action shapes themselves)
        """

        self.MAX_SIZE = max_steps

        # Per-step we append tensor references to Python lists (free) and t.stack() once at the
        # end, instead of 5 indexed-write kernels per step into a preallocated buffer. Each stored
        # tensor is freshly produced per step (the env returns a new state tensor each step), so
        # holding references is safe. This removes ~2500 tiny kernel launches per full rollout.
        self._obs, self._actions, self._logprobs, self._rewards, self._dones = [], [], [], [], []
        self.timestep = 0

    def add_step(
        self,
        obs: Float[Tensor, " num_envs *obs_shape"],
        actions: Int[Tensor, " num_envs"],
        logprobs: Float[Tensor, " num_envs"],
        rewards: Float[Tensor, " num_envs"],
        dones: Bool[Tensor, " num_envs"],
    ):
        """
        Stores one timestep of the rollout (the same transition for all `num_envs` environments at once)
        and advances `self.timestep`.

        Args:
            obs:      (num_envs, *obs_shape)  the observation (o_t) the agent saw *before* acting
            actions:  (num_envs,)             the action (a_t) it took in each env after seeing o_t
            logprobs: (num_envs,)             log pi(a_t | o_t) of the action that was *taken* 
                                              Recorded at collection time, so
                                              these are the behaviour policy's logprobs (used for importance weights)
            rewards:  (num_envs,)             the reward (r_{t+1}) received for taking a_t
            dones:    (num_envs,)  bool       True if that env's episode ended on this step (d_{t+1} = terminated OR truncated),
                                              i.e. the env auto-reset and the *next* stored obs starts a new episode

        Raises ValueError if the rollout already holds `max_steps` timesteps.
        """

        if self.timestep >= self.MAX_SIZE:
            raise ValueError("Rollout is full, cannot add more steps")

        self._obs.append(obs)
        self._actions.append(actions)
        self._logprobs.append(logprobs)
        self._rewards.append(rewards)
        self._dones.append(dones)
        self.timestep += 1

    def reset(self):
        self._obs.clear(); self._actions.clear(); self._logprobs.clear()
        self._rewards.clear(); self._dones.clear()
        self.timestep = 0

    def get(self) -> tuple[Tensor, ...]:
        """
        Stack the per-step lists from (num_envs, *tensor_shape) into (num_envs, timestep, *tensor_shape) tensors. Rollouts can stop early
        (see gen_rollout), so the time dimension is however many steps were actually collected.
        """
        assert self.timestep > 0, "Rollout is empty"
        return RolloutTensors(
            t.stack(self._obs, dim=1),
            t.stack(self._actions, dim=1),
            t.stack(self._logprobs, dim=1),
            t.stack(self._rewards, dim=1).float(),
            t.stack(self._dones, dim=1),
        )

    def get_batches(self, batch_size: int) -> list[RolloutTensors]:
        """
        Splits the rollout buffer into batches of size `batch_size`, and returns a list of
        `RolloutTensors` objects, each containing `batch_size` many trajectories.
        """

        tau = self.get()  # filled portion only
        obs = t.split(tau.obs, batch_size, dim=0)
        acts = t.split(tau.actions, batch_size, dim=0)
        logprobs = t.split(tau.logprobs, batch_size, dim=0)
        rewards = t.split(tau.rewards, batch_size, dim=0)
        dones = t.split(tau.dones, batch_size, dim=0)

        batches = [RolloutTensors(*tensors) for tensors in zip(obs, acts, logprobs, rewards, dones)]

        return batches


tests.test_rollout(Rollout)

VPG Args

We've provided a dataclass for the training arguments, and will explain as needed later on.

@dataclass
class VPGArgs:
    # Basic / global
    seed: int = 1
    env_id: str = "CartPole-gpu"

    # Wandb / logging
    use_wandb: bool = False
    wandb_project_name: str = "VPGCartPole"
    wandb_entity: str | None = None
    video_log_freq: int | None = 50   # every N gradient steps, render a 4x4 grid video of the latest rollout
                                      # (logged to wandb if use_wandb, shown inline if live_viz)

    # Duration of different phases / buffer memory settings
    total_timesteps: int = 500_000
    # max_rollout_steps: int = 500
    # min_rollout_steps: int = 64
    num_envs: int = 4

    num_steps_per_rollout: int = 128

    lr: float = 2.5e-4
    gamma: float = 1
    rollout_use_count: int = 1
    clip_coef: float = 0.2
    device: str = "cpu"
    normalize_returns: bool = True
    num_batches_per_rollout: int = 1
    # LR decay settings
    use_lr_decay: bool = False
    lr_end: Optional[float] = None
    lr_frac: Optional[float] = None
    use_iw: bool = False
    full_reset: bool = True   # fully reset all envs at the start of each rollout
    live_viz: bool = True     # display the 4x4 grid video inline every `video_log_freq` gradient steps (notebook only)
    overwrite_video: bool = True  # each live video replaces the previous one in the notebook output (False appends them)

    def __post_init__(self):
        self.batch_size = self.num_envs // self.num_batches_per_rollout
        self.device = t.device(self.device)

        if self.use_lr_decay:
            assert self.lr_end is not None, "lr_end must be set if use_lr_decay is True"
            assert self.lr_frac is not None, "lr_frac must be set if use_lr_decay is True"

        self.env_steps_per_update = self.num_steps_per_rollout * self.num_envs // self.num_batches_per_rollout

        if not self.use_iw:
            assert self.rollout_use_count == 1, "rollout_use_count must be 1 if use_iw is False"
            assert self.num_batches_per_rollout == 1, "num_batches_per_rollout must be 1 if use_iw is False"

VPG Agent

The following class will be our agent, that will generate rollouts via interaction between the agent and environment, as well as generate actions by sampling them from the policy network. Recall that the policy network now maps observations to logits for each action, so we can sample actions from the distribution.

Exercise - implement VPGAgent

Difficulty: 🔴🔴🔴⚪⚪
Importance: 🔵🔵🔵🔵⚪
You should spend up to 10-15 minutes on this exercise.

Implement the functions: * gen_rollout - this function computes the episode rollout, by interacting with the environment for args.num_steps_per_rollout steps. If an episode terminates, the environment resets itself and we keep going, so a rollout usually contains several episodes per environment.

  • get_actions - this function takes in an observation, and returns the actions and their logprobs. Sample the actions with t.multinomial on the softmax of the policy network's logits (see the docs), then gather the logprob of each sampled action. The third return value is the entropy of the action distribution, but we don't need it here (it's recomputed in compute_logprobs_and_entropy, which is what the loss actually uses), so you can just return None for it.

To see how well the agent is doing we don't need any extra bookkeeping during the rollout: afterwards, the training loop sums each environment's rewards up to (and including) its first done to get the undiscounted return of its first episode, and reports the mean over environments with its standard error. For CartPole (+1 per step alive, max 500) that number is the episode length.

class VPGAgent:
    """Base Agent class handling the interaction with the environment."""

    def __init__(
        self,
        envs: gym.Env,
        policy_network: PolicyNetwork,
        args: VPGArgs,
    ):

        self.envs = envs
        self.policy_network = policy_network
        self.args = args
        self.obs_shape = envs.observation_space.shape
        self.action_shape = envs.action_space.shape

    @t.no_grad()
    def gen_rollout(self, rollout: Rollout) -> Rollout:
        """
        Steps all environments in parallel for `args.num_steps_per_rollout` steps, adding each step to the rollout
        buffer (which is reset first), and returns the filled buffer.
        """
        if self.args.full_reset and hasattr(self.envs, "_terminated"):
            self.envs._terminated[:] = True
            self.envs._truncated[:] = True
        obs, _ = self.envs.reset()  # Need a starting observation
        rollout.reset()

        raise NotImplementedError()

        return rollout

    def get_actions(
        self, obs: Float[Tensor, " num_envs *obs_shape"]
    ) -> tuple[Int[Tensor, " num_envs"], Float[Tensor, " num_envs"]]:
        """
        Computes the agent's turn: given an observation for each environment,
        sample the action the agent takes, along with the log_probs of that action.
        Use t.multinomial to sample the actions.
        """
        raise NotImplementedError()


tests.test_get_actions(VPGAgent, PolicyNetwork)
tests.test_gen_rollout(VPGAgent, PolicyNetwork, VPGArgs, Rollout)
Solution 
class VPGAgent:
    """Base Agent class handling the interaction with the environment."""

    def __init__(
        self,
        envs: gym.Env,
        policy_network: PolicyNetwork,
        args: VPGArgs,
    ):

        self.envs = envs
        self.policy_network = policy_network
        self.args = args
        self.obs_shape = envs.observation_space.shape
        self.action_shape = envs.action_space.shape

    @t.no_grad()
    def gen_rollout(self, rollout: Rollout) -> Rollout:
        """
        Steps all environments in parallel for `args.num_steps_per_rollout` steps, adding each step to the rollout
        buffer (which is reset first), and returns the filled buffer.
        """
        if self.args.full_reset and hasattr(self.envs, "_terminated"):
            self.envs._terminated[:] = True
            self.envs._truncated[:] = True
        obs, _ = self.envs.reset()  # Need a starting observation
        rollout.reset()

        for timestep in range(self.args.num_steps_per_rollout):
            actions, logprobs = self.get_actions(obs)
            new_obs, rewards, terminates, truncates, info = self.envs.step(actions)
            # Mask returns at episode boundaries on EITHER termination or truncation: the env
            # auto-resets on both, so returns must not sum across the reset boundary.
            done = terminates | truncates
            rollout.add_step(obs, actions, logprobs, rewards, done)
            obs = new_obs

        return rollout

    def get_actions(
        self, obs: Float[Tensor, " num_envs *obs_shape"]
    ) -> tuple[Int[Tensor, " num_envs"], Float[Tensor, " num_envs"]]:
        """
        Computes the agent's turn: given an observation for each environment,
        sample the action the agent takes, along with the log_probs of that action.
        Use t.multinomial to sample the actions.
        """
        logits = self.policy_network(obs)  # (num_envs, num_actions) unnormalised scores
        log_probs = F.log_softmax(logits, dim=-1)  # log pi(a|s): logits - logsumexp, numerically stable & differentiable
        # One action index per env, drawn with probability pi(a|s) (multinomial renormalises each row, so we just need
        # non-negative weights). The result is an integer tensor: no gradient flows through a discrete draw, and none should.
        actions = t.multinomial(log_probs.exp(), num_samples=1).squeeze(-1)
        # alternative with eindex
        # logprobs_per_act = eindex(log_probs, actions, "env [env] -> env")
        # alternative with gather
        logprobs_per_act = log_probs.gather(-1, actions.unsqueeze(-1)).squeeze(-1)        
        return actions, logprobs_per_act


tests.test_get_actions(VPGAgent, PolicyNetwork)
tests.test_gen_rollout(VPGAgent, PolicyNetwork, VPGArgs, Rollout)