import time
import numpy as np
import torch
from onpolicy.runner.shared.base_runner import Runner
import wandb
import imageio
import os
from torch.utils.data import DataLoader

from onpolicy.algorithms.r_mappo.r_mappo import R_MAPPO as TrainAlgo
from onpolicy.algorithms.r_mappo.algorithm.rMAPPOPolicy import R_MAPPOPolicy as Policy
from onpolicy.algorithms.r_mappo.algorithm.r_actor_critic import Discriminator
from onpolicy.utils.shared_buffer import SharedReplayBuffer
from onpolicy.utils.trajectory_dataset import TrajectoryDataset

from onpolicy.runner.shared.mpe_runner import MPERunner

np.set_printoptions(precision=4)


def _t2n(x):
    return x.detach().cpu().numpy()


def kl_div(p1, p2):
    return (p1 * (p1.log() - p2.log())).sum(-1)


def js_div(p1, p2):
    m = (p1 + p2) / 2
    return (kl_div(p1, m) + kl_div(p2, m)) / 2


class StagHuntRunner(MPERunner):
    """Runner class to perform training, evaluation. and data collection for the MPEs. See parent class for details."""

    def run(self):
        for rspo_iter in range(self.start_iter, self.all_args.rspo_iterations):

            # policy network
            self.policy = Policy(self.all_args,
                                 self.envs.observation_space[0],
                                 self.share_observation_space,
                                 self.envs.action_space[0],
                                 device=self.device,
                                 critic_dim=rspo_iter + 1)
            self.policy.actor.load_state_dict(
                self.rnd_policy.actor.state_dict())
            self.policy.critics[0].load_state_dict(
                self.rnd_policy.critics[0].state_dict())

            for policy in self.ref_policies:
                policy.actor.eval()
                for critic in policy.critics:
                    critic.eval()

            if self.model_dir is not None:
                self.restore()

            # algorithm
            self.trainer = TrainAlgo(self.all_args,
                                     self.policy,
                                     device=self.device)

            # buffer
            self.buffer = SharedReplayBuffer(rspo_iter, self.all_args,
                                             self.num_agents,
                                             self.envs.observation_space[0],
                                             self.share_observation_space,
                                             self.envs.action_space[0])

            self.logger.info("\n" + "-" * 20 +
                             f"  RSPO iteration {rspo_iter} starts...  " +
                             "-" * 20 + "\n")
            self.warmup()

            start = time.time()
            episodes = int(self.num_env_steps
                           ) // self.episode_length // self.n_rollout_threads

            self.rspo_iter = rspo_iter
            # initialize lagerangian lambda
            if rspo_iter > 0:
                self.ll = torch.zeros(rspo_iter)
            else:
                self.ll = torch.tensor([])

            for episode in range(episodes):
                ep_cnt = 0
                total_div = []
                env_infos = {
                    'collective_return': [],
                    'coop&coop_num': [],
                    'gore1_num': [],
                    'gore2_num': [],
                    'hare1_num': [],
                    'hare2_num': [],
                }

                if self.use_linear_lr_decay:
                    self.trainer.policy.lr_decay(episode, episodes)

                # prepare positive batch for wasserstein distance
                traj_data_p = []
                p_scores = np.zeros((self.all_args.num_agents, rspo_iter),
                                    dtype=np.float32)
                for d_i, (gen, loader, d) in enumerate(
                        zip(self.traj_generators, self.traj_loaders,
                            self.discriminators)):
                    try:
                        data_p = next(gen)
                    except StopIteration:
                        self.traj_generators[d_i] = gen = iter(loader)
                        data_p = next(gen)
                    with torch.no_grad():
                        share_obs, action, reward = data_p
                        p_score = d(share_obs, action,
                                    reward).mean(0).squeeze(-1).cpu().numpy()
                        assert p_score.shape == (
                            self.all_args.num_agents, ), p_score.shape
                        p_scores[:, d_i] = p_score
                    traj_data_p.append(data_p)

                for step in range(self.episode_length):
                    # Sample actions
                    (values, actions, action_log_probs, rnn_states,
                     rnn_states_critic, actions_env) = self.collect(step)

                    # Obser reward and next obs
                    obs, rewards, dones, infos = self.envs.step(actions_env)
                    assert rewards.shape == (self.all_args.n_rollout_threads,
                                             self.all_args.num_agents,
                                             1), rewards.shape
                    rewards *= self.all_args.reward_scaling

                    if self.rspo_iter > 0:
                        # compute divergence reward at the current step
                        share_obs = self.buffer.share_obs[step]
                        divs = np.zeros(
                            (self.all_args.n_rollout_threads,
                             self.all_args.num_agents, self.rspo_iter),
                            dtype=np.float32)
                        for d_i, d in enumerate(self.discriminators):
                            with torch.no_grad():
                                n_score = d(share_obs, actions,
                                            rewards).squeeze(-1).cpu().numpy()
                            assert n_score.shape == (
                                self.all_args.n_rollout_threads,
                                self.all_args.num_agents), n_score.shape
                            w_dist = p_scores[:, d_i] - n_score
                            divs[..., d_i] = w_dist
                        divs *= self.all_args.intrinsic_reward_scaling

                        # NOTE: div_ret should not be immediately reset in env.reset
                        div_rets = self.envs.update_div(divs)
                        rewards = np.concatenate([rewards, divs], -1)
                        assert rewards.shape[-1] == self.rspo_iter + 1

                    # divs = []
                    # if self.rspo_iter > 0:
                    #     pi = self.trainer.policy.get_pi(
                    #         obs, rnn_states, masks).probs
                    #     ref_pis = [
                    #         policy.get_pi(obs, rnn_states, masks).probs
                    #         for policy in self.ref_policies
                    #     ]
                    #     div = torch.stack(
                    #         [js_div(pi, ref_pi) for ref_pi in ref_pis], -1)
                    #     divs = np.array(
                    #         np.split(_t2n(div), self.n_rollout_threads))
                    # else:
                    #     divs = [None for _ in range(obs.shape[0])]

                    for env_i, (done, info) in enumerate(zip(dones, infos)):
                        if done.all():
                            ep_cnt += 1
                            for k in info:
                                env_infos[k].append(info[k])
                            if self.rspo_iter > 0:
                                assert div_rets[env_i].shape == (
                                    self.all_args.num_agents, self.rspo_iter)
                                total_div.append(div_rets[env_i].mean(0))

                    data = obs, rewards, dones, infos, values, actions, action_log_probs, rnn_states, rnn_states_critic

                    # insert data into buffer
                    self.insert(data)

                # update wasserstein discriminator
                if self.rspo_iter > 0:
                    d_losses = np.zeros(rspo_iter, dtype=np.float32)
                    d_p_scores = np.zeros(rspo_iter, dtype=np.float32)
                    d_n_scores = np.zeros(rspo_iter, dtype=np.float32)
                    for d_i, (loader, data_p, d, d_opt) in enumerate(
                            zip(self.traj_loaders, traj_data_p,
                                self.discriminators, self.discriminator_opts)):
                        share_obs_p, action_p, reward_p = data_p
                        share_obs_n, action_n, reward_n = map(
                            loader.dataset.cast, [
                                self.buffer.share_obs[:-1],
                                self.buffer.actions,
                                self.buffer.rewards[..., :1]
                            ])

                        p_score = d(share_obs_p, action_p, reward_p)
                        n_score = d(share_obs_n, action_n, reward_n)
                        d_loss = -(p_score - n_score).mean()

                        d_opt.zero_grad()
                        d_loss.backward()
                        d_opt.step()

                        for p in d.parameters():
                            p.data.clamp_(-0.01, 0.01)

                        d_losses[d_i] = d_loss.item()
                        d_p_scores[d_i] = p_score.mean().item()
                        d_n_scores[d_i] = n_score.mean().item()

                # update lagrange lambda
                if self.rspo_iter > 0 and ep_cnt > 0:
                    self.ll = self.ll + self.all_args.ll_lr * (
                        self.all_args.threshold_eps -
                        sum(total_div) / len(total_div))
                    self.ll.clamp_(0, self.all_args.ll_max)
                    # self.logger.info((self.ll, sum(total_div) / len(total_div)))

                # compute return and update policy
                self.compute()
                train_infos = self.train(episode)

                # post process
                total_num_steps = (
                    episode + 1) * self.episode_length * self.n_rollout_threads

                # save model
                if (episode % self.save_interval == 0
                        or episode == episodes - 1):
                    self.save()

                # log information
                if (episode + 1) % self.log_interval == 0:
                    end = time.time()
                    self.logger.info(
                        "\n Env {} Algo {} Exp {} updates {}/{} episodes, total num timesteps {}/{}, RSPO iteration {}/{}. FPS {}."
                        .format(
                            self.all_args.env_name,
                            self.algorithm_name,
                            self.experiment_name,
                            episode + 1,
                            episodes,
                            total_num_steps,
                            self.num_env_steps,
                            self.rspo_iter + 1,
                            self.all_args.rspo_iterations,
                            int(total_num_steps / (end - start)),
                        ))

                    train_infos["average_episode_rewards"] = np.mean(
                        self.buffer.rewards[..., 0]) * self.episode_length
                    if ep_cnt > 0:
                        if self.rspo_iter > 0:
                            valid_ep_cnt = sum([
                                (div > self.all_args.threshold_eps).all()
                                for div in total_div
                            ])
                        else:
                            valid_ep_cnt = ep_cnt
                        train_infos['acceptance_ratio'] = valid_ep_cnt / ep_cnt
                        self.logger.info(
                            "average episode rewards is {:.2f}, acceptance ratio {:.2%} ({:d}/{:d}), env info {}"
                            .format(
                                train_infos["average_episode_rewards"],
                                train_infos['acceptance_ratio'],
                                valid_ep_cnt,
                                ep_cnt,
                                {k: np.mean(v)
                                 for k, v in env_infos.items()},
                            ))

                    if self.rspo_iter > 0:
                        if ep_cnt > 0:
                            self.logger.info(
                                f"Average diversity episode return {sum(total_div) / len(total_div)}"
                                f", threshold {self.all_args.threshold_eps}.",
                            )
                        self.logger.info(
                            f"Wasserstein Discriminator loss {d_losses}, "
                            f"positive scores {d_p_scores}, negative scores{d_n_scores}."
                        )
                        for ref_idx in range(self.rspo_iter):
                            avg_r = self.buffer.rewards[...,
                                                        1:].mean(axis=(0, 1,
                                                                       2))
                            train_infos[
                                f'episode_intrinsic_rewards/{ref_idx}'] = avg_r[
                                    ref_idx] * self.episode_length
                            train_infos[
                                f'lagrangian_lambda/{ref_idx}'] = self.ll[
                                    ref_idx].item()
                        self.logger.info(
                            "average episode intrinsic reward is {}, lagrangian lambda is {} (multiplied {}, average {:.2f})"
                            .format(
                                np.array([
                                    train_infos[
                                        f'episode_intrinsic_rewards/{idx}']
                                    for idx in range(self.rspo_iter)
                                ]),
                                np.array([
                                    train_infos[f'lagrangian_lambda/{idx}']
                                    for idx in range(self.rspo_iter)
                                ]),
                                np.array([
                                    train_infos[f'lagrangian_lambda/{idx}'] *
                                    train_infos[
                                        f'episode_intrinsic_rewards/{idx}']
                                    for idx in range(self.rspo_iter)
                                ]),
                                sum([
                                    train_infos[f'lagrangian_lambda/{idx}'] *
                                    train_infos[
                                        f'episode_intrinsic_rewards/{idx}']
                                    for idx in range(self.rspo_iter)
                                ]) / self.rspo_iter,
                            ))
                    self.log_train(
                        train_infos,
                        total_num_steps + rspo_iter * self.num_env_steps)
                    if ep_cnt > 0:
                        self.log_env(
                            env_infos,
                            total_num_steps + rspo_iter * self.num_env_steps)

                # eval
                if episode % self.eval_interval == 0 and self.use_eval:
                    self.eval(total_num_steps)

            self.save_trajectory_samples()

            traj_dataset = TrajectoryDataset(self.save_dir, self.rspo_iter)
            assert len(
                traj_dataset) % self.all_args.n_rollout_threads == 0, len(
                    traj_dataset)
            assert not self.all_args.use_recurrent_policy and not self.all_args.use_naive_recurrent_policy
            traj_loader = DataLoader(
                traj_dataset,
                batch_size=self.all_args.episode_length *
                self.all_args.n_rollout_threads,
                shuffle=True,
            )
            d = Discriminator(self.all_args,
                              self.share_observation_space,
                              self.action_space,
                              device=self.device).to(self.device)
            for p in d.parameters():
                p.data.clamp_(-0.01, 0.01)

            self.ref_policies.append(self.trainer.policy)
            self.traj_loaders.append(traj_loader)
            self.traj_generators.append(iter(traj_loader))
            self.discriminators.append(d)
            self.discriminator_opts.append(
                torch.optim.RMSprop(d.parameters(), lr=self.all_args.d_lr))

            self.save()

