added lunar lander agent, ac & pg

Author: Andrewzh112

policy/agent.py

@@ -1,4 +1,6 @@
 import numpy as np
+import torch
+from policy.networks import ActorCritic
 
 
 class BlackJackAgent:
@@ -122,3 +124,75 @@ def update(self, state, action, reward, state_):
 
     def decrease_eps(self):
         self.epsilon = max(0.01, self.epsilon - 1e-5)
+
+
+class PolicyGradientAgent:
+    def __init__(self, input_dim, action_dim, hidden_dim, gamma, lr):
+        self.gamma = gamma
+        self.policy = ActorCritic(*input_dim, action_dim, hidden_dim)
+        self.optimizer = torch.optim.Adam(self.policy.parameters(), lr)
+        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        self.reward_history, self.action_logprob_history = [], []
+
+    def choose_action(self, state):
+        state = torch.from_numpy(state).to(self.device)
+        action_proba = torch.softmax(self.policy(state), dim=-1)
+        action_dist = torch.distributions.Categorical(action_proba)
+        action = action_dist.sample()
+        if self.policy.training:
+            log_probas = action_dist.log_prob(action)
+            self.action_logprob_history.append(log_probas)
+        return action.item()
+
+    def store_reward(self, reward):
+        self.reward_history.append(reward)
+
+    def update(self):
+        # calculate MC returns & loss
+        T = len(self.reward_history)
+        discounts = torch.logspace(0, T, steps=T + 1, base=self.gamma, device=self.device)[:T]
+        returns = torch.tensor([torch.tensor(
+            self.reward_history[t:], dtype=torch.float, device=self.device) @ discounts[t:] for t in range(T)])
+        loss = 0
+        for g, log_prob in zip(returns, self.action_logprob_history):
+            loss += - g * log_prob
+
+        # sgd + reset history
+        self.optimizer.zero_grad()
+        loss.backward()
+        self.optimizer.step()
+        self.reward_history, self.action_logprob_history = [], []
+
+
+class ActorCriticAgent:
+    def __init__(self, input_dim, action_dim, hidden_dim, gamma, lr):
+        self.gamma = gamma
+        self.actor_critic = ActorCritic(*input_dim, action_dim, hidden_dim)
+        self.optimizer = torch.optim.Adam(self.actor_critic.parameters(), lr)
+        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+        self.log_proba, self.value = None, None
+
+
+    def choose_action(self, state):
+        state = torch.from_numpy(state).to(self.device)
+        self.value, action_logits = self.actor_critic(state)
+        action_proba = torch.softmax(action_logits, dim=-1)
+        action_dist = torch.distributions.Categorical(action_proba)
+        action = action_dist.sample()
+        self.log_proba = action_dist.log_prob(action)
+        return action.item()
+
+    def update(self, reward, state_, done):
+        # calculate TD loss
+        state_ = torch.from_numpy(state_).unsqueeze(0).to(self.device)
+        value_, _ = self.actor_critic(state_)
+        critic_loss = (reward + self.gamma * value_ * ~done - self.value).pow(2)
+
+        # actor loss
+        actor_loss = - self.value.detach() * self.log_proba
+
+        # sgd + reset history
+        loss = critic_loss + actor_loss
+        self.optimizer.zero_grad()
+        loss.backward()
+        self.optimizer.step()

policy/lunarlander/main.py

@@ -0,0 +1,53 @@
+import gym
+import argparse
+import torch
+import numpy as np
+from tqdm import tqdm
+from collections import deque
+from policy import agent as Agent
+
+
+parser = argparse.ArgumentParser(description='Lunar Lander Agents')
+parser.add_argument('--agent', type=str, default='Actor Critic', help='Agent style')
+parser.add_argument('--n_episodes', type=int, default=3000, help='Number of episodes you wish to run for')
+parser.add_argument('--hidden_dim', type=int, default=2048, help='Hidden dimension of FC layers')
+parser.add_argument('--lr', '--learning_rate', type=float, default=1e-4, help='Learning rate for Adam optimizer')
+parser.add_argument('--gamma', type=float, default=0.99, help='Reward discount factor')
+
+parser.add_argument('--render', action="store_true", default=False, help='Render environment while training')
+parser.add_argument('--window_legnth', type=int, default=100, help='Length of window to keep track scores')
+args = parser.parse_args()
+
+
+def main():
+    env = gym.make('LunarLander-v2')
+    agent_ = getattr(Agent, args.agent.replace(' ', '') + 'Agent')
+    agent = agent_(input_dim=env.observation_space.shape,
+                   action_dim=env.action_space.n,
+                   hidden_dim=args.hidden_dim,
+                   gamma=args.gamma,
+                   lr=args.lr)
+    pbar = tqdm(range(args.n_episodes))
+    score_history = deque(maxlen=args.window_legnth)
+    for e in pbar:
+        done, score, observation = False, 0, env.reset()
+        while not done:
+            if args.render:
+                env.render()
+            action = agent.choose_action(observation)
+            next_observation, reward, done, _ = env.step(action)
+            if args.agent == 'Actor Critic':
+                agent.update(reward, next_observation, done)
+            else:
+                agent.store_reward(reward)
+            observation = next_observation
+            score += reward
+        if args.agent == 'Policy Gradient':
+            agent.update()
+        score_history.append(score)
+        tqdm.write(
+            f'Episode: {e + 1}/{args.n_episodes}, Score: {score}, Average Score: {np.mean(score_history)}')
+
+
+if __name__ == '__main__':
+    main()

policy/networks.py

@@ -0,0 +1,19 @@
+import torch
+from torch import nn
+
+
+class ActorCritic(nn.Module):
+    def __init__(self, input_dim, n_actions, hidden_dim):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(True),
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.ReLU(True)
+        )
+        self.v = nn.Linear(hidden_dim // 2, 1)
+        self.pi = nn.Linear(hidden_dim // 2, n_actions)
+
+    def forward(self, state):
+        features = self.encoder(state)
+        return self.v(features), self.pi(features)