<jupyter_start><jupyter_text>Enable GPU<jupyter_code>import torch
device = torch.device('cuda:0' if torch.cuda.is_available else 'cpu')<jupyter_output><empty_output><jupyter_text>Actor and Critic Network<jupyter_code>import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical

class Actor_Net(nn.Module):
  def __init__(self, input_dims, output_dims, num_neurons = 128):
    super(Actor_Net, self).__init__()
    self.fc1 = nn.Linear(input_dims, num_neurons)
    self.actor = nn.Linear(num_neurons, output_dims)
    self.log_probs = []
    self.entropies = []

  def forward(self, state):
    x = F.relu(self.fc1(state))
    x = F.softmax(self.actor(x), dim = 1)

    return x

  def get_action(self, state):
    with torch.no_grad():
      probs = self.forward(state)
      dist = Categorical(probs = probs)
      action = dist.sample()
    return action
  
  def eval_action(self, state):
    probs = self.forward(state)
    dist = Categorical(probs = probs)
    action = dist.sample().to(device)
    log_prob = dist.log_prob(action)
    entropy = dist.entropy()
    self.log_probs.append(log_prob)
    self.entropies.append(entropy)

    return action

class Critic_Net(nn.Module):
  def __init__ (self, input_dims, output_dims, num_neurons = 128):
    super(Critic_Net, self).__init__()

    self.values = []
    self.next_values = []
    
    self.fc1 = nn.Linear(input_dims, num_neurons)
    self.critic = nn.Linear(num_neurons, 1)

  def forward (self, state):
    x = F.relu(self.fc1(state))
    x = self.critic(x)

    return x
import torch.optim as optim
import numpy as np
import gym

class Actor_Critic_Agent(nn.Module):
  def __init__(self, input_dims, output_dims, optimizer = 'RMSprop', num_neurons = 128 , gamma = 0.99, actor_lr=0.001, critic_lr = 0.01):
    super(Actor_Critic_Agent, self).__init__()
    self.actor_net = Actor_Net(input_dims= input_dims, output_dims= output_dims, num_neurons= num_neurons).to(device)
    self.critic_net = Critic_Net(input_dims=input_dims, output_dims= output_dims, num_neurons= num_neurons).to(device)
    self.gamma = gamma
    if optimizer == 'RMSprop':
      self.actor_optimizer = optim.RMSprop(params = self.actor_net.parameters(), lr =actor_lr)
      self.critic_optimizer = optim.RMSprop(params = self.critic_net.parameters(), lr = critic_lr)
    else:
      self.actor_optimizer = optim.Adam(params = self.actor_net.parameters(), lr = actor_lr)
      self.critic_optimizer = optim.Adam(params = self.critic_net.parameters(), lr = critic_lr)

  def learn_mean(self, rewards, dones):
    value_criteration = nn.MSELoss()
    value_losses = []
    actor_losses = []
    self.critic_net.next_values = torch.cat(self.critic_net.next_values, dim = 0).squeeze(0)
    self.critic_net.values = torch.cat(self.critic_net.values, dim = 0).squeeze(0)
    self.actor_net.log_probs = torch.cat(self.actor_net.log_probs, dim = 0)
    self.actor_net.entropies = torch.cat(self.actor_net.entropies, dim = 0)

    for reward, entropy, log_prob, v, v_next, done in zip(rewards ,self.actor_net.entropies, self.actor_net.log_probs, self.critic_net.values, self.critic_net.next_values, dones):
      td_target = reward + self.gamma * v_next * done
      td_error = td_target - v
      value_loss = value_criteration(v, td_target.detach())- 0.001 * entropy.detach()
      actor_loss = - log_prob * td_error.detach() 
      value_losses.append(value_loss)
      actor_losses.append(actor_loss)

    self.critic_optimizer.zero_grad()
    value_losses = torch.stack(value_losses).sum()
    value_losses.backward()
    self.critic_optimizer.step()  

    self.actor_optimizer.zero_grad()
    actor_losses = torch.stack(actor_losses).sum()
    actor_losses.backward()
    self.actor_optimizer.step()

      
    # clear out memory 
    self.actor_net.log_probs = []
    self.actor_net.entropies = []
    self.critic_net.values = []
    self.critic_net.next_values = []<jupyter_output><empty_output><jupyter_text>Without Wandb<jupyter_code>import gym
import time
import pdb

env = gym.make('CartPole-v1')
env.seed(543)
torch.manual_seed(543)
state_dims = env.observation_space.shape[0]
action_dims = env.action_space.n
agent = Actor_Critic_Agent(input_dims= state_dims, output_dims = action_dims)

def train():

  num_ep = 2000
  print_every = 100
  running_score = 10
  start = time.time()

  rewards = []
  dones = []

  for ep in range(1, num_ep + 1):
    state = env.reset()
    score = 0
    done = False
    rewards = []
    dones = []

    while not done:
      state = torch.tensor([state]).float().to(device)
      action = agent.actor_net.eval_action(state)
      v = agent.critic_net(state)

      next_state, reward, done, _ = env.step(action.item())
      v_next = agent.critic_net(torch.tensor([next_state]).float().to(device))
  
      agent.critic_net.values.append(v.squeeze(0))
      agent.critic_net.next_values.append(v_next.squeeze(0))
      rewards.append(reward)
      dones.append(1 - done)
      
      # update episode
      score += reward
      state = next_state

      if done:
        break

    # update agent
    #pdb.set_trace()
    agent.learn_mean(rewards,dones)
   
    # calculating score and running score
    running_score = 0.05 * score + (1 - 0.05) * running_score

    if ep % print_every == 0:
      print('episode: {}, running score: {}, time elapsed: {}'.format(ep, running_score, time.time() - start))




