if save_model:
        torch.save(self.policy_network.state_dict(), self.model_path)
        print("Saved model to disk")

    if save_video:
        self.video.close()

    self.env.close()
    return sum(total_rewards) / len(total_rewards)