while not done:
            if RENDER_MODE == 'human':
                self.env.render()

            if save_video:
                self.video.capture_frame()

            state = torch.FloatTensor(state).unsqueeze(0).to(device)
            action_probs = self.policy_network(state)
            action = torch.multinomial(action_probs, 1).item()
            log_prob = torch.log(action_probs.squeeze(0)[action])
            log_probs.append(log_prob)

            next_state, reward, done, _, _ = self.env.step(action)
            rewards.append(reward)
            state = next_state
            total_reward += reward