Created
May 24, 2018 17:09
-
-
Save andrewliao11/3514dfaa4c2c3e9a1cd808ad1ed4bbeb to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import argparse | |
| import gym | |
| import numpy as np | |
| from itertools import count | |
| from collections import namedtuple, deque | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torch.optim as optim | |
| from torch.distributions import Normal | |
| import ipdb | |
| parser = argparse.ArgumentParser(description='PyTorch actor-critic example') | |
| parser.add_argument('--gamma', type=float, default=0.99, metavar='G', | |
| help='discount factor (default: 0.99)') | |
| parser.add_argument('--seed', type=int, default=543, metavar='N', | |
| help='random seed (default: 1)') | |
| parser.add_argument('--render', action='store_true', | |
| help='render the environment') | |
| parser.add_argument('--log-interval', type=int, default=10, metavar='N', | |
| help='interval between training status logs (default: 10)') | |
| args = parser.parse_args() | |
| env = gym.make('InvertedPendulum-v1') | |
| env.seed(args.seed) | |
| torch.manual_seed(args.seed) | |
| SavedAction = namedtuple('SavedAction', ['log_prob', 'value']) | |
| class Policy(nn.Module): | |
| def __init__(self, ob_space, ac_space): | |
| super(Policy, self).__init__() | |
| self.affine1 = nn.Linear(ob_space.shape[0], 128) | |
| self.mean_head = nn.Linear(128, ac_space.shape[0]) | |
| self.logstd_head = nn.Linear(128, ac_space.shape[0]) | |
| self.value_head = nn.Linear(128, 1) | |
| self.saved_actions = [] | |
| self.rewards = [] | |
| def forward(self, x): | |
| x = F.relu(self.affine1(x)) | |
| mean = self.mean_head(x) | |
| logstd = self.logstd_head(x) | |
| std = logstd.exp() | |
| pdparam = (mean, std) | |
| state_values = self.value_head(x) | |
| return pdparam, state_values | |
| model = Policy(env.observation_space, env.action_space) | |
| optimizer = optim.Adam(model.parameters(), lr=3e-4) | |
| eps = np.finfo(np.float32).eps.item() | |
| def select_action(state): | |
| state = torch.from_numpy(state).float() | |
| pdparam, state_value = model(state) | |
| m = Normal(*pdparam) | |
| action = m.sample() | |
| model.saved_actions.append(SavedAction(m.log_prob(action), state_value)) | |
| return action.item() | |
| def finish_episode(): | |
| R = 0 | |
| saved_actions = model.saved_actions | |
| policy_losses = [] | |
| value_losses = [] | |
| rewards = [] | |
| for r in model.rewards[::-1]: | |
| R = r + args.gamma * R | |
| rewards.insert(0, R) | |
| rewards = torch.tensor(rewards) | |
| rewards = (rewards - rewards.mean()) / (rewards.std() + eps) | |
| for (log_prob, value), r in zip(saved_actions, rewards): | |
| reward = r - value.item() | |
| policy_losses.append(-log_prob * reward) | |
| value_losses.append(F.smooth_l1_loss(value, torch.tensor([r]))) | |
| optimizer.zero_grad() | |
| loss = torch.stack(policy_losses).sum() + torch.stack(value_losses).sum() | |
| loss.backward() | |
| optimizer.step() | |
| del model.rewards[:] | |
| del model.saved_actions[:] | |
| def main(): | |
| reward_list = deque([], maxlen=100) | |
| for i_episode in count(1): | |
| state = env.reset() | |
| ep_reward = 0 | |
| for t in range(10000): # Don't infinite loop while learning | |
| action = select_action(state) | |
| state, reward, done, _ = env.step([action]) | |
| ep_reward += reward | |
| if args.render: | |
| env.render() | |
| model.rewards.append(reward) | |
| if done: | |
| reward_list.append(ep_reward) | |
| break | |
| average_reward = sum(reward_list)/len(reward_list) | |
| finish_episode() | |
| if i_episode % args.log_interval == 0: | |
| print('Episode {}\tLast length: {:5d}\tAverage reward: {:.2f}'.format( | |
| i_episode, t, average_reward)) | |
| if __name__ == '__main__': | |
| main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment