rlagent.py

from rlcore.algo import PPO
from rlcore.storage import RolloutStorage


class Neo(object):

  def __init__(self, args, policy, obs_shape, action_space):
    super().__init__()

    self.obs_shape = obs_shape
    self.action_space = action_space
    self.actor_critic = policy
    self.rollouts = RolloutStorage(args.num_steps, args.num_processes, self.obs_shape, self.action_space, 
                                   recurrent_hidden_state_size=1)
    self.args = args
    self.trainer = PPO(self.actor_critic, args.clip_param, args.ppo_epoch, args.num_mini_batch, args.value_loss_coef,
                       args.entropy_coef, lr=args.lr,max_grad_norm=args.max_grad_norm)

  def load_model(self, policy_state):
      self.actor_critic.load_state_dict(policy_state)

  def initialize_obs(self, obs):
    # this function is called at the start of episode
    self.rollouts.obs[0].copy_(obs)

  def update_rollout(self, obs, reward, mask):
    self.rollouts.insert(obs, self.states, self.action, self.action_log_prob, self.value, reward, mask)

  def act(self, step, deterministic=False):
    self.value, self.action, self.action_log_prob, self.states = self.actor_critic.act(self.rollouts.obs[step],
              self.rollouts.recurrent_hidden_states[step],self.rollouts.masks[step],deterministic=deterministic)
    return self.action

  def wrap_horizon(self, next_value):
    self.rollouts.compute_returns(next_value, True, self.args.gamma, self.args.tau)

  def after_update(self):
    self.rollouts.after_update()

  def update(self):
    return self.trainer.update(self.rollouts)