text: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from nn_model.replay_memory import ReplayBuffer
  
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  
  
  class Actor(nn.Module):
      def __init__(self, state_dim, action_dim, max_action=1):
          super(Actor, self).__init__()
  
          self.l1 = nn.Linear(state_dim, 256)
          self.l2 = nn.Linear(256, 256)
          self.l3 = nn.Linear(256, action_dim)
  
          self.max_action = max_action
  
      def forward(self, state):
          a = F.relu(self.l1(state))
          a = F.relu(self.l2(a))
          return self.max_action * torch.tanh(self.l3(a))
  
      def save_checkpoint(self):
          print('... saving checkpoint ...')
          torch.save(self.state_dict(), self.checkpoint_file)
  
      def load_checkpoint(self):
          print('... loading checkpoint ...')
          self.load_state_dict(torch.load(self.checkpoint_file))
  
  
  class Critic(nn.Module):
      def __init__(self, state_dim, action_dim):
          super(Critic, self).__init__()
  
          # Q1 architecture
          self.l1 = nn.Linear(state_dim + action_dim, 256)
          self.l2 = nn.Linear(256, 256)
          self.l3 = nn.Linear(256, 1)
  
      def forward(self, state, action):
          sa = torch.cat([state, action], 1)
  
          q2 = F.relu(self.l4(sa))
          q2 = F.relu(self.l5(q2))
          q2 = self.l6(q2)
          return q2
  
      def Q1(self, state, action):
          sa = torch.cat([state, action], 1)
  
          q1 = F.relu(self.l1(sa))
          q1 = F.relu(self.l2(q1))
          q1 = self.l3(q1)
          return q1
  
      def save_checkpoint(self):
          print('... saving checkpoint ...')
          torch.save(self.state_dict(), self.checkpoint_file)
  
      def load_checkpoint(self):
          print('... loading checkpoint ...')
          self.load_state_dict(torch.load(self.checkpoint_file))
  
  
  class AC(object):
      """ Agent class for use with separate actor and critic networks.
          This is appropriate for very simple environments, such as the mountaincar
      """
  
      def __init__(self,
                   state_dim,
                   action_dim,
                   max_action,
                   batch_size=64,
                   discount=0.99,
                   policy_freq=2,
                   policy_noise=0.2,
                   noise_clip=0.5, ):
          self.discount = discount
          self.batch_size = batch_size
          self.max_action = max_action
          self.policy_noise = policy_noise
          self.noise_clip = noise_clip
  
          self.actor = Actor(state_dim, action_dim, max_action)
          self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=3e-4)
  
          self.critic = Critic(state_dim, action_dim)
          self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=3e-4)
  
          self.replay_buffer = ReplayBuffer(state_dim, action_dim)
  
          self.policy_freq = policy_freq
          self.total_it = 0
  
      def store_transition(self, curr_state, action, next_state, reward, done):
          self.replay_buffer.add(curr_state, action, next_state, reward, done)
  
      def select_action(self, state):
          state = torch.FloatTensor(state.reshape(1, -1)).to(device)
          return self.actor(state).cpu().data.numpy().flatten()
  
      def train(self):
          self.total_it += 1
  
          # Sample replay buffer
          state, action, next_state, reward, not_done = self.replay_buffer.sample(self.batch_size)
  
          with torch.no_grad():
              # Select action according to policy and add clipped noise
              noise = (
                      torch.randn_like(action) * self.policy_noise
              ).clamp(-self.noise_clip, self.noise_clip)
  
              next_action = (
                      self.actor(next_state) + noise
              ).clamp(-self.max_action, self.max_action)
  
              # Compute the target Q value
              target_Q = self.critic(next_state, next_action)
              target_Q = reward + self.discount * target_Q
  
          # Get current Q estimates
          current_Q = self.critic(state, action)
  
          # Compute critic loss
          critic_loss = F.mse_loss(current_Q, target_Q)
  
          # Optimize the critic
          self.critic_optimizer.zero_grad()
          critic_loss.backward()
          self.critic_optimizer.step()
  
          # Delayed policy updates
          if self.total_it % self.policy_freq == 0:
              # Compute actor losse
              actor_loss = -self.critic.Q1(state, self.actor(state)).mean()
  
              # Optimize the actor
              self.actor_optimizer.zero_grad()
              actor_loss.backward()
              self.actor_optimizer.step()
  
      def save_models(self):
          self.actor.save_checkpoint()
          self.critic.save_checkpoint()
  
      def load_models(self):
          self.actor.load_checkpoint()
          self.critic.load_checkpoint()