<jupyter_start><jupyter_text>Implementing Advantage-Actor Critic (A2C) In this notebook you will implement Advantage Actor Critic algorithm that trains on a batch of Atari 2600 environments running in parallel. Firstly, we will use environment wrappers implemented in file `atari_wrappers.py`. These wrappers preprocess observations (resize, grayscal, take max between frames, skip frames and stack them together) and rewards. Some of the wrappers help to reset the environment and pass `done` flag equal to `True` when agent dies.File `env_batch.py` includes implementation of `ParallelEnvBatch` class that allows to run multiple environments in parallel. To create an environment we can use `nature_dqn_env` function. Note that if you are using PyTorch and not using `tensorboardX` you will need to implement a wrapper that will log **raw** total rewards that the *unwrapped* environment returns and redefine the implemention of `nature_dqn_env` function here.<jupyter_code>import numpy as np
from atari_wrappers import nature_dqn_env, NumpySummaries

nenvs = 8

env = nature_dqn_env("SpaceInvadersNoFrameskip-v4", nenvs=nenvs, 
                     summaries='Numpy')
obs = env.reset()
n_actions = env.action_space.n
print('num of actions: {}'.format(n_actions))
assert obs.shape == (8, 84, 84, 4)
assert obs.dtype == np.uint8

# for evaluating the agent


single_env = nature_dqn_env("SpaceInvadersNoFrameskip-v4", nenvs=1, 
                     summaries='Numpy')

