<filename>rl/algorithms/mbrl/mbppo_old.py
import os, subprocess, sys
import argparse
import importlib
import datetime
import random

import time
import wandb

import numpy as np
import torch as T
import torch.nn.functional as F

# T.multiprocessing.set_sharing_strategy('file_system')

from rl.algorithms.mbrl.mbrl import MBRL
from rl.algorithms.mfrl.ppo import PPO
from rl.world_models.fake_world import FakeWorld
import rl.environments.mbpo.static as mbpo_static
# from rl.data.dataset import RLDataModule





class MBPPO(MBRL, PPO):
    """
    Greek alphabet:
        Α α, Β β, Γ γ, Δ δ, Ε ε, Ζ ζ, Η η, Θ θ, Ι ι, Κ κ, Λ λ, Μ μ,
        Ν ν, Ξ ξ, Ο ο, Π π, Ρ ρ, Σ σ/ς, Τ τ, Υ υ, Φ φ, Χ χ, Ψ ψ, Ω ω.

    Algorithm: Model-Based Game: Policy As Leader (PAL) – Practical Version

        01: Initialize: Models parameters( Policy net πφ0, Value net Vψ0, ensemble of MDP world models {M^θ0}_{1:nM} )
        02. Initialize: Replay buffer D.
        03: Hyperparameters: Initial samples Ni = 2 epochs, samples per update E=1000, buffer size B = 2500 ≈ N, number of NPG steps K = 4 ≈ 1
        04: Initial Data: Collect N_init samples from the environment by interacting with initial policy. Store data in buffer D.
        05: for n = 0, 1, 2, ..., N (nEpochs) do
        06:    Learn dynamics model(s) M^_n+1 using data in the buffer.
        07:    Policy updates: π_n+1; V_n+1 = MB-PPO(π_n, V_n, M^_n+1) // call K times
        08:    Collect dataset of N samples from World by interacting with πn+1.
        09.    Add data to replay buffer D, discarding old data if size is larger than B.
        10: end for


    Algorithm: Model-Based Game: Model As Leader (MAL) – Practical Version
        1: Initialize: Policy network π_0, model network(s) Mhat_0, value network V_0.
        2: Hyperparameters: Initial samples N_init, samples per update N, number of PPO steps K >> 1
        3: Initial Data: Collect N_init samples from the environment by interacting with initial policy. Store data in buffer D.
        4: Initial Model: Learn model(s) Mhat_0 using data in D.
        5: for k = 0, 1, 2, ... do
        6:    Optimize π_k+1 using Mck by running K >> 1 steps of model-based PPO (Subroutine 1).
        7:    Collect dataset D_k+1 of N samples from world using πk+1. Aggregate data D = D U D_k+1.
        8:    Learn dynamics model(s) Mhat_k+1 using data in D.
        9: end for


    Algorithm: Model-Based Proximal Policy Optimization (On-Policy, Dyna-style, Model-Based)

        01. Inputs: Models parameters( Policy net πφn, Value net Vψn, ensemble of MDP world models {M^θn+1}_{1:nM}, Replay buffer D)
        02. Hyperparameters: Disc. factor γ, GAE λ, num traj's Nτ, model rollout horizon H
        03. Initialize: Trajectory buffer Dτ = {}
        04. for k = 0, 1, 2, ..., Nτ:
        05.     Sample init p(s^k_0) from distribution/buffer
        05.     Rollout a traj {τ^k_π/M} from {M^θ}_{1:nM} by πn = π(φn) for e = 0, 1, 2, ..., H
        06.     Aggregate traj τ^k in the traj buffer, Dτ = Dτ U {τ^k_π/M}
        07.     Compute RTG R^_t, GAE A^_t based on Vn = V(θn)
        08. end for
        09. Update πφ by maxz Jπ
                φ = arg max_φ {(1/T|Dk|) sum sum min((π/πk), 1 +- eps) Aπk }
        10. Fit Vθ by MSE(Jv)
                θ = arg min_θ {(1/T|Dk|) sum sum (Vθ(st) - RTG)^2 }
        11. Return: Policy net πθn+1, value net Vφn+1

    Subroutine 1: Model-Based Natural Policy Gradient Update Step

        1: Require: Policy (stochastic) network πθ, value/baseline network V , ensemble of MDP dynamics models fMcφg,
        reward function R, initial state distribution or buffer.
        2: Hyperparameters: Discount factor γ, GAE λ, number of trajectories Nτ, rollout horizon H, normalized NPG step
        size δ
        3: Initialize trajectory buffer Dτ = fg
        4: for k = 1; 2; : : : ; Nτ do
        5: Sample initial state sk 0 from initial state distribution/buffer
        6: Perform H step rollout from sk 0 with πθ to get τjk = (sk 0; ak 0; sk 1; ak 2; : : : sk H; ak H), one for each model Mcφj in the
        ensemble.
        7: Query reward function to obtain rewards for each step of the trajectories
        8: Truncate trajectories if termination/truncation conditions are part of the environment
        9: Aggregate the trajectories in trajectory buffer, Dτ = Dτ [ fτg
        10: end for
        11: Compute advantages for each trajectory using V and GAE (Schulman et al., 2016).
        12: Compute vanilla policy gradient using the dataset
        g = E(s;a)∼Dτ [rθ log πθ(ajs)Aπ(s; a)]
        13: Perform normalized NPG update (F denotes the Fisher matrix)
        θ = θ + sgT Fδ−1g F −1g
        14: Update value/baseline network V to fit the computed returns in Dτ.
        15: Return Policy network πθ, value network V

    """
    def __init__(self, exp_prefix, configs, seed, device, wb) -> None:
        super(MBPPO, self).__init__(exp_prefix, configs, seed, device)
        # print('init MBPPO Algorithm!')
        self.configs = configs
        self.seed = seed
        self._device_ = device
        self.WandB = wb
        self._build()


    def __init__(self, exp_prefix, configs, seed, device, wb) -> None:
        super(MBPPO, self).__init__(exp_prefix, configs, seed, device)
        # print('init MBPPO Algorithm!')
        self.configs = configs
        self.seed = seed
        self._device_ = device
        self.WandB = wb
        self._build()


    ## build MBPPO components: (env, D, AC, alpha)
    def _build(self):
        super(MBPPO, self)._build()
        self._set_ppo()
        # self._set_fake_world()


    ## PPO
    def _set_ppo(self):
        PPO._build_ppo(self)


    ## FakeEnv
    def _set_fake_world(self):
        env_name = self.configs['environment']['name']
        device = self._device_
        if self.configs['environment']['name'][:4] == 'pddm':
        	static_fns = None
        else:
        	static_fns = mbpo_static[env_name[:-3].lower()]

        self.fake_world = FakeWorld(self.world_model)


    def learn(self):
        N = self.configs['algorithm']['learning']['epochs']
        NT = self.configs['algorithm']['learning']['epoch_steps']
        Ni = self.configs['algorithm']['learning']['init_epochs']
        Nx = self.configs['algorithm']['learning']['expl_epochs']

        E = self.configs['algorithm']['learning']['env_steps']
        G_ppo = self.configs['algorithm']['learning']['grad_PPO_steps']

        model_train_frequency = self.configs['world_model']['model_train_freq']
        batch_size_m = self.configs['world_model']['network']['batch_size'] # bs_m
        wm_epochs = self.configs['algorithm']['learning']['grad_WM_steps']
        real_ratio = self.configs['data']['real_ratio'] # rr
        batch_size = self.configs['data']['batch_size'] # bs
        mini_batch_size = self.configs['data']['mini_batch_size'] # bs
        rollout_trajectories = self.configs['data']['rollout_trajectories'] # bs_ro
        K = self.configs['data']['rollout_horizon']

        global_step = 0
        start_time = time.time()
        o, Z, el, t = self.learn_env.reset(), 0, 0, 0
        oldJs = [0, 0]
        JVList, JPiList = [0]*Ni, [0]*Ni
        logs = dict()
        lastEZ, lastES = 0, -2

        start_time_real = time.time()
        for n in range(1, N+1):
            if self.configs['experiment']['print_logs']:
                print('=' * 50)
                if n > Nx:
                    print(f'\n[ Epoch {n}   Learning ]'+(' '*50))
                    oldJs = [0, 0]
                    JVList, JPiList, KLList = [0], [0], [0]
                    JValList = [0]
                elif n > Ni:
                    print(f'\n[ Epoch {n}   Exploration + Learning ]'+(' '*50))
                    JVList, JPiList, KLList = [], [], []
                    JValList = []
                else:
                    print(f'\n[ Epoch {n}   Inintial Exploration ]'+(' '*50))
                    oldJs = [0, 0]
                    JVList, JPiList, KLList = [0], [0], [0]
                    JValList = [0]

            # print(f'[ Replay Buffer ] Size: {self.buffer.total_size()}')
            nt = 0
            # self.buffer.clean_buffer()
            o, d, Z, el, = self.learn_env.reset(), 0, 0, 0

            learn_start_real = time.time()
            while nt < NT: # full epoch
                # Interaction steps
                for e in range(1, E+1):
                    # o, Z, el, t = self.internact(n, o, Z, el, t)
                    o, Z, el, t = self.internact_opB(n, o, Z, el, t)
                    print(f'[ Epoch {n}   Interaction ] Env Steps: {e} | Return: {round(Z, 2)}'+(" "*10), end='\r')
                with T.no_grad(): v = self.actor_critic.get_v(T.Tensor(o)).cpu()
                # self.buffer.traj_tail(d, v, el)
                self.buffer.finish_path(el, v)

                # print(f'Buffer size = {self.buffer.total_size()}')
                # Taking gradient steps after exploration
                if n > Ni:
                    # if nt % model_train_frequency == 0:
                    # 03. Train model pθ on Denv via maximum likelihood
                    # print(f'\n[ Epoch {n}   Training World Model ]'+(' '*50))

                    ho_mean = 0. #self.fake_world.train_fake_world(self.buffer)
                    JValList.append(ho_mean) # ho: holdout

                    # model_fit_bs = self.configs['data']['buffer_size']
                    # model_fit_batch = self.buffer.sample_batch(model_fit_bs, self._device_)
                    # s, a, sp, _, _, _, _ = model_fit_batch.values()
                    # if n == Ni+1:
                    #     samples_to_collect = 4000
                    # else:
                    #     samples_to_collect = 1000

                    # for i, model in enumerate(self.models):
                    #     print(f'\n[ Epoch {n}   Training World Model {i+1} ]'+(' '*50))
                    #     loss_general = model.compute_loss(s[-samples_to_collect:],
                    #                                       a[-samples_to_collect:],
                    #                                       sp[-samples_to_collect:]) # generalization error
                    #     dynamics_loss = model.fit_dynamics(s, a, sp, **job_data)
                    #     logger.log_kv('dyn_loss_' + str(i), dynamics_loss[-1])
                    #     logger.log_kv('dyn_loss_gen_' + str(i), loss_general)
                    #     if job_data['learn_reward']:
                    #         reward_loss = model.fit_reward(s, a, r.reshape(-1, 1), **job_data)
                    #         logger.log_kv('rew_loss_' + str(i), reward_loss[-1])

                    KstepsList = []
                    self.initialize_model_buffer()
                    # PPO-P >>>>
                    for g in range(1, G_ppo+1):
                        print('\n')
                        # print(f'Actor-Critic Grads...{g}', end='\r')
                        # print(f'[ Epoch {n}   Training Actor-Critic ] AC Grads: {g}'+(" "*50), end='\r')
                        # # Reallocate model buffer
                        # self.initialize_model_buffer()
                        self.model_buffer.reset()
                        # # Generate M k-steps imaginary rollouts for SAC traingin
                        # self.rollout_world_model_op(rollout_trajectories, K, n)
                        # k_avg = self.rollout_world_model_trajectories(rollout_trajectories, K, n)
                        k_avg = self.rollout_real_world_trajectories(rollout_trajectories, K, g, n)
                        KstepsList.append(k_avg)
                        batch_size = int(self.model_buffer.total_size())
                        print('\n')
                        stop_pi = False
                        kl = 0
                        for gg in range(1, 81): # 101
                            print(f'[ Epoch {n} ] AC: {g} | ac: {gg} || stopPG={stop_pi} | KL={round(kl, 4)}'+(' '*40), end='\r')
                            batch = self.model_buffer.sample_batch(batch_size, self._device_)
                            Jv, Jpi, kl, stop_pi = self.trainAC(g, batch, oldJs)
                            # print(f'[ Epoch {n} ] AC: {g} | ac: {gg} || PG={stop_pi} | KL={kl}'+(' '*80), end='\r')
                            oldJs = [Jv, Jpi]
                            JVList.append(Jv.item())
                            JPiList.append(Jpi.item())
                            KLList.append(kl)
                    # PPO-P <<<<

                nt += E

            print('\n')
            # logs['time/training                  '] = time.time() - learn_start_real

            # logs['training/wm/Jtrain_mean        '] = np.mean(JMeanTrainList)
            # logs['training/wm/Jtrain             '] = np.mean(JTrainList)
            logs['training/wm/Jval               '] = np.mean(JValList)
            # logs['training/wm/test_mse           '] = np.mean(LossTestList)

            logs['training/ppo/Jv                '] = np.mean(JVList)
            logs['training/ppo/Jpi               '] = np.mean(JPiList)
            logs['training/ppo/KL                '] = np.mean(KLList)

            logs['data/env_buffer                '] = self.buffer.total_size()
            logs['data/env_rollout_steps         '] = self.buffer.average_horizon()
            if hasattr(self, 'model_buffer'):
                logs['data/model_buffer              '] = self.model_buffer.total_size()
                logs['data/model_rollout_steps       '] = np.mean(KstepsList)
            else:
                logs['data/model_buffer              '] = 0.
                logs['data/model_rollout_steps       '] = 0.
            # else:
            #     logs['data/model_buffer              '] = 0
            # logs['data/rollout_length            '] = K

            eval_start_real = time.time()
            EZ, ES, EL = self.evaluate()
            # EZ, ES, EL = self.evaluate_op()

            # logs['time/evaluation                '] = time.time() - eval_start_real

            if self.configs['environment']['type'] == 'mujoco-pddm-shadowhand':
                logs['evaluation/episodic_score_mean '] = np.mean(ES)
                logs['evaluation/episodic_score_std  '] = np.std(ES)
            else:
                logs['evaluation/episodic_return_mean'] = np.mean(EZ)
                logs['evaluation/episodic_return_std '] = np.std(EZ)
            logs['evaluation/episodic_length_mean'] = np.mean(EL)

            logs['time/total                     '] = time.time() - start_time_real

            # if n > (N - 50):
            #     if self.configs['environment']['type'] == 'mujoco-pddm-shadowhand':
            #         if np.mean(ES) > lastES:
            #             print(f'[ Epoch {n}   Agent Saving ]                    ')
            #             env_name = self.configs['environment']['name']
            #             alg_name = self.configs['algorithm']['name']
            #             T.save(self.actor_critic.actor,
            #             f'./saved_agents/{env_name}-{alg_name}-seed:{self.seed}-epoch:{n}.pTtar')
            #             lastES = np.mean(ES)
            #     else:
            #         if np.mean(EZ) > lastEZ:
            #             print(f'[ Epoch {n}   Agent Saving ]                    ')
            #             env_name = self.configs['environment']['name']
            #             alg_name = self.configs['algorithm']['name']
            #             T.save(self.actor_critic.actor,
            #             f'./saved_agents/{env_name}-{alg_name}-seed:{self.seed}-epoch:{n}.pTtar')
            #             lastEZ = np.mean(EZ)

            # Printing logs
            if self.configs['experiment']['print_logs']:
                for k, v in logs.items():
                    print(f'{k}  {round(v, 4)}'+(' '*10))

            # WandB
            if self.WandB:
                wandb.log(logs)

        self.learn_env.close()
        self.eval_env.close()


    def rollout_real_world_trajectories(self, rollout_trajectories, K, g, n):
    	# 07. Sample st uniformly from Denv
    	device = self._device_
    	# Nτ = 400 * 2 # number of trajectories x number of models
    	Nτ = 40 # number of trajectories x number of models
    	# Nτ = self.configs['data']['rollout_trajectories']
    	# o, d, el = self.traj_env.reset(), 0, 0

        # 08. Perform k-step model rollout starting from st using policy πφ; add to Dmodel
    	k_end_total = 0
    	for nτ in range(1, Nτ+1): # Generate trajectories
            # o, d, el = self.traj_env.reset(), 0, 0
            o, Z, el = self.traj_env.reset(), 0, 0
            # print(f'nτ = {nτ}'+(' '*70), end='\r')
            for k in range(1, K+1): # Generate rollouts
                print(f'[ Epoch {n} ] Model Rollout: nτ = {nτ} | k = {k} | Buffer = {self.model_buffer.total_size()} | Return = {round(Z, 2)}', end='\r')
                # print(f'[ Epoch {n} ] AC Training Grads: {g} || Model Rollout: nτ = {nτ} | k = {k} | Buffer size = {self.model_buffer.total_size()}'+(' '*10))
                with T.no_grad(): a, log_pi, v = self.actor_critic.get_a_and_v_np(T.Tensor(o))
                # o_next, r, d_next, _ = self.traj_env.step(a)
                o_next, r, d, _ = self.traj_env.step(a)
                Z += r
                el += 1

                # self.model_buffer.store_transition(o, a, r, d, v, log_pi)
                self.model_buffer.store(o, a, r, o_next, v, log_pi, el)
                o = o_next
                # if d_next or (el == K):
                #     o_next, d_next, el = self.traj_env.reset(), 0, 0
                #     print(f'[ Model Rollout ] Breaking early: {k}'+(' '*50), end='\r')
                #     # k_end_total += k
                #     break
                if d or (el == K):
                    # print(f'[ Model Rollout ] Breaking early: {k}'+(' '*50), end='\r')
                    # if el == K:
                    #     with T.no_grad(): v = self.actor_critic.get_v(T.Tensor(o)).cpu()
                    # else:
                    #     v = T.Tensor([0.0])
                    # self.model_buffer.finish_path(el, v)
                    break
                    # o, el = self.traj_env.reset(), 0
                # else:
                # o, d = o_next, d_next
            if el == K:
                # print(f'[ Model Rollout ] Average Breaking : {k_end_total//Nτ}'+(' '*50))
                with T.no_grad(): v = self.actor_critic.get_v(T.Tensor(o)).cpu()
            else:
                v = T.Tensor([0.0])
            # self.model_buffer.traj_tail(d, v)
            self.model_buffer.finish_path(el, v)
            k_end_total += k
    	return k_end_total//Nτ


    def rollout_world_model_trajectories(self, rollout_trajectories, K, n):
    	# 07. Sample st uniformly from Denv
    	device = self._device_
    	batch_size = rollout_trajectories
    	B_ro = self.buffer.sample_batch(batch_size) # Torch
    	O = B_ro['observations_next']
    	D = B_ro['terminals']

        # 08. Perform k-step model rollout starting from st using policy πφ; add to Dmodel
    	k_end_total = 0
    	for o, d in zip(O, D):
            # print('o: ', o.shape)
            # print('d: ', d.shape)
            for k in range(1, K+1):
                # print(f'k = {k}', end='\r')
                with T.no_grad(): a, log_pi, _, v = self.actor_critic.get_pi_and_v(o)
                a, log_pi, v = a.cpu(), log_pi.cpu(), v.cpu()

                o_next, r, d_next, _ = self.fake_world.step(o, a, deterministic=True) # ip: Tensor, op: Tensor

                self.model_buffer.store_transition(o, a, r, d, v, log_pi)

                if d_next:
    	            print(f'[ Model Rollout ] Breaking early: {k}'+(' '*50), end='\r')
    	            k_end_total += k
    	            break

                o, d = o_next, d_next

    	print(f'[ Model Rollout ] Average Breaking : {k_end_total//batch_size}'+(' '*50))
    	with T.no_grad(): v = self.actor_critic.get_v(o)
    	self.model_buffer.traj_tail(d, v)

    	return k_end_total//batch_size


    def rollout_world_model_op(self, rollout_trajectories, K, n):
    	# 07. Sample st uniformly from Denv
    	device = self._device_
    	batch_size = rollout_trajectories
    	# print(f'[ Epoch {n}   Model Rollout ] Batch Size: {batch_size} | Rollout Horizon: {K}'+(' '*50))
    	B_ro = self.buffer.sample_batch(batch_size) # Torch
    	O = B_ro['observations']
    	D = B_ro['terminals']

        # 08. Perform k-step model rollout starting from st using policy πφ; add to Dmodel
    	for k in range(1, K+1):
            # print('k = ', k)
            # print('k = ', k, end='\r')
            A = self.actor_critic.get_action(O) # Stochastic action | No reparameterization
            with T.no_grad(): A, log_Pi, _, V = self.actor_critic.get_pi_and_v(O)

            O_next, R, D_next, _ = self.fake_world.step(O, A) # ip: Tensor, op: Tensor
        	# O_next, R, D, _ = self.fake_world.step_np(O, A) # ip: Tensor, op: Numpy

            self.model_buffer.store_batch(O, A, R, D, V, log_Pi)

            O_next = T.Tensor(O_next)
            D = T.tensor(D, dtype=T.bool)
            nonD = ~D.squeeze(-1)
            if nonD.sum() == 0:
        	    print(f'[ Epoch {n}   Model Rollout ] Breaking early: {k} | {nonD.sum()} / {nonD.shape}')
        	    break

            O, D = O_next[nonD], D_next[nonD]

    	with T.no_grad(): V = self.actor_critic.get_v(O)
    	self.model_buffer.traj_tail(D, V)


    def ppo_mini_batch(self, mini_batch_size):
    	B = self.model_buffer.sample_batch(mini_batch_size, self._device_)
    	return B





