Refactored rl.ppo.Memory to use deques

from typing import Any, Union, Dict

from typing import Any, OrderedDict, Union, Dict

from copy import deepcopy

from torch.nn import Module

def clone(model: Union[BaseEstimator, Module], attrs: Dict[str, Any]=None) -> Union[BaseEstimator, Module]:

def clone(model: Union[BaseEstimator, Module, OrderedDict, Dict],

          attrs: Dict[str, Any]=None) -> Union[BaseEstimator, Module, OrderedDict, Dict]:

    """

    Copy a scikit-learn or pytorch model.

"""

Proximal Policy Optimization (PPO)

"""

from typing import Callable, List

from typing import Callable, List, Union, Tuple

from collections import deque

import torch

import torch.nn as nn

class Memory:

    def __init__(self):

        self.actions = []

        self.states = []

        self.logprobs = []

        self.rewards = []

        self.is_terminals = []

    def __init__(self, states=(), actions=(), rewards=(), logprobs=(), is_terminals=(),

                 maxlen=None):

        self.maxlen = maxlen

        self.states = deque(states, maxlen=maxlen)

        self.actions = deque(actions, maxlen=maxlen)

        self.rewards = deque(rewards, maxlen=maxlen)

        self.logprobs = deque(logprobs, maxlen=maxlen)

        self.is_terminals = deque(is_terminals, maxlen=maxlen)

        self._lists = (self.states, self.actions, self.rewards, self.logprobs,

                       self.is_terminals)

    def __len__(self):

        return len(self.states)

    def add(self, state, action, logprob, reward, done):

    def clear(self):

        del self.actions[:]

        del self.states[:]

        del self.logprobs[:]

        del self.rewards[:]

        del self.is_terminals[:]

        """

        Clear all entries.

        """

        for l in self._lists:

                    for _ in range(len(l)):

                        del l[0]

    def flush(self):

        """

        Delete all entries except those belonging to the last, unfinished episode.

        """

        size = len(self.states)

        for i, done in enumerate(reversed(self.is_terminals)):

            if done:

                truncate = size - i

                del self.actions[:truncate]

                del self.states[:truncate]

                del self.logprobs[:truncate]

                del self.rewards[:truncate]

                del self.is_terminals[:truncate]

                for l in self._lists:

                    for idx in range(truncate):

                        del l[0]

                break

    def as_array(self):

        return \

            np.asarray(self.states), \

            np.asarray(self.actions), \

            np.asarray(self.rewards), \

            np.asarray(self.logprobs), \

            np.asarray(self.is_terminals)

class Policy(nn.Module):

    def update(self, policy, memory, epochs: int=1, optimizer=None, summary=None,

    def update(self, policy, memory: Memory, epochs: int=1, optimizer=None, summary=None,

        grad_callback=None):

        rewards = returns(memory.rewards, memory.is_terminals, self.gamma, truncate=self.truncate)

        states, actions, rewards, logprobs, is_terminals = memory.as_array()

        rewards = returns(rewards, is_terminals, self.gamma, truncate=self.truncate)

        truncate = len(rewards)

        # If the returns calculated are zero length, i.e. when memory does not

        # contain a single full episode, because returns() truncated incompleted

        # Casting to correct data type and DEVICE

        # pylint: disable=not-callable

        rewards = torch.tensor(rewards[:truncate]).float().to(self.device)

        old_states = torch.tensor(memory.states[:truncate]).float().to(self.device).detach()

        old_actions = torch.tensor(memory.actions[:truncate]).float().to(self.device).detach()

        old_logprobs = torch.tensor(memory.logprobs[:truncate]).float().to(self.device).detach()

        old_states = torch.tensor(states[:truncate]).float().to(self.device).detach()

        old_actions = torch.tensor(actions[:truncate]).float().to(self.device).detach()

        old_logprobs = torch.tensor(logprobs[:truncate]).float().to(self.device).detach()

        # If states/actions are 1D arrays of single number states/actions,

        # convert them to 2D matrix of 1 column where each row is one timestep.

            for t in trange(1, int(timesteps) + 1, leave=False):

                # Running policy:

                memory.states.append(state)

                action, logprob = policy.predict(state)

                state, reward, done, info = self.env.step(action)

                new_state, reward, done, info = self.env.step(action)

                episodic_rewards[-1] += reward

                interval_rewards[-1] += reward

                t_episode += 1

                if done:

                    state = self.env.reset()

                    new_state = self.env.reset()

                    if reward_aggregation.endswith('normalized'):

                        episodic_rewards[-1] /= t_episode

                    episodic_rewards.append(0.)

                    t_episode = 0

                memory.actions.append(action)

                memory.logprobs.append(logprob)

                memory.rewards.append(reward)

                memory.is_terminals.append(done)

                memory.add(state, action, logprob, reward, done)

                if step_callback is not None:

                    step_callback(locals())

                        lr_scheduler()

                    memory.flush()

                state = new_state

            self.meta_policy = policy if track_higher_grads else None

        self.policy.load_state_dict(policy.state_dict())

            return interval_rewards

    def predict(self, state):

    def predict(self, state) -> Tuple[Union[np.ndarray, int], float]:

        return self.policy.predict(state)