@torch.no_grad()<EOL><INDENT>def rollout(self, batch_info: BatchInfo, model: Model, number_of_steps: int) -> Rollout:<DEDENT>
accumulator = TensorAccumulator()<EOL>episode_information = []  <EOL>if self.hidden_state is None and model.is_recurrent:<EOL><INDENT>self.hidden_state = model.zero_state(self.last_observation.size(<NUM_LIT:0>)).to(self.device)<EOL><DEDENT>initial_hidden_state = self.hidden_state<EOL>for step_idx in range(number_of_steps):<EOL><INDENT>if model.is_recurrent:<EOL><INDENT>step = model.step(self.last_observation.to(self.device), state=self.hidden_state)<EOL>self.hidden_state = step['<STR_LIT:state>']<EOL><DEDENT>else:<EOL><INDENT>step = model.step(self.last_observation.to(self.device))<EOL><DEDENT>for name, tensor in step.items():<EOL><INDENT>accumulator.add(name, tensor.cpu())<EOL><DEDENT>accumulator.add('<STR_LIT>', self.last_observation)<EOL>actions_numpy = step['<STR_LIT>'].detach().cpu().numpy()<EOL>new_obs, new_rewards, new_dones, new_infos = self.environment.step(actions_numpy)<EOL>dones_tensor = torch.from_numpy(new_dones.astype(np.float32)).clone()<EOL>accumulator.add('<STR_LIT>', dones_tensor)<EOL>self.last_observation = torch.from_numpy(new_obs).clone()<EOL>if model.is_recurrent:<EOL><INDENT>self.hidden_state = self.hidden_state * (<NUM_LIT:1.0> - dones_tensor.unsqueeze(-<NUM_LIT:1>)).to(self.device)<EOL><DEDENT>accumulator.add('<STR_LIT>', torch.from_numpy(new_rewards.astype(np.float32)).clone())<EOL>episode_information.append(new_infos)<EOL><DEDENT>if model.is_recurrent:<EOL><INDENT>final_values = model.value(self.last_observation.to(self.device), state=self.hidden_state).cpu()<EOL><DEDENT>else:<EOL><INDENT>final_values = model.value(self.last_observation.to(self.device)).cpu()<EOL><DEDENT>accumulated_tensors = accumulator.result()<EOL>return Trajectories(<EOL>num_steps=accumulated_tensors['<STR_LIT>'].size(<NUM_LIT:0>),<EOL>num_envs=accumulated_tensors['<STR_LIT>'].size(<NUM_LIT:1>),<EOL>environment_information=episode_information,<EOL>transition_tensors=accumulated_tensors,<EOL>rollout_tensors={<EOL>'<STR_LIT>': initial_hidden_state.cpu() if initial_hidden_state is not None else None,<EOL>'<STR_LIT>': final_values<EOL>}<EOL>)<EOL>
Calculate env rollout