def main(exp_prefix, config, seed, device, wb):

    print('Start an MBPPO experiment...')
    print('\n')

    configs = config.configurations

    if seed:
        random.seed(seed), np.random.seed(seed), T.manual_seed(seed)

    alg_name = configs['algorithm']['name']
    alg_mode = configs['algorithm']['mode']
    env_name = configs['environment']['name']
    env_type = configs['environment']['type']
    wm_epochs = configs['algorithm']['learning']['grad_WM_steps']
    DE = configs['world_model']['num_ensembles']

    group_name = f"{env_name}-{alg_name}-{alg_mode}-D"
    exp_prefix = f"seed:{seed}"

    if wb:
        wandb.init(
            name=exp_prefix,
            group=group_name,
            # project='test',
            project='AMMI-RL-2022',
            config=configs
        )

    agent = MBPPO(exp_prefix, configs, seed, device, wb)

    agent.learn()

    print('\n')
    print('... End the MBPPO experiment')


if __name__ == "__main__":

    import argparse

    parser = argparse.ArgumentParser()

    parser.add_argument('-exp_prefix', type=str)
    parser.add_argument('-cfg', type=str)
    parser.add_argument('-seed', type=str)
    parser.add_argument('-device', type=str)
    parser.add_argument('-wb', type=str)

    args = parser.parse_args()

    exp_prefix = args.exp_prefix
    sys.path.append(f"{os.getcwd()}/configs")
    config = importlib.import_module(args.cfg)
    seed = int(args.seed)
    device = args.device
    wb = eval(args.wb)

    main(exp_prefix, config, seed, device, wb)

