# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/07_algorithms.ipynb (unless otherwise specified).

__all__ = ['PPO']

# Cell
import numpy as np
import gym
import torch
import torch.nn as nn
import torch.nn.functional as F
from rl_bolts import neuralnets as nns
from rl_bolts import losses as l
from .env_wrappers import BestPracticesWrapper, ToTorchWrapper, StateNormalizeWrapper
from .buffers import PGBuffer
from .datasets import PolicyGradientRLDataset
from .loops import polgrad_interaction_loop
import rl_bolts.utils as utils
import pytorch_lightning as pl
from argparse import Namespace
from typing import Optional, Union

# Cell
class PPO(pl.LightningModule):
    """
    Implementation of the Proximal Policy Optimization (PPO) algorithm. See the paper: https://arxiv.org/abs/1707.06347

    It is a PyTorch Lightning Module. See their docs: https://pytorch-lightning.readthedocs.io/en/latest/

    Args:
    - env (str): Environment to run in. Handles vector observation environments with either gym.spaces.Box or gym.spaces.Discrete
    action space.
    - hidden_sizes (tuple): Hidden layer sizes for actor-critic network.
    - gamma (float): Discount factor.
    - lam (float): Lambda factor for GAE-Lambda calculation.
    - clipratio (float): Clip ratio for PPO-clip objective.
    - train_iters (int): How many steps to take over the latest data batch.
    - batch_size (int): How many interactions to collect per update.
    - pol_lr (float): Learning rate for the policy optimizer.
    - val_lr (float): Learning rate for the value optimizer.
    - maxkl (float): Max allowed KL divergence between policy updates.
    - seed (int): Random seed for pytorch and numpy
    - evaluate (bool): Whether to run eval episodes at the end of each epoch. Saves episodes using gym.wrappers.Monitor.
    - monitor_dir (str): Directory for monitor to write to. Default is /tmp
    """
    def __init__(
        self,
        env: str,
        hidden_sizes: Optional[tuple] = (32, 32),
        gamma: Optional[float] = 0.99,
        lam: Optional[float] = 0.97,
        clipratio: Optional[float] = 0.2,
        train_iters: Optional[int] = 80,
        batch_size: Optional[int] = 4000,
        pol_lr: Optional[float] = 3e-4,
        val_lr: Optional[float] = 1e-3,
        maxkl: Optional[float] = 0.01,
        seed: Optional[int] = 0,
        evaluate: Optional[bool] = True,
        monitor_dir: Optional[str] = 'video_results'
    ):
        super().__init__()

        np.random.seed(seed)
        torch.manual_seed(seed)


        hparams = Namespace(
             **{
                'env':env,
                'hidden_sizes':hidden_sizes,
                'gamma':gamma,
                'lam':lam,
                'clipratio':clipratio,
                'train_iters':train_iters,
                'batch_size':batch_size,
                'pol_lr':pol_lr,
                'val_lr':val_lr,
                'maxkl':maxkl
             }
        )

        self.hparams = hparams

        env = gym.make(env)
        self.env = ToTorchWrapper(env)

        self.actor_critic = nns.ActorCritic(
            self.env.observation_space.shape[0],
            self.env.action_space,
            hidden_sizes=hidden_sizes,
        )

        self.gamma = gamma
        self.clipratio = clipratio
        self.train_iters = train_iters
        self.batch_size = batch_size
        self.pol_lr = pol_lr
        self.val_lr = val_lr
        self.maxkl = maxkl
        self.evaluate = evaluate

        if self.evaluate:
            eval_env = gym.wrappers.Monitor(env, monitor_dir, force=True)
            self.eval_env = ToTorchWrapper(eval_env)

        self.tracker_dict = {}

        self.buffer = PGBuffer(
            self.env.observation_space.shape,
            self.env.action_space.shape,
            size = self.batch_size,
            gamma = self.gamma
        )

        self.inner_loop()

    def configure_optimizers(self):
        self.policy_optimizer = torch.optim.Adam(self.actor_critic.policy.parameters(), lr=self.pol_lr)
        self.value_optimizer = torch.optim.Adam(self.actor_critic.value_f.parameters(), lr=self.val_lr)
        return self.policy_optimizer, self.value_optimizer

    def forward(self, x, a = None):
        out = self.actor_critic(x, a)
        return out

    def training_step(self, batch, batch_idx, optimizer_idx):
        states, actions, advs, rets, logps_old = batch

        if optimizer_idx == 0:
            stops = 0
            stopslst = []
            policy, logps = self.actor_critic.policy(states, actions)
            pol_loss_old, kl = l.ppo_clip_policy_loss(logps, logps_old, advs, clipratio=self.clipratio)

            for i in range(self.train_iters):
                self.policy_optimizer.zero_grad()
                policy, logps = self.actor_critic.policy(states, actions)
                pol_loss, kl = l.ppo_clip_policy_loss(logps, logps_old, advs, clipratio=self.clipratio)
                if kl > 1.5 * self.maxkl:
                    stops += 1
                    stopslst.append(i)
                    break
                pol_loss.backward()
                self.policy_optimizer.step()

            log = {
                "PolicyLoss": pol_loss_old.item(),
                "DeltaPolLoss": (pol_loss - pol_loss_old).item(),
                "KL": kl,
                "Entropy": policy.entropy().mean().item(),
                "TimesEarlyStopped": stops,
                "AvgEarlyStopStep": np.mean(stopslst) if len(stopslst) > 0 else 0
            }
            loss = pol_loss_old

        elif optimizer_idx == 1:
            values_old = self.actor_critic.value_f(states)
            val_loss_old = l.actor_critic_value_loss(values_old, rets)
            for i in range(self.train_iters):
                self.value_optimizer.zero_grad()
                values = self.actor_critic.value_f(states)
                val_loss = l.actor_critic_value_loss(values, rets)
                val_loss.backward()
                self.value_optimizer.step()

            delta_val_loss = (val_loss - val_loss_old).item()
            log = {"ValueLoss": val_loss_old.item(), "DeltaValLoss": delta_val_loss}
            loss = val_loss

        self.tracker_dict.update(log)
        log.update(self.tracker_dict)
        return {"loss": loss, "log": log, "progress_bar": log}

    def inner_loop(self) -> None:
        buffer, infos, _ = polgrad_interaction_loop(self.env, self.actor_critic, self.buffer, self.batch_size)
        self.data = buffer.get()
        self.tracker_dict.update(infos)

    def on_epoch_end(self):
        utils.printdict(self.tracker_dict)
        self.tracker_dict = {}
        self.inner_loop()
        self.eval_episodes(n_episodes=1)

    def train_dataloader(self):
        dataset = PolicyGradientRLDataset(self.data)
        dataloader = torch.utils.data.DataLoader(dataset, batch_size=self.batch_size, sampler=None, num_workers=4)
        return dataloader

    def backward(self, *args, **kwargs):
        pass

    def eval_episodes(self, n_episodes = 3):
        if self.evaluate:
            with torch.no_grad():
                rets_lst = []
                lens_lst = []
                for i in range(n_episodes):
                    episode_return = 0
                    episode_length = 0
                    obs = self.eval_env.reset()
                    done = False
                    while not done:
                        action, logp, value = self.actor_critic.step(obs)

                        obs, r, done, _ = self.eval_env.step(action)
                        episode_return += r
                        episode_length += 1

                        if done:
                            rets_lst.append(episode_return)
                            lens_lst.append(episode_length)

                            episode_return = 0
                            episode_length = 0
                            obs = self.eval_env.reset()

                dct = {
                    "NumEvalEpisodes": n_episodes,
                    "MeanEvalEpReturn": np.mean(rets_lst),
                    "StdEvalEpReturn": np.std(rets_lst),
                    "MeanEvalEpLength": np.mean(lens_lst),
                    "StdEvalEpLength": np.std(lens_lst)
                }

            self.tracker_dict.update(dct)
# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/07_algorithms.ipynb (unless otherwise specified).

__all__ = ['PPO']

# Cell
import numpy as np
import gym
import torch
import torch.nn as nn
import torch.nn.functional as F
from rl_bolts import neuralnets as nns
from rl_bolts import losses as l
from .env_wrappers import BestPracticesWrapper, ToTorchWrapper, StateNormalizeWrapper
from .buffers import PGBuffer
from .datasets import PolicyGradientRLDataset
from .loops import polgrad_interaction_loop
import rl_bolts.utils as utils
import pytorch_lightning as pl
from argparse import Namespace
from typing import Optional, Union

# Cell
class PPO(pl.LightningModule):
    """
    Implementation of the Proximal Policy Optimization (PPO) algorithm. See the paper: https://arxiv.org/abs/1707.06347

    It is a PyTorch Lightning Module. See their docs: https://pytorch-lightning.readthedocs.io/en/latest/

    Args:
    - env (str): Environment to run in. Handles vector observation environments with either gym.spaces.Box or gym.spaces.Discrete
    action space.
    - hidden_sizes (tuple): Hidden layer sizes for actor-critic network.
    - gamma (float): Discount factor.
    - lam (float): Lambda factor for GAE-Lambda calculation.
    - clipratio (float): Clip ratio for PPO-clip objective.
    - train_iters (int): How many steps to take over the latest data batch.
    - batch_size (int): How many interactions to collect per update.
    - pol_lr (float): Learning rate for the policy optimizer.
    - val_lr (float): Learning rate for the value optimizer.
    - maxkl (float): Max allowed KL divergence between policy updates.
    - seed (int): Random seed for pytorch and numpy
    - evaluate (bool): Whether to run eval episodes at the end of each epoch. Saves episodes using gym.wrappers.Monitor.
    - monitor_dir (str): Directory for monitor to write to. Default is /tmp
    """
    def __init__(
        self,
        env: str,
        hidden_sizes: Optional[tuple] = (32, 32),
        gamma: Optional[float] = 0.99,
        lam: Optional[float] = 0.97,
        clipratio: Optional[float] = 0.2,
        train_iters: Optional[int] = 80,
        batch_size: Optional[int] = 4000,
        pol_lr: Optional[float] = 3e-4,
        val_lr: Optional[float] = 1e-3,
        maxkl: Optional[float] = 0.01,
        seed: Optional[int] = 0,
        evaluate: Optional[bool] = True,
        monitor_dir: Optional[str] = 'video_results'
    ):
        super().__init__()

        np.random.seed(seed)
        torch.manual_seed(seed)


        hparams = Namespace(
             **{
                'env':env,
                'hidden_sizes':hidden_sizes,
                'gamma':gamma,
                'lam':lam,
                'clipratio':clipratio,
                'train_iters':train_iters,
                'batch_size':batch_size,
                'pol_lr':pol_lr,
                'val_lr':val_lr,
                'maxkl':maxkl
             }
        )

        self.hparams = hparams

        env = gym.make(env)
        self.env = ToTorchWrapper(env)

        self.actor_critic = nns.ActorCritic(
            self.env.observation_space.shape[0],
            self.env.action_space,
            hidden_sizes=hidden_sizes,
        )

        self.gamma = gamma
        self.clipratio = clipratio
        self.train_iters = train_iters
        self.batch_size = batch_size
        self.pol_lr = pol_lr
        self.val_lr = val_lr
        self.maxkl = maxkl
        self.evaluate = evaluate

        if self.evaluate:
            eval_env = gym.wrappers.Monitor(env, monitor_dir, force=True)
            self.eval_env = ToTorchWrapper(eval_env)

        self.tracker_dict = {}

        self.buffer = PGBuffer(
            self.env.observation_space.shape,
            self.env.action_space.shape,
            size = self.batch_size,
            gamma = self.gamma
        )

        self.inner_loop()

    def configure_optimizers(self):
        self.policy_optimizer = torch.optim.Adam(self.actor_critic.policy.parameters(), lr=self.pol_lr)
        self.value_optimizer = torch.optim.Adam(self.actor_critic.value_f.parameters(), lr=self.val_lr)
        return self.policy_optimizer, self.value_optimizer

    def forward(self, x, a = None):
        out = self.actor_critic(x, a)
        return out

    def training_step(self, batch, batch_idx, optimizer_idx):
        states, actions, advs, rets, logps_old = batch

        if optimizer_idx == 0:
            stops = 0
            stopslst = []
            policy, logps = self.actor_critic.policy(states, actions)
            pol_loss_old, kl = l.ppo_clip_policy_loss(logps, logps_old, advs, clipratio=self.clipratio)

            for i in range(self.train_iters):
                self.policy_optimizer.zero_grad()
                policy, logps = self.actor_critic.policy(states, actions)
                pol_loss, kl = l.ppo_clip_policy_loss(logps, logps_old, advs, clipratio=self.clipratio)
                if kl > 1.5 * self.maxkl:
                    stops += 1
                    stopslst.append(i)
                    break
                pol_loss.backward()
                self.policy_optimizer.step()

            log = {
                "PolicyLoss": pol_loss_old.item(),
                "DeltaPolLoss": (pol_loss - pol_loss_old).item(),
                "KL": kl,
                "Entropy": policy.entropy().mean().item(),
                "TimesEarlyStopped": stops,
                "AvgEarlyStopStep": np.mean(stopslst) if len(stopslst) > 0 else 0
            }
            loss = pol_loss_old

        elif optimizer_idx == 1:
            values_old = self.actor_critic.value_f(states)
            val_loss_old = l.actor_critic_value_loss(values_old, rets)
            for i in range(self.train_iters):
                self.value_optimizer.zero_grad()
                values = self.actor_critic.value_f(states)
                val_loss = l.actor_critic_value_loss(values, rets)
                val_loss.backward()
                self.value_optimizer.step()

            delta_val_loss = (val_loss - val_loss_old).item()
            log = {"ValueLoss": val_loss_old.item(), "DeltaValLoss": delta_val_loss}
            loss = val_loss

        self.tracker_dict.update(log)
        log.update(self.tracker_dict)
        return {"loss": loss, "log": log, "progress_bar": log}

    def inner_loop(self) -> None:
        buffer, infos, _ = polgrad_interaction_loop(self.env, self.actor_critic, self.buffer, self.batch_size)
        self.data = buffer.get()
        self.tracker_dict.update(infos)

    def on_epoch_end(self):
        utils.printdict(self.tracker_dict)
        self.tracker_dict = {}
        self.inner_loop()
        self.eval_episodes(n_episodes=1)

    def train_dataloader(self):
        dataset = PolicyGradientRLDataset(self.data)
        dataloader = torch.utils.data.DataLoader(dataset, batch_size=self.batch_size, sampler=None, num_workers=4)
        return dataloader

    def backward(self, *args, **kwargs):
        pass

    def eval_episodes(self, n_episodes = 3):
        if self.evaluate:
            with torch.no_grad():
                rets_lst = []
                lens_lst = []
                for i in range(n_episodes):
                    episode_return = 0
                    episode_length = 0
                    obs = self.eval_env.reset()
                    done = False
                    while not done:
                        action, logp, value = self.actor_critic.step(obs)

                        obs, r, done, _ = self.eval_env.step(action)
                        episode_return += r
                        episode_length += 1

                        if done:
                            rets_lst.append(episode_return)
                            lens_lst.append(episode_length)

                            episode_return = 0
                            episode_length = 0
                            obs = self.eval_env.reset()

                dct = {
                    "NumEvalEpisodes": n_episodes,
                    "MeanEvalEpReturn": np.mean(rets_lst),
                    "StdEvalEpReturn": np.std(rets_lst),
                    "MeanEvalEpLength": np.mean(lens_lst),
                    "StdEvalEpLength": np.std(lens_lst)
                }

            self.tracker_dict.update(dct)
# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/07_algorithms.ipynb (unless otherwise specified). # Cell # Cell Implementation of the Proximal Policy Optimization (PPO) algorithm. See the paper: https://arxiv.org/abs/1707.06347      It is a PyTorch Lightning Module. See their docs: https://pytorch-lightning.readthedocs.io/en/latest/      Args:     - env (str): Environment to run in. Handles vector observation environments with either gym.spaces.Box or gym.spaces.Discrete     action space.     - hidden_sizes (tuple): Hidden layer sizes for actor-critic network.     - gamma (float): Discount factor.     - lam (float): Lambda factor for GAE-Lambda calculation.     - clipratio (float): Clip ratio for PPO-clip objective.     - train_iters (int): How many steps to take over the latest data batch.     - batch_size (int): How many interactions to collect per update.     - pol_lr (float): Learning rate for the policy optimizer.     - val_lr (float): Learning rate for the value optimizer.     - maxkl (float): Max allowed KL divergence between policy updates.     - seed (int): Random seed for pytorch and numpy     - evaluate (bool): Whether to run eval episodes at the end of each epoch. Saves episodes using gym.wrappers.Monitor.     - monitor_dir (str): Directory for monitor to write to. Default is /tmp