# Copyright (c) 2021-2025, ETH Zurich and NVIDIA CORPORATION # All rights reserved. # # SPDX-License-Identifier: BSD-3-Clause from __future__ import annotations import torch from tensordict import TensorDict from rsl_rl.utils import split_and_pad_trajectories class RolloutStorage: class Transition: def __init__(self): self.observations = None self.actions = None self.privileged_actions = None self.rewards = None self.dones = None self.values = None self.actions_log_prob = None self.action_mean = None self.action_sigma = None self.hidden_states = None def clear(self): self.__init__() def __init__( self, training_type, num_envs, num_transitions_per_env, obs, actions_shape, device="cpu", ): # store inputs self.training_type = training_type self.device = device self.num_transitions_per_env = num_transitions_per_env self.num_envs = num_envs self.actions_shape = actions_shape # Core self.observations = TensorDict( {key: torch.zeros(num_transitions_per_env, *value.shape, device=device) for key, value in obs.items()}, batch_size=[num_transitions_per_env, num_envs], device=self.device, ) self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte() # for distillation if training_type == "distillation": self.privileged_actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) # for reinforcement learning if training_type == "rl": self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device) self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device) # For RNN networks self.saved_hidden_states_a = None self.saved_hidden_states_c = None # counter for the number of transitions stored self.step = 0 def add_transitions(self, transition: Transition): # check if the transition is valid if self.step >= self.num_transitions_per_env: raise OverflowError("Rollout buffer overflow! You should call clear() before adding new transitions.") # Core self.observations[self.step].copy_(transition.observations) self.actions[self.step].copy_(transition.actions) self.rewards[self.step].copy_(transition.rewards.view(-1, 1)) self.dones[self.step].copy_(transition.dones.view(-1, 1)) # for distillation if self.training_type == "distillation": self.privileged_actions[self.step].copy_(transition.privileged_actions) # for reinforcement learning if self.training_type == "rl": self.values[self.step].copy_(transition.values) self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1)) self.mu[self.step].copy_(transition.action_mean) self.sigma[self.step].copy_(transition.action_sigma) # For RNN networks self._save_hidden_states(transition.hidden_states) # increment the counter self.step += 1 def _save_hidden_states(self, hidden_states): if hidden_states is None or hidden_states == (None, None): return # make a tuple out of GRU hidden state sto match the LSTM format hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],) hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],) # initialize if needed if self.saved_hidden_states_a is None: self.saved_hidden_states_a = [ torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a)) ] self.saved_hidden_states_c = [ torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c)) ] # copy the states for i in range(len(hid_a)): self.saved_hidden_states_a[i][self.step].copy_(hid_a[i]) self.saved_hidden_states_c[i][self.step].copy_(hid_c[i]) def clear(self): self.step = 0 def compute_returns(self, last_values, gamma, lam, normalize_advantage: bool = True): advantage = 0 for step in reversed(range(self.num_transitions_per_env)): # if we are at the last step, bootstrap the return value if step == self.num_transitions_per_env - 1: next_values = last_values else: next_values = self.values[step + 1] # 1 if we are not in a terminal state, 0 otherwise next_is_not_terminal = 1.0 - self.dones[step].float() # TD error: r_t + gamma * V(s_{t+1}) - V(s_t) delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step] # Advantage: A(s_t, a_t) = delta_t + gamma * lambda * A(s_{t+1}, a_{t+1}) advantage = delta + next_is_not_terminal * gamma * lam * advantage # Return: R_t = A(s_t, a_t) + V(s_t) self.returns[step] = advantage + self.values[step] # Compute the advantages self.advantages = self.returns - self.values # Normalize the advantages if flag is set # This is to prevent double normalization (i.e. if per minibatch normalization is used) if normalize_advantage: self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8) # for distillation def generator(self): if self.training_type != "distillation": raise ValueError("This function is only available for distillation training.") for i in range(self.num_transitions_per_env): yield self.observations[i], self.actions[i], self.privileged_actions[i], self.dones[i] # for reinforcement learning with feedforward networks def mini_batch_generator(self, num_mini_batches, num_epochs=8): if self.training_type != "rl": raise ValueError("This function is only available for reinforcement learning training.") batch_size = self.num_envs * self.num_transitions_per_env mini_batch_size = batch_size // num_mini_batches indices = torch.randperm(num_mini_batches * mini_batch_size, requires_grad=False, device=self.device) # Core observations = self.observations.flatten(0, 1) actions = self.actions.flatten(0, 1) values = self.values.flatten(0, 1) returns = self.returns.flatten(0, 1) # For PPO old_actions_log_prob = self.actions_log_prob.flatten(0, 1) advantages = self.advantages.flatten(0, 1) old_mu = self.mu.flatten(0, 1) old_sigma = self.sigma.flatten(0, 1) for epoch in range(num_epochs): for i in range(num_mini_batches): # Select the indices for the mini-batch start = i * mini_batch_size end = (i + 1) * mini_batch_size batch_idx = indices[start:end] # Create the mini-batch # -- Core obs_batch = observations[batch_idx] actions_batch = actions[batch_idx] # -- For PPO target_values_batch = values[batch_idx] returns_batch = returns[batch_idx] old_actions_log_prob_batch = old_actions_log_prob[batch_idx] advantages_batch = advantages[batch_idx] old_mu_batch = old_mu[batch_idx] old_sigma_batch = old_sigma[batch_idx] # yield the mini-batch yield obs_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( None, None, ), None # for reinfrocement learning with recurrent networks def recurrent_mini_batch_generator(self, num_mini_batches, num_epochs=8): if self.training_type != "rl": raise ValueError("This function is only available for reinforcement learning training.") padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones) mini_batch_size = self.num_envs // num_mini_batches for ep in range(num_epochs): first_traj = 0 for i in range(num_mini_batches): start = i * mini_batch_size stop = (i + 1) * mini_batch_size dones = self.dones.squeeze(-1) last_was_done = torch.zeros_like(dones, dtype=torch.bool) last_was_done[1:] = dones[:-1] last_was_done[0] = True trajectories_batch_size = torch.sum(last_was_done[:, start:stop]) last_traj = first_traj + trajectories_batch_size masks_batch = trajectory_masks[:, first_traj:last_traj] obs_batch = padded_obs_trajectories[:, first_traj:last_traj] actions_batch = self.actions[:, start:stop] old_mu_batch = self.mu[:, start:stop] old_sigma_batch = self.sigma[:, start:stop] returns_batch = self.returns[:, start:stop] advantages_batch = self.advantages[:, start:stop] values_batch = self.values[:, start:stop] old_actions_log_prob_batch = self.actions_log_prob[:, start:stop] # reshape to [num_envs, time, num layers, hidden dim] (original shape: [time, num_layers, num_envs, hidden_dim]) # then take only time steps after dones (flattens num envs and time dimensions), # take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim] last_was_done = last_was_done.permute(1, 0) hid_a_batch = [ saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] .transpose(1, 0) .contiguous() for saved_hidden_states in self.saved_hidden_states_a ] hid_c_batch = [ saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj] .transpose(1, 0) .contiguous() for saved_hidden_states in self.saved_hidden_states_c ] # remove the tuple for GRU hid_a_batch = hid_a_batch[0] if len(hid_a_batch) == 1 else hid_a_batch hid_c_batch = hid_c_batch[0] if len(hid_c_batch) == 1 else hid_c_batch yield obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, ( hid_a_batch, hid_c_batch, ), masks_batch first_traj = last_traj