<filename>rl/algorithms/mbrl/mbppo_old.py
import os, subprocess, sys
import argparse
import importlib
import datetime
import random

import time
import wandb

import numpy as np
import torch as T
import torch.nn.functional as F

# T.multiprocessing.set_sharing_strategy('file_system')

from rl.algorithms.mbrl.mbrl import MBRL
from rl.algorithms.mfrl.ppo import PPO
from rl.world_models.fake_world import FakeWorld
import rl.environments.mbpo.static as mbpo_static
# from rl.data.dataset import RLDataModule





class MBPPO(MBRL, PPO):
    """
    Greek alphabet:
        Α α, Β β, Γ γ, Δ δ, Ε ε, Ζ ζ, Η η, Θ θ, Ι ι, Κ κ, Λ λ, Μ μ,
        Ν ν, Ξ ξ, Ο ο, Π π, Ρ ρ, Σ σ/ς, Τ τ, Υ υ, Φ φ, Χ χ, Ψ ψ, Ω ω.

    Algorithm: Model-Based Game: Policy As Leader (PAL) – Practical Version

        01: Initialize: Models parameters( Policy net πφ0, Value net Vψ0, ensemble of MDP world models {M^θ0}_{1:nM} )
        02. Initialize: Replay buffer D.
        03: Hyperparameters: Initial samples Ni = 2 epochs, samples per update E=1000, buffer size B = 2500 ≈ N, number of NPG steps K = 4 ≈ 1
        04: Initial Data: Collect N_init samples from the environment by interacting with initial policy. Store data in buffer D.
        05: for n = 0, 1, 2, ..., N (nEpochs) do
        06:    Learn dynamics model(s) M^_n+1 using data in the buffer.
        07:    Policy updates: π_n+1; V_n+1 = MB-PPO(π_n, V_n, M^_n+1) // call K times
        08:    Collect dataset of N samples from World by interacting with πn+1.
        09.    Add data to replay buffer D, discarding old data if size is larger than B.
        10: end for


    Algorithm: Model-Based Game: Model As Leader (MAL) – Practical Version
        1: Initialize: Policy network π_0, model network(s) Mhat_0, value network V_0.
        2: Hyperparameters: Initial samples N_init, samples per update N, number of PPO steps K >> 1
        3: Initial Data: Collect N_init samples from the environment by interacting with initial policy. Store data in buffer D.
        4: Initial Model: Learn model(s) Mhat_0 using data in D.
        5: for k = 0, 1, 2, ... do
        6:    Optimize π_k+1 using Mck by running K >> 1 steps of model-based PPO (Subroutine 1).
        7:    Collect dataset D_k+1 of N samples from world using πk+1. Aggregate data D = D U D_k+1.
        8:    Learn dynamics model(s) Mhat_k+1 using data in D.
        9: end for


    Algorithm: Model-Based Proximal Policy Optimization (On-Policy, Dyna-style, Model-Based)

        01. Inputs: Models parameters( Policy net πφn, Value net Vψn, ensemble of MDP world models {M^θn+1}_{1:nM}, Replay buffer D)
        02. Hyperparameters: Disc. factor γ, GAE λ, num traj's Nτ, model rollout horizon H
        03. Initialize: Trajectory buffer Dτ = {}
        04. for k = 0, 1, 2, ..., Nτ:
        05.     Sample init p(s^k_0) from distribution/buffer
        05.     Rollout a traj {τ^k_π/M} from {M^θ}_{1:nM} by πn = π(φn) for e = 0, 1, 2, ..., H
        06.     Aggregate traj τ^k in the traj buffer, Dτ = Dτ U {τ^k_π/M}
        07.     Compute RTG R^_t, GAE A^_t based on Vn = V(θn)
        08. end for
        09. Update πφ by maxz Jπ
                φ = arg max_φ {(1/T|Dk|) sum sum min((π/πk), 1 +- eps) Aπk }
        10. Fit Vθ by MSE(Jv)
                θ = arg min_θ {(1/T|Dk|) sum sum (Vθ(st) - RTG)^2 }
        11. Return: Policy net πθn+1, value net Vφn+1

    Subroutine 1: Model-Based Natural Policy Gradient Update Step

        1: Require: Policy (stochastic) network πθ, value/baseline network V , ensemble of MDP dynamics models fMcφg,
        reward function R, initial state distribution or buffer.
        2: Hyperparameters: Discount factor γ, GAE λ, number of trajectories Nτ, rollout horizon H, normalized NPG step
        size δ
        3: Initialize trajectory buffer Dτ = fg
        4: for k = 1; 2; : : : ; Nτ do
        5: Sample initial state sk 0 from initial state distribution/buffer
        6: Perform H step rollout from sk 0 with πθ to get τjk = (sk 0; ak 0; sk 1; ak 2; : : : sk H; ak H), one for each model Mcφj in the
        ensemble.
        7: Query reward function to obtain rewards for each step of the trajectories
        8: Truncate trajectories if termination/truncation conditions are part of the environment
        9: Aggregate the trajectories in trajectory buffer, Dτ = Dτ [ fτg
        10: end for
        11: Compute advantages for each trajectory using V and GAE (Schulman et al., 2016).
        12: Compute vanilla policy gradient using the dataset
        g = E(s;a)∼Dτ [rθ log πθ(ajs)Aπ(s; a)]
        13: Perform normalized NPG update (F denotes the Fisher matrix)
        θ = θ + sgT Fδ−1g F −1g
        14: Update value/baseline network V to fit the computed returns in Dτ.
        15: Return Policy network πθ, value network V

    """
    def __init__(self, exp_prefix, configs, seed, device, wb) -> None:
        super(MBPPO, self).__init__(exp_prefix, configs, seed, device)
        # print('init MBPPO Algorithm!')
        self.configs = configs
        self.seed = seed
        self._device_ = device
        self.WandB = wb
        self._build()


    def __init__(self, exp_prefix, configs, seed, device, wb) -> None:
        super(MBPPO, self).__init__(exp_prefix, configs, seed, device)
        # print('init MBPPO Algorithm!')
        self.configs = configs
        self.seed = seed
        self._device_ = device
        self.WandB = wb
        self._build()


    ## build MBPPO components: (env, D, AC, alpha)
    def _build(self):
        super(MBPPO, self)._build()
        self._set_ppo()
        # self._set_fake_world()


    ## PPO
    def _set_ppo(self):
        PPO._build_ppo(self)


    ## FakeEnv
    def _set_fake_world(self):
        env_name = self.configs['environment']['name']
        device = self._device_
        if self.configs['environment']['name'][:4] == 'pddm':
        	static_fns = None
        else:
        	static_fns = mbpo_static[env_name[:-3].lower()]

        self.fake_world = FakeWorld(self.world_model)


    def learn(self):
        N = self.configs['algorithm']['learning']['epochs']
        NT = self.configs['algorithm']['learning']['epoch_steps']
        Ni = self.configs['algorithm']['learning']['init_epochs']
        Nx = self.configs['algorithm']['learning']['expl_epochs']

        E = self.configs['algorithm']['learning']['env_steps']
        G_ppo = self.configs['algorithm']['learning']['grad_PPO_steps']

        model_train_frequency = self.configs['world_model']['model_train_freq']
        batch_size_m = self.configs['world_model']['network']['batch_size'] # bs_m
        wm_epochs = self.configs['algorithm']['learning']['grad_WM_steps']
        real_ratio = self.configs['data']['real_ratio'] # rr
        batch_size = self.configs['data']['batch_size'] # bs
        mini_batch_size = self.configs['data']['mini_batch_size'] # bs
        rollout_trajectories = self.configs['data']['rollout_trajectories'] # bs_ro
        K = self.configs['data']['rollout_horizon']

        global_step = 0
        start_time = time.time()
        o, Z, el, t = self.learn_env.reset(), 0, 0, 0
        oldJs = [0, 0]
        JVList, JPiList = [0]*Ni, [0]*Ni
        logs = dict()
        lastEZ, lastES = 0, -2

        start_time_real = time.time()
        for n in range(1, N+1):
            if self.configs['experiment']['print_logs']:
                print('=' * 50)
                if n > Nx:
                    print(f'\n[ Epoch {n}   Learning ]'+(' '*50))
                    oldJs = [0, 0]
                    JVList, JPiList, KLList = [0], [0], [0]
                    JValList = [0]
                elif n > Ni:
                    print(f'\n[ Epoch {n}   Exploration + Learning ]'+(' '*50))
                    JVList, JPiList, KLList = [], [], []
                    JValList = []
                else:
                    print(f'\n[ Epoch {n}   Inintial Exploration ]'+(' '*50))
                    oldJs = [0, 0]
                    JVList, JPiList, KLList = [0], [0], [0]
                    JValList = [0]

            # print(f'[ Replay Buffer ] Size: {self.buffer.total_size()}')
            nt = 0
            # self.buffer.clean_buffer()
            o, d, Z, el, = self.learn_env.reset(), 0, 0, 0

            learn_start_real = time.time()
            while nt < NT: # full epoch
                # Interaction steps
                for e in range(1, E+1):
                    # o, Z, el, t = self.internact(n, o, Z, el, t)
                    o, Z, el, t = self.internact_opB(n, o, Z, el, t)
                    print(f'[ Epoch {n}   Interaction ] Env Steps: {e} | Return: {round(Z, 2)}'+(" "*10), end='\r')
                with T.no_grad(): v = self.actor_critic.get_v(T.Tensor(o)).cpu()
                # self.buffer.traj_tail(d, v, el)
                self.buffer.finish_path(el, v)

                # print(f'Buffer size = {self.buffer.total_size()}')
                # Taking gradient steps after exploration
                if n > Ni:
                    # if nt % model_train_frequency == 0:
                    # 03. Train model pθ on Denv via maximum likelihood
                    # print(f'\n[ Epoch {n}   Training World Model ]'+(' '*50))

                    ho_mean = 0. #self.fake_world.train_fake_world(self.buffer)
                    JValList.append(ho_mean) # ho: holdout

                    # model_fit_bs = self.configs['data']['buffer_size']
                    # model_fit_batch = self.buffer.sample_batch(model_fit_bs, self._device_)
                    # s, a, sp, _, _, _, _ = model_fit_batch.values()
                    # if n == Ni+1:
                    #     samples_to_collect = 4000
                    # else:
                    #     samples_to_collect = 1000

                    # for i, model in enumerate(self.models):
                    #     print(f'\n[ Epoch {n}   Training World Model {i+1} ]'+(' '*50))
                    #     loss_general = model.compute_loss(s[-samples_to_collect:],
                    #                                       a[-samples_to_collect:],
                    #                                       sp[-samples_to_collect:]) # generalization error
                    #     dynamics_loss = model.fit_dynamics(s, a, sp, **job_data)
                    #     logger.log_kv('dyn_loss_' + str(i), dynamics_loss[-1])
                    #     logger.log_kv('dyn_loss_gen_' + str(i), loss_general)
                    #     if job_data['learn_reward']:
                    #         reward_loss = model.fit_reward(s, a, r.reshape(-1, 1), **job_data)
                    #         logger.log_kv('rew_loss_' + str(i), reward_loss[-1])

                    KstepsList = []
                    self.initialize_model_buffer()
                    # PPO-P >>>>
                    for g in range(1, G_ppo+1):
                        print('\n')
                        # print(f'Actor-Critic Grads...{g}', end='\r')
                        # print(f'[ Epoch {n}   Training Actor-Critic ] AC Grads: {g}'+(" "*50), end='\r')
                        # # Reallocate model buffer
                        # self.initialize_model_buffer()
                        self.model_buffer.reset()
                        # # Generate M k-steps imaginary rollouts for SAC traingin
                        # self.rollout_world_model_op(rollout_trajectories, K, n)
                        # k_avg = self.rollout_world_model_trajectories(rollout_trajectories, K, n)
                        k_avg = self.rollout_real_world_trajectories(rollout_trajectories, K, g, n)
                        KstepsList.append(k_avg)
                        batch_size = int(self.model_buffer.total_size())
                        print('\n')
                        stop_pi = False
                        kl = 0
                        for gg in range(1, 81): # 101
                            print(f'[ Epoch {n} ] AC: {g} | ac: {gg} || stopPG={stop_pi} | KL={round(kl, 4)}'+(' '*40), end='\r')
                            batch = self.model_buffer.sample_batch(batch_size, self._device_)
                            Jv, Jpi, kl, stop_pi = self.trainAC(g, batch, oldJs)
                            # print(f'[ Epoch {n} ] AC: {g} | ac: {gg} || PG={stop_pi} | KL={kl}'+(' '*80), end='\r')
                            oldJs = [Jv, Jpi]
                            JVList.append(Jv.item())
                            JPiList.append(Jpi.item())
                            KLList.append(kl)
                    # PPO-P <<<<

                nt += E

            print('\n')
            # logs['time/training                  '] = time.time() - learn_start_real

            # logs['training/wm/Jtrain_mean        '] = np.mean(JMeanTrainList)
            # logs['training/wm/Jtrain             '] = np.mean(JTrainList)
            logs['training/wm/Jval               '] = np.mean(JValList)
            # logs['training/wm/test_mse           '] = np.mean(LossTestList)

            logs['training/ppo/Jv                '] = np.mean(JVList)
            logs['training/ppo/Jpi               '] = np.mean(JPiList)
            logs['training/ppo/KL                '] = np.mean(KLList)

            logs['data/env_buffer                '] = self.buffer.total_size()
            logs['data/env_rollout_steps         '] = self.buffer.average_horizon()
            if hasattr(self, 'model_buffer'):
                logs['data/model_buffer              '] = self.model_buffer.total_size()
                logs['data/model_rollout_steps       '] = np.mean(KstepsList)
            else:
                logs['data/model_buffer              '] = 0.
                logs['data/model_rollout_steps       '] = 0.
            # else:
            #     logs['data/model_buffer              '] = 0
            # logs['data/rollout_length            '] = K

            eval_start_real = time.time()
            EZ, ES, EL = self.evaluate()
            # EZ, ES, EL = self.evaluate_op()

            # logs['time/evaluation                '] = time.time() - eval_start_real

            if self.configs['environment']['type'] == 'mujoco-pddm-shadowhand':
                logs['evaluation/episodic_score_mean '] = np.mean(ES)
                logs['evaluation/episodic_score_std  '] = np.std(ES)
            else:
                logs['evaluation/episodic_return_mean'] = np.mean(EZ)
                logs['evaluation/episodic_return_std '] = np.std(EZ)
            logs['evaluation/episodic_length_mean'] = np.mean(EL)

            logs['time/total                     '] = time.time() - start_time_real

            # if n > (N - 50):
            #     if self.configs['environment']['type'] == 'mujoco-pddm-shadowhand':
            #         if np.mean(ES) > lastES:
            #             print(f'[ Epoch {n}   Agent Saving ]                    ')
            #             env_name = self.configs['environment']['name']
            #             alg_name = self.configs['algorithm']['name']
            #             T.save(self.actor_critic.actor,
            #             f'./saved_agents/{env_name}-{alg_name}-seed:{self.seed}-epoch:{n}.pTtar')
            #             lastES = np.mean(ES)
            #     else:
            #         if np.mean(EZ) > lastEZ:
            #             print(f'[ Epoch {n}   Agent Saving ]                    ')
            #             env_name = self.configs['environment']['name']
            #             alg_name = self.configs['algorithm']['name']
            #             T.save(self.actor_critic.actor,
            #             f'./saved_agents/{env_name}-{alg_name}-seed:{self.seed}-epoch:{n}.pTtar')
            #             lastEZ = np.mean(EZ)

            # Printing logs
            if self.configs['experiment']['print_logs']:
                for k, v in logs.items():
                    print(f'{k}  {round(v, 4)}'+(' '*10))

            # WandB
            if self.WandB:
                wandb.log(logs)

        self.learn_env.close()
        self.eval_env.close()


    def rollout_real_world_trajectories(self, rollout_trajectories, K, g, n):
    	# 07. Sample st uniformly from Denv
    	device = self._device_
    	# Nτ = 400 * 2 # number of trajectories x number of models
    	Nτ = 40 # number of trajectories x number of models
    	# Nτ = self.configs['data']['rollout_trajectories']
    	# o, d, el = self.traj_env.reset(), 0, 0

        # 08. Perform k-step model rollout starting from st using policy πφ; add to Dmodel
    	k_end_total = 0
    	for nτ in range(1, Nτ+1): # Generate trajectories
            # o, d, el = self.traj_env.reset(), 0, 0
            o, Z, el = self.traj_env.reset(), 0, 0
            # print(f'nτ = {nτ}'+(' '*70), end='\r')
            for k in range(1, K+1): # Generate rollouts
                print(f'[ Epoch {n} ] Model Rollout: nτ = {nτ} | k = {k} | Buffer = {self.model_buffer.total_size()} | Return = {round(Z, 2)}', end='\r')
                # print(f'[ Epoch {n} ] AC Training Grads: {g} || Model Rollout: nτ = {nτ} | k = {k} | Buffer size = {self.model_buffer.total_size()}'+(' '*10))
                with T.no_grad(): a, log_pi, v = self.actor_critic.get_a_and_v_np(T.Tensor(o))
                # o_next, r, d_next, _ = self.traj_env.step(a)
                o_next, r, d, _ = self.traj_env.step(a)
                Z += r
                el += 1

                # self.model_buffer.store_transition(o, a, r, d, v, log_pi)
                self.model_buffer.store(o, a, r, o_next, v, log_pi, el)
                o = o_next
                # if d_next or (el == K):
                #     o_next, d_next, el = self.traj_env.reset(), 0, 0
                #     print(f'[ Model Rollout ] Breaking early: {k}'+(' '*50), end='\r')
                #     # k_end_total += k
                #     break
                if d or (el == K):
                    # print(f'[ Model Rollout ] Breaking early: {k}'+(' '*50), end='\r')
                    # if el == K:
                    #     with T.no_grad(): v = self.actor_critic.get_v(T.Tensor(o)).cpu()
                    # else:
                    #     v = T.Tensor([0.0])
                    # self.model_buffer.finish_path(el, v)
                    break
                    # o, el = self.traj_env.reset(), 0
                # else:
                # o, d = o_next, d_next
            if el == K:
                # print(f'[ Model Rollout ] Average Breaking : {k_end_total//Nτ}'+(' '*50))
                with T.no_grad(): v = self.actor_critic.get_v(T.Tensor(o)).cpu()
            else:
                v = T.Tensor([0.0])
            # self.model_buffer.traj_tail(d, v)
            self.model_buffer.finish_path(el, v)
            k_end_total += k
    	return k_end_total//Nτ


    def rollout_world_model_trajectories(self, rollout_trajectories, K, n):
    	# 07. Sample st uniformly from Denv
    	device = self._device_
    	batch_size = rollout_trajectories
    	B_ro = self.buffer.sample_batch(batch_size) # Torch
    	O = B_ro['observations_next']
    	D = B_ro['terminals']

        # 08. Perform k-step model rollout starting from st using policy πφ; add to Dmodel
    	k_end_total = 0
    	for o, d in zip(O, D):
            # print('o: ', o.shape)
            # print('d: ', d.shape)
            for k in range(1, K+1):
                # print(f'k = {k}', end='\r')
                with T.no_grad(): a, log_pi, _, v = self.actor_critic.get_pi_and_v(o)
                a, log_pi, v = a.cpu(), log_pi.cpu(), v.cpu()

                o_next, r, d_next, _ = self.fake_world.step(o, a, deterministic=True) # ip: Tensor, op: Tensor

                self.model_buffer.store_transition(o, a, r, d, v, log_pi)

                if d_next:
    	            print(f'[ Model Rollout ] Breaking early: {k}'+(' '*50), end='\r')
    	            k_end_total += k
    	            break

                o, d = o_next, d_next

    	print(f'[ Model Rollout ] Average Breaking : {k_end_total//batch_size}'+(' '*50))
    	with T.no_grad(): v = self.actor_critic.get_v(o)
    	self.model_buffer.traj_tail(d, v)

    	return k_end_total//batch_size


    def rollout_world_model_op(self, rollout_trajectories, K, n):
    	# 07. Sample st uniformly from Denv
    	device = self._device_
    	batch_size = rollout_trajectories
    	# print(f'[ Epoch {n}   Model Rollout ] Batch Size: {batch_size} | Rollout Horizon: {K}'+(' '*50))
    	B_ro = self.buffer.sample_batch(batch_size) # Torch
    	O = B_ro['observations']
    	D = B_ro['terminals']

        # 08. Perform k-step model rollout starting from st using policy πφ; add to Dmodel
    	for k in range(1, K+1):
            # print('k = ', k)
            # print('k = ', k, end='\r')
            A = self.actor_critic.get_action(O) # Stochastic action | No reparameterization
            with T.no_grad(): A, log_Pi, _, V = self.actor_critic.get_pi_and_v(O)

            O_next, R, D_next, _ = self.fake_world.step(O, A) # ip: Tensor, op: Tensor
        	# O_next, R, D, _ = self.fake_world.step_np(O, A) # ip: Tensor, op: Numpy

            self.model_buffer.store_batch(O, A, R, D, V, log_Pi)

            O_next = T.Tensor(O_next)
            D = T.tensor(D, dtype=T.bool)
            nonD = ~D.squeeze(-1)
            if nonD.sum() == 0:
        	    print(f'[ Epoch {n}   Model Rollout ] Breaking early: {k} | {nonD.sum()} / {nonD.shape}')
        	    break

            O, D = O_next[nonD], D_next[nonD]

    	with T.no_grad(): V = self.actor_critic.get_v(O)
    	self.model_buffer.traj_tail(D, V)


    def ppo_mini_batch(self, mini_batch_size):
    	B = self.model_buffer.sample_batch(mini_batch_size, self._device_)
    	return B