import time
import numpy as np
import torch
from onpolicy.runner.shared.base_runner import Runner
import wandb
import imageio
import os
from torch.utils.data import DataLoader

from onpolicy.algorithms.r_mappo.r_mappo import R_MAPPO as TrainAlgo
from onpolicy.algorithms.r_mappo.algorithm.rMAPPOPolicy import R_MAPPOPolicy as Policy
from onpolicy.algorithms.r_mappo.algorithm.r_actor_critic import Discriminator
from onpolicy.utils.shared_buffer import SharedReplayBuffer
from onpolicy.utils.trajectory_dataset import TrajectoryDataset

from onpolicy.runner.shared.mpe_runner import MPERunner

np.set_printoptions(precision=4)


def _t2n(x):
    return x.detach().cpu().numpy()


def kl_div(p1, p2):
    return (p1 * (p1.log() - p2.log())).sum(-1)


def js_div(p1, p2):
    m = (p1 + p2) / 2
    return (kl_div(p1, m) + kl_div(p2, m)) / 2


class StagHuntRunner(MPERunner):
    """Runner class to perform training, evaluation. and data collection for the MPEs. See parent class for details."""

    def run(self):
        for rspo_iter in range(self.start_iter, self.all_args.rspo_iterations):

            # policy network
            self.policy = Policy(self.all_args,
                                 self.envs.observation_space[0],
                                 self.share_observation_space,
                                 self.envs.action_space[0],
                                 device=self.device,
                                 critic_dim=rspo_iter + 1)
            self.policy.actor.load_state_dict(
                self.rnd_policy.actor.state_dict())
            self.policy.critics[0].load_state_dict(
                self.rnd_policy.critics[0].state_dict())

            for policy in self.ref_policies:
                policy.actor.eval()
                for critic in policy.critics:
                    critic.eval()

            if self.model_dir is not None:
                self.restore()

            # algorithm
            self.trainer = TrainAlgo(self.all_args,
                                     self.policy,
                                     device=self.device)

            # buffer
            self.buffer = SharedReplayBuffer(rspo_iter, self.all_args,
                                             self.num_agents,
                                             self.envs.observation_space[0],
                                             self.share_observation_space,
                                             self.envs.action_space[0])

            self.logger.info("\n" + "-" * 20 +
                             f"  RSPO iteration {rspo_iter} starts...  " +
                             "-" * 20 + "\n")
            self.warmup()

            start = time.time()
            episodes = int(self.num_env_steps
                           ) // self.episode_length // self.n_rollout_threads

            self.rspo_iter = rspo_iter
            # initialize lagerangian lambda
            if rspo_iter > 0:
                self.ll = torch.zeros(rspo_iter)
            else:
                self.ll = torch.tensor([])

            for episode in range(episodes):
                ep_cnt = 0
                total_div = []
                env_infos = {
                    'collective_return': [],
                    'coop&coop_num': [],
                    'gore1_num': [],
                    'gore2_num': [],
                    'hare1_num': [],
                    'hare2_num': [],
                }

                if self.use_linear_lr_decay:
                    self.trainer.policy.lr_decay(episode, episodes)

                # prepare positive batch for wasserstein distance
                traj_data_p = []
                p_scores = np.zeros((self.all_args.num_agents, rspo_iter),
                                    dtype=np.float32)
                for d_i, (gen, loader, d) in enumerate(
                        zip(self.traj_generators, self.traj_loaders,
                            self.discriminators)):
                    try:
                        data_p = next(gen)
                    except StopIteration:
                        self.traj_generators[d_i] = gen = iter(loader)
                        data_p = next(gen)
                    with torch.no_grad():
                        share_obs, action, reward = data_p
                        p_score = d(share_obs, action,
                                    reward).mean(0).squeeze(-1).cpu().numpy()
                        assert p_score.shape == (
                            self.all_args.num_agents, ), p_score.shape
                        p_scores[:, d_i] = p_score
                    traj_data_p.append(data_p)

                for step in range(self.episode_length):
                    # Sample actions
                    (values, actions, action_log_probs, rnn_states,
                     rnn_states_critic, actions_env) = self.collect(step)

                    # Obser reward and next obs
                    obs, rewards, dones, infos = self.envs.step(actions_env)
                    assert rewards.shape == (self.all_args.n_rollout_threads,
                                             self.all_args.num_agents,
                                             1), rewards.shape
                    rewards *= self.all_args.reward_scaling

                    if self.rspo_iter > 0:
                        # compute divergence reward at the current step
                        share_obs = self.buffer.share_obs[step]
                        divs = np.zeros(
                            (self.all_args.n_rollout_threads,
                             self.all_args.num_agents, self.rspo_iter),
                            dtype=np.float32)
                        for d_i, d in enumerate(self.discriminators):
                            with torch.no_grad():
                                n_score = d(share_obs, actions,
                                            rewards).squeeze(-1).cpu().numpy()
                            assert n_score.shape == (
                                self.all_args.n_rollout_threads,
                                self.all_args.num_agents), n_score.shape
                            w_dist = p_scores[:, d_i] - n_score
                            divs[..., d_i] = w_dist
                        divs *= self.all_args.intrinsic_reward_scaling

                        # NOTE: div_ret should not be immediately reset in env.reset
                        div_rets = self.envs.update_div(divs)
                        rewards = np.concatenate([rewards, divs], -1)
                        assert rewards.shape[-1] == self.rspo_iter + 1

                    # divs = []
                    # if self.rspo_iter > 0:
                    #     pi = self.trainer.policy.get_pi(
                    #         obs, rnn_states, masks).probs
                    #     ref_pis = [
                    #         policy.get_pi(obs, rnn_states, masks).probs
                    #         for policy in self.ref_policies
                    #     ]
                    #     div = torch.stack(
                    #         [js_div(pi, ref_pi) for ref_pi in ref_pis], -1)
                    #     divs = np.array(
                    #         np.split(_t2n(div), self.n_rollout_threads))
                    # else:
                    #     divs = [None for _ in range(obs.shape[0])]

                    for env_i, (done, info) in enumerate(zip(dones, infos)):
                        if done.all():
                            ep_cnt += 1
                            for k in info:
                                env_infos[k].append(info[k])
                            if self.rspo_iter > 0:
                                assert div_rets[env_i].shape == (
                                    self.all_args.num_agents, self.rspo_iter)
                                total_div.append(div_rets[env_i].mean(0))

                    data = obs, rewards, dones, infos, values, actions, action_log_probs, rnn_states, rnn_states_critic

                    # insert data into buffer
                    self.insert(data)

                # update wasserstein discriminator
                if self.rspo_iter > 0:
                    d_losses = np.zeros(rspo_iter, dtype=np.float32)
                    d_p_scores = np.zeros(rspo_iter, dtype=np.float32)
                    d_n_scores = np.zeros(rspo_iter, dtype=np.float32)
                    for d_i, (loader, data_p, d, d_opt) in enumerate(
                            zip(self.traj_loaders, traj_data_p,
                                self.discriminators, self.discriminator_opts)):
                        share_obs_p, action_p, reward_p = data_p
                        share_obs_n, action_n, reward_n = map(
                            loader.dataset.cast, [
                                self.buffer.share_obs[:-1],
                                self.buffer.actions,
                                self.buffer.rewards[..., :1]
                            ])

                        p_score = d(share_obs_p, action_p, reward_p)
                        n_score = d(share_obs_n, action_n, reward_n)
                        d_loss = -(p_score - n_score).mean()

                        d_opt.zero_grad()
                        d_loss.backward()
                        d_opt.step()

                        for p in d.parameters():
                            p.data.clamp_(-0.01, 0.01)

                        d_losses[d_i] = d_loss.item()
                        d_p_scores[d_i] = p_score.mean().item()
                        d_n_scores[d_i] = n_score.mean().item()

                # update lagrange lambda
                if self.rspo_iter > 0 and ep_cnt > 0:
                    self.ll = self.ll + self.all_args.ll_lr * (
                        self.all_args.threshold_eps -
                        sum(total_div) / len(total_div))
                    self.ll.clamp_(0, self.all_args.ll_max)
                    # self.logger.info((self.ll, sum(total_div) / len(total_div)))

                # compute return and update policy
                self.compute()
                train_infos = self.train(episode)

                # post process
                total_num_steps = (
                    episode + 1) * self.episode_length * self.n_rollout_threads

                # save model
                if (episode % self.save_interval == 0
                        or episode == episodes - 1):
                    self.save()

                # log information
                if (episode + 1) % self.log_interval == 0:
                    end = time.time()
                    self.logger.info(
                        "\n Env {} Algo {} Exp {} updates {}/{} episodes, total num timesteps {}/{}, RSPO iteration {}/{}. FPS {}."
                        .format(
                            self.all_args.env_name,
                            self.algorithm_name,
                            self.experiment_name,
                            episode + 1,
                            episodes,
                            total_num_steps,
                            self.num_env_steps,
                            self.rspo_iter + 1,
                            self.all_args.rspo_iterations,
                            int(total_num_steps / (end - start)),
                        ))

                    train_infos["average_episode_rewards"] = np.mean(
                        self.buffer.rewards[..., 0]) * self.episode_length
                    if ep_cnt > 0:
                        if self.rspo_iter > 0:
                            valid_ep_cnt = sum([
                                (div > self.all_args.threshold_eps).all()
                                for div in total_div
                            ])
                        else:
                            valid_ep_cnt = ep_cnt
                        train_infos['acceptance_ratio'] = valid_ep_cnt / ep_cnt
                        self.logger.info(
                            "average episode rewards is {:.2f}, acceptance ratio {:.2%} ({:d}/{:d}), env info {}"
                            .format(
                                train_infos["average_episode_rewards"],
                                train_infos['acceptance_ratio'],
                                valid_ep_cnt,
                                ep_cnt,
                                {k: np.mean(v)
                                 for k, v in env_infos.items()},
                            ))

                    if self.rspo_iter > 0:
                        if ep_cnt > 0:
                            self.logger.info(
                                f"Average diversity episode return {sum(total_div) / len(total_div)}"
                                f", threshold {self.all_args.threshold_eps}.",
                            )
                        self.logger.info(
                            f"Wasserstein Discriminator loss {d_losses}, "
                            f"positive scores {d_p_scores}, negative scores{d_n_scores}."
                        )
                        for ref_idx in range(self.rspo_iter):
                            avg_r = self.buffer.rewards[...,
                                                        1:].mean(axis=(0, 1,
                                                                       2))
                            train_infos[
                                f'episode_intrinsic_rewards/{ref_idx}'] = avg_r[
                                    ref_idx] * self.episode_length
                            train_infos[
                                f'lagrangian_lambda/{ref_idx}'] = self.ll[
                                    ref_idx].item()
                        self.logger.info(
                            "average episode intrinsic reward is {}, lagrangian lambda is {} (multiplied {}, average {:.2f})"
                            .format(
                                np.array([
                                    train_infos[
                                        f'episode_intrinsic_rewards/{idx}']
                                    for idx in range(self.rspo_iter)
                                ]),
                                np.array([
                                    train_infos[f'lagrangian_lambda/{idx}']
                                    for idx in range(self.rspo_iter)
                                ]),
                                np.array([
                                    train_infos[f'lagrangian_lambda/{idx}'] *
                                    train_infos[
                                        f'episode_intrinsic_rewards/{idx}']
                                    for idx in range(self.rspo_iter)
                                ]),
                                sum([
                                    train_infos[f'lagrangian_lambda/{idx}'] *
                                    train_infos[
                                        f'episode_intrinsic_rewards/{idx}']
                                    for idx in range(self.rspo_iter)
                                ]) / self.rspo_iter,
                            ))
                    self.log_train(
                        train_infos,
                        total_num_steps + rspo_iter * self.num_env_steps)
                    if ep_cnt > 0:
                        self.log_env(
                            env_infos,
                            total_num_steps + rspo_iter * self.num_env_steps)

                # eval
                if episode % self.eval_interval == 0 and self.use_eval:
                    self.eval(total_num_steps)

            self.save_trajectory_samples()

            traj_dataset = TrajectoryDataset(self.save_dir, self.rspo_iter)
            assert len(
                traj_dataset) % self.all_args.n_rollout_threads == 0, len(
                    traj_dataset)
            assert not self.all_args.use_recurrent_policy and not self.all_args.use_naive_recurrent_policy
            traj_loader = DataLoader(
                traj_dataset,
                batch_size=self.all_args.episode_length *
                self.all_args.n_rollout_threads,
                shuffle=True,
            )
            d = Discriminator(self.all_args,
                              self.share_observation_space,
                              self.action_space,
                              device=self.device).to(self.device)
            for p in d.parameters():
                p.data.clamp_(-0.01, 0.01)

            self.ref_policies.append(self.trainer.policy)
            self.traj_loaders.append(traj_loader)
            self.traj_generators.append(iter(traj_loader))
            self.discriminators.append(d)
            self.discriminator_opts.append(
                torch.optim.RMSprop(d.parameters(), lr=self.all_args.d_lr))

            self.save()

Runner class to perform training, evaluation. and data collection for the MPEs. See parent class for details. # policy network # algorithm # buffer # initialize lagerangian lambda # prepare positive batch for wasserstein distance # Sample actions # Obser reward and next obs # compute divergence reward at the current step # NOTE: div_ret should not be immediately reset in env.reset # divs = [] # if self.rspo_iter > 0: #     pi = self.trainer.policy.get_pi( #         obs, rnn_states, masks).probs #     ref_pis = [ #         policy.get_pi(obs, rnn_states, masks).probs #         for policy in self.ref_policies #     ] #     div = torch.stack( #         [js_div(pi, ref_pi) for ref_pi in ref_pis], -1) #     divs = np.array( #         np.split(_t2n(div), self.n_rollout_threads)) # else: #     divs = [None for _ in range(obs.shape[0])] # insert data into buffer # update wasserstein discriminator # update lagrange lambda # self.logger.info((self.ll, sum(total_div) / len(total_div))) # compute return and update policy # post process # save model # log information # eval