print(single_env.reset())
import gym
one_env = gym.make('SpaceInvadersNoFrameskip-v4')
one_obs = one_env.reset()
print(one_obs.shape)<jupyter_output>(210, 160, 3)<jupyter_text>Next, we will need to implement a model that predicts logits and values. It is suggested that you use the same model as in [Nature DQN paper](https://web.stanford.edu/class/psych209/Readings/MnihEtAlHassibis15NatureControlDeepRL.pdf) with a modification that instead of having a single output layer, it will have two output layers taking as input the output of the last hidden layer. **Note** that this model is different from the model you used in homework where you implemented DQN. You can use your favorite deep learning framework here. We suggest that you use orthogonal initialization with parameter $\sqrt{2}$ for kernels and initialize biases with zeros.<jupyter_code># import tensorflow as torch
# import torch as tf

import torch.nn as nn
import torch.nn.functional as F


def conv2d_size_out(size, kernel_size, stride):
    """
    common use case:
    cur_layer_img_w = conv2d_size_out(cur_layer_img_w, kernel_size, stride)
    cur_layer_img_h = conv2d_size_out(cur_layer_img_h, kernel_size, stride)
    to understand the shape for dense layer's input
    """
    return (size - (kernel_size - 1) - 1) // stride  + 1

class DQN(nn.Module):
    def __init__(self, obs_dim, n_actions):
        """
        obs_dim: image dimension of the observation 
            The input image must have 2d dim 
            obs_dim x obs_dim
        """
        super().__init__()
        

        kernel_size = 3
        stride = 2
        self.conv1 = nn.Conv2d(4, 16, kernel_size, stride)
        out_size = conv2d_size_out(obs_dim, kernel_size, stride)
        self.conv2 = nn.Conv2d(16, 32, kernel_size, stride)
        out_size = conv2d_size_out(out_size, kernel_size, stride)
        self.conv3 = nn.Conv2d(32, 64, kernel_size, stride)
        out_size = conv2d_size_out(out_size, kernel_size, stride)

        # size of the output tensor after convolution batch_size x 64 x out_size x out_size
        self.linear = nn.Linear(64*out_size*out_size, 256)
        
        
        # value head
        self.value = nn.Linear(256, 1)
        
        # policy head
        self.policy = nn.Linear(256, n_actions)
        

    def forward(self, state_t):
        """
        takes agent's observation (tensor), returns qvalues (tensor)
        :param state_t: a batch of 4-frame buffers, shape = [batch_size, 4, h, w]
        """
        # Use your network to compute qvalues for given state
        # qvalues = <YOUR CODE>
        
        #print('== type of input tensor ==', type(state_t))
        #print(state_t.shape)
        t = self.conv1(state_t)
        t = F.relu(t)
        t = self.conv2(t)
        t = F.relu(t)
        t = self.conv3(t)
        t = F.relu(t)

        t = t.view(state_t.shape[0], -1)
        t = self.linear(t)
        t = F.relu(t)
        
        values = self.value(t)
        
        # log proba of actions 
        logits = F.log_softmax(self.policy(t), dim=1)
        return values, logits<jupyter_output><empty_output><jupyter_text>You will also need to define and use a policy that wraps the model. While the model computes logits for all actions, the policy will sample actions and also compute their log probabilities.  `policy.act` should return a dictionary of all the arrays that are needed to interact with an environment and train the model. Note that actions must be an `np.ndarray` while the othertensors need to have the type determined by your deep learning framework.<jupyter_code>import torch
from torch.distributions import Categorical

device = 'cuda' if torch.cuda.is_available() else 'cpu'

class Policy:
    def __init__(self, model):
        self.model = model
        self.model.to(device)
        
    def parse_state(self, inputs):
        #inputs = torch.Tensor(inputs, device=device)
        #inputs = inputs.permute(0, 3, 1, 2).contiguous()
        
        # to float
        inputs = inputs.astype(np.float32)
        
        # depth major
        inputs = inputs.transpose(0, 3, 1, 2)
        
        # normalize the inputs
        mean = np.mean(inputs, axis=(2, 3), keepdims=True)
        std = np.std(inputs, axis=(2, 3), keepdims=True)
        inputs = (inputs- mean) / std
        
        # to torch tensor
        inputs = torch.Tensor(inputs, device=device)
        return inputs
    
    def act(self, inputs):
        # input dim (8, 84, 84, 4)
        # 8: number of parallel env 
        # 4: last 4 frames to make POMDP a MDP
        # 84: frame dim
        
        # Should return a dict containing keys ['actions', 'logits', 'log_probs', 'values'].
        # inputs are batched input from different environment
        # each action in actions repr the action taken in one env
        # Same thing with log probability
        
        # convert input to depth major      
        inputs = self.parse_state(inputs)

        # value estimates of states and 
        # log proba of actions at the states
        values, logits = self.model(inputs)
        # sample from log probabilities
        m = Categorical(logits=logits)
    
        # sample an action for each env
        actions = m.sample()
        
        return {
            'actions': actions,
            'log_probs': m.log_prob(actions),
            'estimated_values': values
        }
    
    
    def choose_action(self, inputs):
        """Choose action greedily in a testing env"""
        inputs = self.parse_state(inputs)
        
        with torch.no_grad():
            _, logits = self.model(inputs)
            actions = torch.argmax(logits, dim=-1)
        return actions
        
    def get_value(self, inputs):
        """estimate state value without keeping the computation graph"""
        inputs = self.parse_state(inputs)
        with torch.no_grad():
            v, _ = self.model(inputs)
        return v<jupyter_output><empty_output><jupyter_text>Next will pass the environment and policy to a runner that collects partial trajectories from the environment. The class that does is is already implemented for you. This runner interacts with the environment for a given number of steps and returns a dictionary containingkeys * 'observations' * 'rewards' * 'resets'* 'actions'* all other keys that you defined in `Policy`under each of these keys there is a python `list` of interactions with the environment of specified length $T$ &mdash; the size of partial trajectory.  To train the part of the model that predicts state values you will need to compute the value targets. Any callable could be passed to `EnvRunner` to be applied to each partial trajectory after it is collected. Thus, we can implement and use `ComputeValueTargets` callable. The formula for the value targets is simple:$$\hat v(s_t) = \left( \sum_{t'=0}^{T - 1} \gamma^{t'}r_{t+t'} \right) + \gamma^T \hat{v}(s_{t+T}),$$In implementation, however, do not forget to use `trajectory['resets']` flags to check if you need to add the value targets at the next step when computing value targets for the current step. You can access `trajectory['state']['latest_observation']`to get last observations in partial trajectory &mdash; $s_{t+T}$.<jupyter_code>class ComputeTargetValues:
    """Explicitly compute target value along the trajectory
    in the forward view
    """
    def __init__(self, policy, gamma=0.99):
        self.policy = policy
        self.gamma = gamma
    
    def __call__(self, trajectory):
        # agent takes `nsteps` actions
        nsteps = len(trajectory['actions'])
        
        # target values forward view
        target_values = []
        r = self.policy.get_value(
            trajectory['state']['latest_observation']).squeeze(dim=1)
        
        for i in range(nsteps-1, -1, -1):
            immediate_reward = trajectory['rewards'][i]
            #print(immediate_reward)
            immediate_reward = torch.Tensor(immediate_reward, 
                                            device=device)
            r = immediate_reward + self.gamma * r * trajectory['resets'][i]
            target_values.append(r)
        
        # reverse discounted rewards
        trajectory['target_values'] = [
            r for r in reversed(target_values)]
        return
# get one trajectory
import runners

model = DQN(obs_dim=84, n_actions=6)
policy = Policy(model)

runner = runners.EnvRunner(
    env, policy, nsteps=5,
    transforms=[
        ComputeTargetValues(policy),
    ])

trajectory = runner.get_next()<jupyter_output><empty_output><jupyter_text>After computing value targets we will transform lists of interactions into tensorswith the first dimension `batch_size` which is equal to `T * nenvs`, i.e. you essentially needto flatten the first two dimensions.<jupyter_code>class MergeTimeBatch:
    """ Merges first two axes typically representing time and env batch. """
    def __call__(self, trajectory):
        # Modify trajectory inplace.
        #<TODO: implement>
        pass
def f(trajectory):
    # agent takes `nsteps` actions
    nsteps = len(trajectory['actions'])

    # discounted rewards forward view
    discounted_rewards = []
    r = policy.get_value(
        trajectory['state']['latest_observation']).squeeze(dim=1)
    
    print(r.shape)
    for i in range(nsteps-1, -1, -1):
        immediate_reward = trajectory['rewards'][i]
        print('== reward shape ==', immediate_reward.shape)
        
        print('== reset shape ==', trajectory['resets'][i].shape)
        immediate_reward = torch.Tensor(immediate_reward, 
                                        device=device)
        r = immediate_reward + 0.99 * r * trajectory['resets'][i]
        discounted_rewards.append(r)
        
        print('== target reward shape ==', r.shape)
    # reverse discounted rewards
    trajectory['discounted_rewards'] = [
        r for r in reversed(discounted_rewards)]
    return<jupyter_output><empty_output><jupyter_text>Now is the time to implement the advantage actor critic algorithm itself. You can look into your lecture,[Mnih et al. 2016](https://arxiv.org/abs/1602.01783) paper, and [lecture](https://www.youtube.com/watch?v=Tol_jw5hWnI&list=PLkFD6_40KJIxJMR-j5A1mkxK26gh_qg37&index=20) by <NAME>. Policy objective along one trajectory$$\frac{1}{T}\sum_{i=t_1}^{T} \log(\pi(a_i | s_i)) (R_i - V(s_i, \theta_v))$$The length of the trajectory is $T$. Maximize likelihood of selected actions based on its advantage. For distributed RL, just take average among all workers<jupyter_code># policy loss along one trajectory

# estimated value into one tensor of shape (n_step, n_workers)
def transform_estimated_values(estimated_values):
    """
    estimated_values: estimated values from env runner
    """
    batched = [x.squeeze(1) for x in estimated_values]
    return torch.vstack(batched)
transform_estimated_values(trajectory['estimated_values'])
# transform target values r + \gamma v(next state) into shape (n_steps, n_workers)
trajectory['target_values']
def transform_target_values(target_values):
    """transform target values r + \gamma v(next state) into shape (n_steps, n_workers)
    
    target_values: target values along the trajectory for all workers from env runner
    """
    return torch.vstack(target_values)

transform_target_values(trajectory['target_values'])
trajectory['log_probs']
def transform_log_probs(log_probs):
    """
    log_probs: log probabilities of actions for all workers from env runner
    """
    return torch.vstack(log_probs)

transform_log_probs(trajectory['log_probs'])
# Compute policy loss
target_values = transform_target_values(trajectory['target_values'])
estimated_values = transform_estimated_values(trajectory['estimated_values'])

# advantages
adv = target_values - estimated_values

# log probabilities
log_probs = transform_log_probs(trajectory['log_probs'])

# policy loss
policy_loss = -torch.mean(log_probs * adv)

print(policy_loss)<jupyter_output>tensor(-0.0856, grad_fn=<NegBackward>)<jupyter_text>Loss for criticWant to minimize the difference between estimated values and target values (just like q-learning)$$\frac{1}{T} \sum_{i=t_1}^T (R_i - V(s_i, \theta))^2$$<jupyter_code># Compute Value loss
target_values = transform_target_values(trajectory['target_values'])
estimated_values = transform_estimated_values(trajectory['estimated_values'])

value_loss = torch.mean((target_values - estimated_values)**2)
print(value_loss)

# Package everything in an object
class A2C:
    def __init__(self,
                 policy,
                 optimizer,
                 value_loss_coef=0.25,
                 entropy_coef=0.01,
                 max_grad_norm=0.5):
        self.policy = policy
        self.optimizer = optimizer
        self.value_loss_coef = value_loss_coef
        self.entropy_coef = entropy_coef
        self.max_grad_norm = max_grad_norm
    
    def policy_loss(self, trajectory):
        # You will need to compute advantages here.
        
        target_values = transform_target_values(trajectory['target_values'])
        estimated_values = transform_estimated_values(trajectory['estimated_values'])

        # advantages
        adv = target_values - estimated_values

        # log probabilities
        log_probs = transform_log_probs(trajectory['log_probs'])

        # policy loss
        policy_loss = -torch.mean(log_probs * adv)
        return policy_loss
    
    def value_loss(self, trajectory):
        target_values = transform_target_values(trajectory['target_values'])
        estimated_values = transform_estimated_values(trajectory['estimated_values'])
        value_loss = torch.mean((target_values - estimated_values)**2)
        return value_loss
        
    
    def loss(self, trajectory):
        return self.policy_loss(trajectory) + self.value_loss(trajectory)
      
    def step(self, trajectory):
        loss = self.loss(trajectory)
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        return loss.detach().cpu().numpy().item()
# evaluate the agent

def evaluate(agent, env, n_games=1):
    """Plays an a game from start till done, returns per-game rewards 
    agent: an agent with .choose_action interface
    test
    
    env: a testing env (no parallelism)
    """

    game_rewards = []
    for _ in range(n_games):
        state = env.reset()

        total_reward = 0
        while True:
            action = agent.choose_action(state)
            state, reward, done, info = env.step(action)
            total_reward += reward.item()
            if done:
                break

        game_rewards.append(total_reward)
    return game_rewards



model = DQN(obs_dim=84, n_actions=6)
policy = Policy(model)

single_env = nature_dqn_env("SpaceInvadersNoFrameskip-v4", nenvs=1, 
                     summaries='Numpy')

game_rewards = evaluate(policy, single_env, n_games=10)

print('Sampled game rewards: {}'.format(game_rewards))<jupyter_output>Sampled game rewards: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]<jupyter_text>Now you can train your model. With reasonable hyperparameters training on a single GTX1080 for 10 million steps across all batched environments (which translates to about 5 hours of wall clock time)it should be possible to achieve *average raw reward over last 100 episodes* (the average is taken over 100 last episodes in each environment in the batch) of about 600. You should plot this quantity with respect to `runner.step_var` &mdash; the number of interactions with all environments. It is highly encouraged to also provide plots of the following quantities (these are useful for debugging as well):* [Coefficient of Determination](https://en.wikipedia.org/wiki/Coefficient_of_determination) between value targets and value predictions* Entropy of the policy $\pi$* Value loss* Policy loss* Value targets* Value predictions* Gradient norm* Advantages* A2C lossFor optimization we suggest you use RMSProp with learning rate starting from 7e-4 and linearly decayed to 0, smoothing constant (alpha in PyTorch and decay in TensorFlow) equal to 0.99 and epsilon equal to 1e-5.<jupyter_code>#https://stackoverflow.com/questions/58686400/can-not-get-pytorch-working-with-tensorboard
from torch.utils.tensorboard import SummaryWriter

model = DQN(obs_dim=84, n_actions=6)
policy = Policy(model)

env = nature_dqn_env("SpaceInvadersNoFrameskip-v4", nenvs=nenvs, 
                     summaries='Numpy')

test_env = nature_dqn_env("SpaceInvadersNoFrameskip-v4", nenvs=1, 
                     summaries='Numpy')

env.reset()
test_env.reset()

runner = runners.EnvRunner(
    env, policy, nsteps=5,
    transforms=[
        ComputeTargetValues(policy),
    ])

optimizer = torch.optim.RMSprop(policy.model.parameters(), 
                               lr=7e-4, 
                               alpha=0.99,
                               eps=1e-5)

a2c = A2C(policy, optimizer)
writer = SummaryWriter(log_dir='a2c_log')
n_iters = 1000*100

losses = []
rewards = []
for i in range(1, n_iters+1):
    trajectory = runner.get_next()
    loss = a2c.step(trajectory)
    if i % 1000 == 0:
        print('loss at iter {} is: {}'.format(i, loss))
        
        # look at agent's performance on test env
        game_rewards = evaluate(policy, test_env, n_games=3)
        avg_reward = sum(game_rewards) / len(game_rewards)
        print('reward at iter {} is: {}'.format(i, avg_reward))
        
        losses.append(loss)
        rewards.append(avg_reward)
        
        writer.add_scalar('loss', loss, i)
        writer.add_scalar('rewards', avg_reward, i)<jupyter_output>loss at iter 1000 is: -0.28313541412353516
reward at iter 1000 is: 5.666666666666667
loss at iter 2000 is: -0.4396350085735321
reward at iter 2000 is: 5.666666666666667
loss at iter 3000 is: -0.21510395407676697
reward at iter 3000 is: 5.666666666666667
loss at iter 4000 is: -0.7474871873855591
reward at iter 4000 is: 5.666666666666667
loss at iter 5000 is: -0.5137770175933838
reward at iter 5000 is: 5.666666666666667
loss at iter 6000 is: -0.7505146265029907
reward at iter 6000 is: 6.0
loss at iter 7000 is: -0.46364039182662964
reward at iter 7000 is: 5.666666666666667
loss at iter 8000 is: -0.21803009510040283
reward at iter 8000 is: 5.666666666666667
loss at iter 9000 is: -0.23464012145996094
reward at iter 9000 is: 5.666666666666667
loss at iter 10000 is: 0.12844032049179077
reward at iter 10000 is: 5.666666666666667
loss at iter 11000 is: -0.16093593835830688
reward at iter 11000 is: 5.666666666666667
loss at iter 12000 is: -0.10033877193927765
reward at iter 12000 is: 5.666666666[...]
{'fraction_non_alphanumeric': 0.0907871864911845, 'fraction_numerical': 0.03749689595232183, 'mean_word_length': 3.1330049261083746, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 33, 'https://': 5, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 11, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '490954', 'license': 'Unlicense', 'n_tokens_mistral': 6604, 'n_tokens_neox': 6005, 'n_words': 1885, 'path': 'week06_policy_based/a2c-optional.ipynb', 'repo_name': 'hsl89/Practical_RL', 'size': '20142'}