train() #RMS<jupyter_output>episode: 100, running score: 43.32507441570408, time elapsed: 4.842878341674805
episode: 200, running score: 129.30332722904944, time elapsed: 19.552313089370728<jupyter_text>Wtih wandb<jupyter_code>!pip install wandb
!wandb login

import wandb
sweep_config = dict()
sweep_config['method'] = 'grid'
sweep_config['metric'] = {'name': 'running_score', 'goal': 'maximize'}
sweep_config['parameters'] = {'learning': {'value': 'learn_mean'}, 'actor_learning_rate': {'values' : [0.01, 0.001, 0.0001,0.0003,0.00001]}, 'critic_learning_rate' : {'values': [0.01, 0.001, 0.0001, 0.0003, 0.00001]}
                              , 'num_neurons': {'value': 128 }, 'optimizer': {'values' : ['RMSprop', 'Adam']}}

sweep_id = wandb.sweep(sweep_config, project = 'Advantage_Actor_Critic')
import gym 
import torch
import time
import wandb



def train():
  wandb.init(config = {'env':'CartPole-v1','algorithm:': 'Actor_Critic','architecture': 'seperate','num_laeyrs':'2'}, project = 'Advantage_Actor_Critic',group = 'Cart_128_neurons_2_layer')
  config = wandb.config

  env = gym.make('CartPole-v1')
  env.seed(543)
  torch.manual_seed(543)

  state_dim = env.observation_space.shape[0]
  action_dim = env.action_space.n

  device = torch.device('cuda:0' if torch.cuda.is_available else 'cpu')
  agent = Actor_Critic_Agent(input_dims= state_dim, output_dims= action_dim, optimizer = config.optimizer, num_neurons= config.num_neurons, actor_lr = config.actor_learning_rate, critic_lr = config.critic_learning_rate)


  num_ep = 3000
  print_interval = 100
  save_interval = 1000
  running_score = 10
  start = time.time()

  
  wandb.watch(agent)
  for ep in range(1,num_ep+1):
    state = env.reset()
    score = 0
    done = False
    rewards = []
    dones = []

    while not done:
      state = torch.tensor([state]).float().to(device)
      action = agent.actor_net.eval_action(state)
      v = agent.critic_net(state)

      next_state, reward, done, _ = env.step(action.item())
      v_next = agent.critic_net(torch.tensor([next_state]).float().to(device))
  
      agent.critic_net.values.append(v.squeeze(0))
      agent.critic_net.next_values.append(v_next.squeeze(0))
      rewards.append(reward)
      dones.append(1 - done)
      
      # update episode
      score += reward
      state = next_state

      if done:
        break

    # update agent
    agent.learn_mean(rewards,dones)
   
    # calculating score and running score
    running_score = 0.05 * score + (1 - 0.05) * running_score

    wandb.log({'episode': ep, 'running_score': running_score})   

    if ep % print_interval == 0:
      print('episode {} average reward {}, ended at {:.01f}'.format(ep, running_score, time.time() - start))    
    
    if ep % save_interval == 0:
      save_name_actor = 'actor_' + str(ep) + '.pt'
      torch.save(agent.actor_net.state_dict(),save_name_actor)
      save_name_critic = 'critic_' + str(ep) + '.pt'
      torch.save(agent.critic_net.state_dict(),save_name_critic)
      wandb.save(save_name_actor)
      wandb.save(save_name_critic)

    if ep == num_ep:
      dummy_input = torch.rand(1,4).to(device)
      torch.onnx.export(agent.actor_net,dummy_input,'final_model_actor.onnx')
      wandb.save('final_model_actor.onnx')
      torch.onnx.export(agent.critic_net, dummy_input, 'final_model_critic.onnx')
      wandb.save('final_model_critic.onnx')
    
wandb.agent(sweep_id, train)<jupyter_output>[34m[1mwandb[0m: Agent Starting Run: wivnmds7 with config:
[34m[1mwandb[0m: 	actor_learning_rate: 0.01
[34m[1mwandb[0m: 	critic_learning_rate: 0.01
[34m[1mwandb[0m: 	learning: learn_mean
[34m[1mwandb[0m: 	num_neurons: 128
[34m[1mwandb[0m: 	optimizer: RMSprop
[34m[1mwandb[0m: [33mWARNING[0m Ignored wandb.init() arg project when running a sweep
{'fraction_non_alphanumeric': 0.13198959687906373, 'fraction_numerical': 0.02990897269180754, 'mean_word_length': 3.799271970878835, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '1535', 'license': 'MIT', 'n_tokens_mistral': 3617, 'n_tokens_neox': 3384, 'n_words': 604, 'path': 'TD Actor Critic/TD_Actor_Critic_seperate_net.ipynb', 'repo_name': 'gt-coar/BrianSURE2021', 'size': '9228'}