def main(exp_prefix, config, seed, device, wb):

    print('Start an MBPPO experiment...')
    print('\n')

    configs = config.configurations

    if seed:
        random.seed(seed), np.random.seed(seed), T.manual_seed(seed)

    alg_name = configs['algorithm']['name']
    alg_mode = configs['algorithm']['mode']
    env_name = configs['environment']['name']
    env_type = configs['environment']['type']
    wm_epochs = configs['algorithm']['learning']['grad_WM_steps']
    DE = configs['world_model']['num_ensembles']

    group_name = f"{env_name}-{alg_name}-{alg_mode}-D"
    exp_prefix = f"seed:{seed}"

    if wb:
        wandb.init(
            name=exp_prefix,
            group=group_name,
            # project='test',
            project='AMMI-RL-2022',
            config=configs
        )

    agent = MBPPO(exp_prefix, configs, seed, device, wb)

    agent.learn()

    print('\n')
    print('... End the MBPPO experiment')


if __name__ == "__main__":

    import argparse

    parser = argparse.ArgumentParser()

    parser.add_argument('-exp_prefix', type=str)
    parser.add_argument('-cfg', type=str)
    parser.add_argument('-seed', type=str)
    parser.add_argument('-device', type=str)
    parser.add_argument('-wb', type=str)

    args = parser.parse_args()

    exp_prefix = args.exp_prefix
    sys.path.append(f"{os.getcwd()}/configs")
    config = importlib.import_module(args.cfg)
    seed = int(args.seed)
    device = args.device
    wb = eval(args.wb)

    main(exp_prefix, config, seed, device, wb)

