pytorch-a2c-ppo-acktr/main.py

import gym
import os
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from baselines.common.vec_env.subproc_vec_env import SubprocVecEnv
from envs import make_env
from model import ActorCritic
from torch.autograd import Variable
import argparse

parser = argparse.ArgumentParser(description='A3C')
parser.add_argument('--lr', type=float, default=7e-4,
                    help='learning rate (default: 7e-4)')
parser.add_argument('--eps', type=float, default=1e-5,
                    help='RMSprop optimizer epsilon (default: 1e-5)')
parser.add_argument('--alpha', type=float, default=0.99,
                    help='RMSprop optimizer apha (default: 0.99)')
parser.add_argument('--gamma', type=float, default=0.99,
                    help='discount factor for rewards (default: 0.99)')
parser.add_argument('--entropy-coef', type=float, default=0.01,
                    help='entropy term coefficient (default: 0.01)')
parser.add_argument('--value-loss-coef', type=float, default=0.5,
                    help='value loss coefficient (default: 0.5)')
parser.add_argument('--max-grad-norm', type=float, default=0.5,
                    help='value loss coefficient (default: 0.5)')
parser.add_argument('--seed', type=int, default=1,
                    help='random seed (default: 1)')
parser.add_argument('--num-processes', type=int, default=16,
                    help='how many training CPU processes to use (default: 16)')
parser.add_argument('--num-steps', type=int, default=5,
                    help='number of forward steps in A2C (default: 5)')
parser.add_argument('--num-stack', type=int, default=4,
                    help='number of frames to stack (default: 4)')
parser.add_argument('--log-interval', type=int, default=10,
                    help='log interval, one log per n updates (default: 10)')
parser.add_argument('--num-frames', type=int, default=10e6,
                    help='number of frames to train (default: 10e6)')
parser.add_argument('--env-name', default='PongNoFrameskip-v4',
                    help='environment to train on (default: PongNoFrameskip-v4)')
parser.add_argument('--log-dir', default='/tmp/gym/',
                    help='directory to save agent logs (default: /tmp/gym)')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='disables CUDA training')


args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()

num_updates = int(args.num_frames) // args.num_steps // args.num_processes

torch.manual_seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)

try:
    os.makedirs(args.log_dir)
except OSError:
    pass

def main():
    os.environ['OMP_NUM_THREADS'] = '1'

    envs = SubprocVecEnv([make_env(args.env_name, args.seed, i, args.log_dir)
                            for i in range(args.num_processes)])

    actor_critic = ActorCritic(envs.observation_space.shape[0] * args.num_stack, envs.action_space)

    if args.cuda:
        actor_critic.cuda()

    optimizer = optim.RMSprop(actor_critic.parameters(), args.lr, eps=args.eps, alpha=args.alpha)
    #optimizer = KFACOptimizer(actor_critic, damping=1e-2, kl_clip=0.01, stat_decay=0.99)

    obs_shape = envs.observation_space.shape
    obs_shape = (obs_shape[0] * args.num_stack, obs_shape[1], obs_shape[2])

    states = torch.zeros(args.num_steps + 1, args.num_processes, *obs_shape)
    current_state = torch.zeros(args.num_processes, *obs_shape)
    counts = 0

    def update_current_state(state):
        state = torch.from_numpy(np.stack(state)).float()
        current_state[:, :-1] = current_state[:, 1:]
        current_state[:, -1] = state

    state = envs.reset()
    update_current_state(state)

    rewards = torch.zeros(args.num_steps, args.num_processes, 1)
    returns = torch.zeros(args.num_steps + 1, args.num_processes, 1)

    actions = torch.LongTensor(args.num_steps, args.num_processes)
    masks = torch.zeros(args.num_steps, args.num_processes, 1)

    # These variables are used to compute average rewards for all processes.
    # Note that rewards are clipped so you need to use a monitor (see envs.py)
    # to get true rewards.
    total_rewards = torch.zeros([args.num_processes, 1])
    final_rewards = torch.zeros([args.num_processes, 1])

    if args.cuda:
        states = states.cuda()
        current_state = current_state.cuda()
        rewards = rewards.cuda()
        returns = returns.cuda()
        actions = actions.cuda()
        masks = masks.cuda()

    for j in range(num_updates):
        for step in range(args.num_steps):
            # Sample actions
            _, logits = actor_critic(Variable(states[step], volatile=True))
            probs = F.softmax(logits)
            log_probs = F.log_softmax(logits).data
            actions[step] = probs.multinomial().data

            cpu_actions = actions[step].cpu()
            cpu_actions = cpu_actions.numpy()

            # Obser reward and next state
            state, reward, done, info = envs.step(cpu_actions)

            reward = torch.from_numpy(np.expand_dims(np.stack(reward), 1)).float()
            total_rewards += reward

            np_masks = np.array([0.0 if done_ else 1.0 for done_ in done])

            # If done then clean the history of observations.
            pt_masks = torch.from_numpy(np_masks.reshape(np_masks.shape[0], 1, 1, 1)).float()
            if args.cuda:
                pt_masks = pt_masks.cuda()
            current_state *= pt_masks

            update_current_state(state)
            states[step + 1].copy_(current_state)
            rewards[step].copy_(reward)
            masks[step].copy_(torch.from_numpy(np_masks))

            final_rewards *= masks[step].cpu()
            final_rewards += (1 - masks[step].cpu()) * total_rewards

            total_rewards *= masks[step].cpu()

        # Reshape to do in a single forward pass for all steps
        values, logits = actor_critic(Variable(states.view(-1, *states.size()[-3:])))
        log_probs = F.log_softmax(logits)
        probs = F.softmax(logits)

        # Unreshape
        logits_size = (args.num_steps + 1, args.num_processes, logits.size(-1))

        log_probs = F.log_softmax(logits).view(logits_size)[:-1]
        probs = F.softmax(logits).view(logits_size)[:-1]

        values = values.view(args.num_steps + 1, args.num_processes, 1)
        logits = logits.view(logits_size)[:-1]

        action_log_probs = log_probs.gather(2, Variable(actions.unsqueeze(2)))

        dist_entropy = -(log_probs * probs).sum(-1).mean()

        returns[-1] = values[-1].data

        for step in reversed(range(args.num_steps)):
            returns[step] = returns[step + 1] * \
                args.gamma * masks[step] + rewards[step]

        value_loss = (values[:-1] - Variable(returns[:-1])).pow(2).mean()

        advantages = returns[:-1] - values[:-1].data
        action_loss = -(Variable(advantages) * action_log_probs).mean()

        optimizer.zero_grad()
        (value_loss * args.value_loss_coef + action_loss - dist_entropy * args.entropy_coef).backward()

        nn.utils.clip_grad_norm(actor_critic.parameters(), args.max_grad_norm)
        optimizer.step()

        states[0].copy_(states[-1])

        if j % args.log_interval == 0:
            print("Updates {}, num frames {}, mean clipped reward {:.5f}, max clipped reward {:.1f}, entropy {:.5f}, value loss {:.5f}, policy loss {:.5f}".format(
                j, j * args.num_processes * args.num_steps, final_rewards.mean(), final_rewards.max(), -dist_entropy.data[0], value_loss.data[0], action_loss.data[0]))

if __name__ == "__main__":
    main()