<jupyter_start><jupyter_text><jupyter_code>!pip install pybullet --upgrade

#see PyBullet Quickstart Guide at
#https://docs.google.com/document/d/10sXEhzFRSnvFcl3XxNGhnD4N2SedqwdAvK3dsihxVUA/edit
import os
import time
import random
import numpy as np
import matplotlib.pyplot as plt
import pybullet_envs
import gym
import torch
import torch.nn as nn
import torch.nn.functional as F
from gym import wrappers
from torch.autograd import Variable
from collections import deque<jupyter_output><empty_output><jupyter_text>STEP1<jupyter_code>class ReplayBuffer(object):
  def __init__(self, max_size = 1e6):
    self.storage = []
    self.max_size = max_size
    self.ptr = 0
  
  def add(self, transition):
    if len(self.storage) == self.max_size:
      self.storage[int(self.ptr)] = transition
      self.ptr = (self.ptr + 1) % self.max_size
    else:
      self.storage.append(transition)

  def sample(self, batch_sze):
    ind = np.random.randint(0, len(self.storage), batch_size)
    batch_states, batch_next_states, batch_actions, batch_rewards, batch_dones = [], [], [], [], []
    for i in ind:
      state, next_state, action, reward, done = self.storage[i]
      batch_states.append(np.array(state, copy = False))
      batch_next_states.append(np.array(next_state, copy = False))
      batch_actions.append(np.array(action, copy = False))
      batch_rewards.append(np.array(reward, copy = False))
      batch_dones.append(np.array(done, copy = False))
    return np.array(batch_states), np.array(batch_next_states), np.array(batch_actions), np.array(batch_rewards).reshape(-1,1), np.array(batch_dones).reshape(-1,1)<jupyter_output><empty_output><jupyter_text>STEP2<jupyter_code>class Actor(nn.Module):

  def __init__(self, state_dims, action_dim, max_action):
    super(Actor, self).__init__()
    self.layer_1 = nn.Linear(state_dims, 400)
    self.layer_2 = nn.Linear(400, 300)
    self.layer_3 = nn.Linear(300, action_dim)
    self.max_action = max_action

  def forward(self, x):
    x = F.relu(self.layer_1(x))
    x = F.relu(self.layer_2(x))
    x = self.max_action * torch.tanh(self.layer_3(x))
    return x<jupyter_output><empty_output><jupyter_text>STEP3<jupyter_code>class Critic(nn.Module):

  def __init__(self, state_dims, action_dim):
    super(Critic, self).__init__()
    # First Critic Network
    self.layer_1 = nn.Linear(state_dims + action_dim, 400)
    self.layer_2 = nn.Linear(400, 300)
    self.layer_3 = nn.Linear(300, action_dim)
    # Second Critic Network
    self.layer_4 = nn.Linear(state_dims + action_dim, 400)
    self.layer_5 = nn.Linear(400, 300)
    self.layer_6 = nn.Linear(300, action_dim)

  def forward(self, x, u):
    xu = torch.cat([x, u], 1)

    x1 = F.relu(self.layer_1(xu))
    x1 = F.relu(self.layer_2(x1))
    x1 = self.layer_3(x1)

    x2 = F.relu(self.layer_4(xu))
    x2 = F.relu(self.layer_5(x2))
    x2 = self.layer_6(x2)

    return x1, x2

  def Q1(self, x, u):
    xu = torch.cat([x, u], 1)
    x1 = F.relu(self.layer_1(xu))
    x1 = F.relu(self.layer_2(x1))
    x1 = self.layer_3(x1)
    return x1<jupyter_output><empty_output><jupyter_text>STEP 4-15<jupyter_code>device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class T3D(object):
  def __init__(self, state_dims, action_dim, max_action):
    self.actor = Actor(state_dims, action_dim, max_action).to(device)
    self.actor_target = Actor(state_dims, action_dim, max_action).to(device)
    self.actor_target_load_state_dict(self.actor.state_dict)

    self.actor_optimizer = torch.optim.Adam(self.actor.parameters())


    self.critic = Critic(state_dims, action_dim).to(device)
    self.critic_target = critic(state_dims, action_dim).to(device)
    self.critic_target.load_state_dict(self.critic.state_dict)

    self.critic_optimizer = torch.optim.Adam(self.critic.parameters())
    self.max_action = max_action

  def select_action(self, state):
    state = torch.Tensor(state.reshape(1, -1)).to(device)
    return self.actor(state).cpu().data.numpy().flatten()

  ##### STEP4
  def train(self, replay_buffer, iterations, batch_size = 100, discount = 0.99, tau = 0.005, policy_noise = 0.2, noise_clip = 0.5, policy_freq = 2):
    for it in range(iterations):
      batch_states, batch_next_states, batch_actions, batch_rewards, batch_dones = replay_buffer.sample(batch_size)
      state = torch.Tensor(batch_states).to(device)
      next_state = torch.Tensor(batch_next_states).to(device)
      action = torch.Tensor(batch_actions).to(device)
      reward = torch.Tensor(batch_rewards).to(device)
      done = torch.Tensor(batch_dones).to(device)
      ##### STEP5
      next_action = self.actor_target.forward(next_state)
      ##### STEP6
      noise = torch.Tensor(batch_actions).data.normal_(0, policy_noise).to(device)
      noise = noise.clamp(-noise_clip, noise_clip)
      next_action = (next_action + noise).clamp(-self.max_action, self.max_action)
      ##### STEP7
      target_Q1, target_Q2 = self.critic_target.forward(net_state, next_action)
      ##### STEP8
      target_Q = torch.min(target_Q1, target_Q2)
      ##### STEP9
      target_Q = reward + ((1 - done) * discount * target_Q).detach()
      ##### STEP10
      current_Q1, current_Q2 = self.critic.forward(state, action)
      ##### STEP11
      critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(current_Q2, target_Q)
      ##### STEP12
      self.critic_optimizer.zero_grad()
      critic_loss.backward()
      self.critic_optimizer.step()
      ##### STEP13
      if it % policy_freq == 0:
        actor_loss = -(self.critic.Q1(state, self.actor(state)).mean())
        self.actor_optimizer.grad_zero()
        actor_loss.backward()
        self.actor_optimizer.step()
        ##### STEP14
        for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()):
          target_param.data.copy_(taru * param.data + (1 - tau) * target_param.data)
        ##### STEP15
        for param, target_param in zip(self.critic.parameters(), self.critic_target.parameters()):
          target_param.data.copy_(taru * param.data + (1 - tau) * target_param.data)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1351483534398435, 'fraction_numerical': 0.02429083795239648, 'mean_word_length': 3.8923444976076556, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 1, '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': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '104082', 'license': 'Apache-2.0', 'n_tokens_mistral': 2314, 'n_tokens_neox': 2228, 'n_words': 380, 'path': 'Phase2/Project9/T3D.ipynb', 'repo_name': 'ksasi/Deep_Vision', 'size': '6134'}