87232aa39c159f04a501ad268012da45a1ff537c,torch_rl/torch_rl/algos/base.py,BaseAlgo,collect_experiences,#BaseAlgo#,68

Before Change

        log = {
            "return_per_episode": self.log_return.cpu().numpy(),
            "reshaped_return_per_episode": self.log_reshaped_return.cpu().numpy(),
            "num_frames_per_episode": self.log_num_frames.cpu().numpy(),
            "num_frames": self.num_frames
        }

After Change

            self.log_episode_reshaped_return += self.rewards[i]
            self.log_episode_num_frames += torch.ones(self.num_procs, device=self.device)

            for i, done_ in enumerate(done):
                if done_:
                    self.log_done_counter += 1
                    self.log_return.append(self.log_episode_return[i].item())
                    self.log_reshaped_return.append(self.log_episode_reshaped_return[i].item())
                    self.log_num_frames.append(self.log_episode_num_frames[i].item())

            self.log_episode_return *= self.mask
            self.log_episode_reshaped_return *= self.mask
            self.log_episode_num_frames *= self.mask

Instances