# T.multiprocessing.set_sharing_strategy('file_system') # from rl.data.dataset import RLDataModule Greek alphabet:         Α α, Β β, Γ γ, Δ δ, Ε ε, Ζ ζ, Η η, Θ θ, Ι ι, Κ κ, Λ λ, Μ μ,         Ν ν, Ξ ξ, Ο ο, Π π, Ρ ρ, Σ σ/ς, Τ τ, Υ υ, Φ φ, Χ χ, Ψ ψ, Ω ω.      Algorithm: Model-Based Game: Policy As Leader (PAL) – Practical Version          01: Initialize: Models parameters( Policy net πφ0, Value net Vψ0, ensemble of MDP world models {M^θ0}_{1:nM} )         02. Initialize: Replay buffer D.         03: Hyperparameters: Initial samples Ni = 2 epochs, samples per update E=1000, buffer size B = 2500 ≈ N, number of NPG steps K = 4 ≈ 1         04: Initial Data: Collect N_init samples from the environment by interacting with initial policy. Store data in buffer D.         05: for n = 0, 1, 2, ..., N (nEpochs) do         06:    Learn dynamics model(s) M^_n+1 using data in the buffer.         07:    Policy updates: π_n+1; V_n+1 = MB-PPO(π_n, V_n, M^_n+1) // call K times         08:    Collect dataset of N samples from World by interacting with πn+1.         09.    Add data to replay buffer D, discarding old data if size is larger than B.         10: end for       Algorithm: Model-Based Game: Model As Leader (MAL) – Practical Version         1: Initialize: Policy network π_0, model network(s) Mhat_0, value network V_0.         2: Hyperparameters: Initial samples N_init, samples per update N, number of PPO steps K >> 1         3: Initial Data: Collect N_init samples from the environment by interacting with initial policy. Store data in buffer D.         4: Initial Model: Learn model(s) Mhat_0 using data in D.         5: for k = 0, 1, 2, ... do         6:    Optimize π_k+1 using Mck by running K >> 1 steps of model-based PPO (Subroutine 1).         7:    Collect dataset D_k+1 of N samples from world using πk+1. Aggregate data D = D U D_k+1.         8:    Learn dynamics model(s) Mhat_k+1 using data in D.         9: end for       Algorithm: Model-Based Proximal Policy Optimization (On-Policy, Dyna-style, Model-Based)          01. Inputs: Models parameters( Policy net πφn, Value net Vψn, ensemble of MDP world models {M^θn+1}_{1:nM}, Replay buffer D)         02. Hyperparameters: Disc. factor γ, GAE λ, num traj's Nτ, model rollout horizon H         03. Initialize: Trajectory buffer Dτ = {}         04. for k = 0, 1, 2, ..., Nτ:         05.     Sample init p(s^k_0) from distribution/buffer         05.     Rollout a traj {τ^k_π/M} from {M^θ}_{1:nM} by πn = π(φn) for e = 0, 1, 2, ..., H         06.     Aggregate traj τ^k in the traj buffer, Dτ = Dτ U {τ^k_π/M}         07.     Compute RTG R^_t, GAE A^_t based on Vn = V(θn)         08. end for         09. Update πφ by maxz Jπ                 φ = arg max_φ {(1/T|Dk|) sum sum min((π/πk), 1 +- eps) Aπk }         10. Fit Vθ by MSE(Jv)                 θ = arg min_θ {(1/T|Dk|) sum sum (Vθ(st) - RTG)^2 }         11. Return: Policy net πθn+1, value net Vφn+1      Subroutine 1: Model-Based Natural Policy Gradient Update Step          1: Require: Policy (stochastic) network πθ, value/baseline network V , ensemble of MDP dynamics models fMcφg,         reward function R, initial state distribution or buffer.         2: Hyperparameters: Discount factor γ, GAE λ, number of trajectories Nτ, rollout horizon H, normalized NPG step         size δ         3: Initialize trajectory buffer Dτ = fg         4: for k = 1; 2; : : : ; Nτ do         5: Sample initial state sk 0 from initial state distribution/buffer         6: Perform H step rollout from sk 0 with πθ to get τjk = (sk 0; ak 0; sk 1; ak 2; : : : sk H; ak H), one for each model Mcφj in the         ensemble.         7: Query reward function to obtain rewards for each step of the trajectories         8: Truncate trajectories if termination/truncation conditions are part of the environment         9: Aggregate the trajectories in trajectory buffer, Dτ = Dτ [ fτg         10: end for         11: Compute advantages for each trajectory using V and GAE (Schulman et al., 2016).         12: Compute vanilla policy gradient using the dataset         g = E(s;a)∼Dτ [rθ log πθ(ajs)Aπ(s; a)]         13: Perform normalized NPG update (F denotes the Fisher matrix)         θ = θ + sgT Fδ−1g F −1g         14: Update value/baseline network V to fit the computed returns in Dτ.         15: Return Policy network πθ, value network V # print('init MBPPO Algorithm!') # print('init MBPPO Algorithm!') ## build MBPPO components: (env, D, AC, alpha) # self._set_fake_world() ## PPO ## FakeEnv # bs_m # rr # bs # bs # bs_ro # print(f'[ Replay Buffer ] Size: {self.buffer.total_size()}') # self.buffer.clean_buffer() # full epoch # Interaction steps # o, Z, el, t = self.internact(n, o, Z, el, t) # self.buffer.traj_tail(d, v, el) # print(f'Buffer size = {self.buffer.total_size()}') # Taking gradient steps after exploration # if nt % model_train_frequency == 0: # 03. Train model pθ on Denv via maximum likelihood # print(f'\n[ Epoch {n}   Training World Model ]'+(' '*50)) #self.fake_world.train_fake_world(self.buffer) # ho: holdout # model_fit_bs = self.configs['data']['buffer_size'] # model_fit_batch = self.buffer.sample_batch(model_fit_bs, self._device_) # s, a, sp, _, _, _, _ = model_fit_batch.values() # if n == Ni+1: #     samples_to_collect = 4000 # else: #     samples_to_collect = 1000 # for i, model in enumerate(self.models): #     print(f'\n[ Epoch {n}   Training World Model {i+1} ]'+(' '*50)) #     loss_general = model.compute_loss(s[-samples_to_collect:], #                                       a[-samples_to_collect:], #                                       sp[-samples_to_collect:]) # generalization error #     dynamics_loss = model.fit_dynamics(s, a, sp, **job_data) #     logger.log_kv('dyn_loss_' + str(i), dynamics_loss[-1]) #     logger.log_kv('dyn_loss_gen_' + str(i), loss_general) #     if job_data['learn_reward']: #         reward_loss = model.fit_reward(s, a, r.reshape(-1, 1), **job_data) #         logger.log_kv('rew_loss_' + str(i), reward_loss[-1]) # PPO-P >>>> # print(f'Actor-Critic Grads...{g}', end='\r') # print(f'[ Epoch {n}   Training Actor-Critic ] AC Grads: {g}'+(" "*50), end='\r') # # Reallocate model buffer # self.initialize_model_buffer() # # Generate M k-steps imaginary rollouts for SAC traingin # self.rollout_world_model_op(rollout_trajectories, K, n) # k_avg = self.rollout_world_model_trajectories(rollout_trajectories, K, n) # 101 # print(f'[ Epoch {n} ] AC: {g} | ac: {gg} || PG={stop_pi} | KL={kl}'+(' '*80), end='\r') # PPO-P <<<< # logs['time/training                  '] = time.time() - learn_start_real # logs['training/wm/Jtrain_mean        '] = np.mean(JMeanTrainList) # logs['training/wm/Jtrain             '] = np.mean(JTrainList) # logs['training/wm/test_mse           '] = np.mean(LossTestList) # else: #     logs['data/model_buffer              '] = 0 # logs['data/rollout_length            '] = K # EZ, ES, EL = self.evaluate_op() # logs['time/evaluation                '] = time.time() - eval_start_real # if n > (N - 50): #     if self.configs['environment']['type'] == 'mujoco-pddm-shadowhand': #         if np.mean(ES) > lastES: #             print(f'[ Epoch {n}   Agent Saving ]                    ') #             env_name = self.configs['environment']['name'] #             alg_name = self.configs['algorithm']['name'] #             T.save(self.actor_critic.actor, #             f'./saved_agents/{env_name}-{alg_name}-seed:{self.seed}-epoch:{n}.pTtar') #             lastES = np.mean(ES) #     else: #         if np.mean(EZ) > lastEZ: #             print(f'[ Epoch {n}   Agent Saving ]                    ') #             env_name = self.configs['environment']['name'] #             alg_name = self.configs['algorithm']['name'] #             T.save(self.actor_critic.actor, #             f'./saved_agents/{env_name}-{alg_name}-seed:{self.seed}-epoch:{n}.pTtar') #             lastEZ = np.mean(EZ) # Printing logs # WandB # 07. Sample st uniformly from Denv # Nτ = 400 * 2 # number of trajectories x number of models # number of trajectories x number of models # Nτ = self.configs['data']['rollout_trajectories'] # o, d, el = self.traj_env.reset(), 0, 0 # 08. Perform k-step model rollout starting from st using policy πφ; add to Dmodel # Generate trajectories # o, d, el = self.traj_env.reset(), 0, 0 # print(f'nτ = {nτ}'+(' '*70), end='\r') # Generate rollouts # print(f'[ Epoch {n} ] AC Training Grads: {g} || Model Rollout: nτ = {nτ} | k = {k} | Buffer size = {self.model_buffer.total_size()}'+(' '*10)) # o_next, r, d_next, _ = self.traj_env.step(a) # self.model_buffer.store_transition(o, a, r, d, v, log_pi) # if d_next or (el == K): #     o_next, d_next, el = self.traj_env.reset(), 0, 0 #     print(f'[ Model Rollout ] Breaking early: {k}'+(' '*50), end='\r') #     # k_end_total += k #     break # print(f'[ Model Rollout ] Breaking early: {k}'+(' '*50), end='\r') # if el == K: #     with T.no_grad(): v = self.actor_critic.get_v(T.Tensor(o)).cpu() # else: #     v = T.Tensor([0.0]) # self.model_buffer.finish_path(el, v) # o, el = self.traj_env.reset(), 0 # else: # o, d = o_next, d_next # print(f'[ Model Rollout ] Average Breaking : {k_end_total//Nτ}'+(' '*50)) # self.model_buffer.traj_tail(d, v) # 07. Sample st uniformly from Denv # Torch # 08. Perform k-step model rollout starting from st using policy πφ; add to Dmodel # print('o: ', o.shape) # print('d: ', d.shape) # print(f'k = {k}', end='\r') # ip: Tensor, op: Tensor # 07. Sample st uniformly from Denv # print(f'[ Epoch {n}   Model Rollout ] Batch Size: {batch_size} | Rollout Horizon: {K}'+(' '*50)) # Torch # 08. Perform k-step model rollout starting from st using policy πφ; add to Dmodel # print('k = ', k) # print('k = ', k, end='\r') # Stochastic action | No reparameterization # ip: Tensor, op: Tensor # O_next, R, D, _ = self.fake_world.step_np(O, A) # ip: Tensor, op: Numpy # project='test',