algorithm.py

import os
import random
import re
import time
from copy import deepcopy

import gymnasium as gym
import matplotlib.pyplot as plt
import numpy as np
import pynvml
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from spikingjelly.activation_based import (functional, layer, learning, neuron,
                                           surrogate)
from torch._prims_common import check
from torch.cuda import seed
from torch.distributions import transform_to
from tqdm import tqdm

device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
pynvml.nvmlInit()
gpu_index=0

def elapsed_time(start_time):
    return time.time()-start_time
class ReplayBuffer:
    def __init__(self,memory_capacity=int(2e6),batch_size=256,num_actions=4,num_states=24):
        self.memory_capacity=memory_capacity
        self.num_states=num_states
        self.num_actions=num_actions
        self.batch_size=batch_size
        self.buffer_counter=0
        self.state_buffer=np.zeros((self.memory_capacity,self.num_states))
        self.action_buffer=np.zeros((self.memory_capacity,self.num_actions))
        self.reward_buffer=np.zeros(self.memory_capacity)
        self.next_state_buffer=np.zeros((self.memory_capacity,self.num_states))
        self.done_buffer=np.zeros(self.memory_capacity)
    def store(self,state,action,reward,next_state,done):
        index=self.buffer_counter%self.memory_capacity
        self.state_buffer[index]=state
        self.action_buffer[index]=action
        self.reward_buffer[index]=reward
        self.next_state_buffer[index]=next_state
        self.done_buffer[index]=done
        self.buffer_counter+=1
    def sample(self):
        max_range=min(self.buffer_counter,self.memory_capacity)
        indices=np.random.randint(0,max_range,size=self.batch_size)
        states=torch.tensor(self.state_buffer[indices],dtype=torch.float32,device=device,requires_grad=False)
        actions=torch.tensor(self.action_buffer[indices],dtype=torch.float32,device=device,requires_grad=False)
        rewards=torch.tensor(self.reward_buffer[indices],dtype=torch.float32,device=device,requires_grad=False)
        next_states=torch.tensor(self.next_state_buffer[indices],dtype=torch.float32,device=device,requires_grad=False)
        dones=torch.tensor(self.done_buffer[indices],dtype=torch.float32,device=device,requires_grad=False)
        return states,actions,rewards,next_states,dones

class Critic(nn.Module):
    def __init__(self,num_states,num_actions,hidden_dim,hidden_dim2):
        super(Critic,self).__init__()
        self.num_actions=num_actions
        self.num_states=num_states
        self.hidden_dim=hidden_dim
        self.hidden_dim2=hidden_dim2
        self.fc1=nn.Linear(self.num_states+self.num_actions,self.hidden_dim)
        self.fc2=nn.Linear(self.hidden_dim,self.hidden_dim2)
        self.fc3=nn.Linear(self.hidden_dim2,1)
    def forward(self,state,action):
        x=torch.cat((state,action),-1)
        x=F.relu(self.fc1(x))
        x=F.relu(self.fc2(x))
        x=self.fc3(x)
        return x

class Actor(nn.Module):
    def __init__(self,num_states,num_actions,action_bound,hidden_dim,hidden_dim2,surrogate_gradient,SAN=False):
        super(Actor,self).__init__()
        self.num_states=num_states
        self.num_actions=num_actions
        self.action_bound=action_bound
        self.hidden_dim=hidden_dim
        self.hidden_dim2=hidden_dim2
        self.surrogate_gradient=surrogate_gradient
        self.SAN=SAN
        self.surrogate_map={
            'Sigmoid':surrogate.Sigmoid,
            'PiecewiseQuadratic':surrogate.PiecewiseQuadratic,
            'PiecewiseExp':surrogate.PiecewiseExp,
            'PiecewiseLeakyReLU':surrogate.PiecewiseLeakyReLU,
            'SoftSign':surrogate.SoftSign,
            'SuperSpike':surrogate.SuperSpike,
            'ATan':surrogate.ATan,
            'Erf':surrogate.Erf,
            'SquarewaveFourierSeries':surrogate.SquarewaveFourierSeries,
            'LogTailedReLU':surrogate.LogTailedReLU,
            'Rect':surrogate.Rect,
            'DeterministicPass':surrogate.DeterministicPass,
        }
        
        if SAN==True:
            surrogate_class=self.surrogate_map.get(self.surrogate_gradient,surrogate.Sigmoid)
            surrogate_func=surrogate_class()

            self.L1=nn.Sequential(
                nn.Linear(self.num_states,self.hidden_dim),
                neuron.LIFNode(surrogate_function=surrogate_func,backend='cupy'),
                nn.Linear(self.hidden_dim,self.hidden_dim2),
                neuron.LIFNode(surrogate_function=surrogate_func,backend='cupy'),
            )
            self.L2=nn.Sequential(
                nn.Linear(self.hidden_dim2,self.num_actions),
                neuron.NonSpikingLIFNode(),
            )
            self.L3=nn.Sequential(
                nn.Linear(self.hidden_dim2,self.num_actions),
                neuron.NonSpikingLIFNode(),
            )
            functional.set_step_mode(self.L1,step_mode='m')
            functional.set_step_mode(self.L2,step_mode='m')
            functional.set_step_mode(self.L3,step_mode='m')
            self.T=16
        else:
            self.fc1=nn.Linear(self.num_states,self.hidden_dim)
            self.fc2=nn.Linear(self.hidden_dim,self.hidden_dim2)
            self.mu=nn.Linear(self.hidden_dim2,self.num_actions)
            self.log_std=nn.Linear(self.hidden_dim2,self.num_actions)
        self.min_log_std=-20
        self.max_log_std=2
        self.sp=nn.Softplus()
    def forward(self,state,SAN=None):
        if SAN is None:
            SAN=self.SAN
        if isinstance(state,np.ndarray):
            state=torch.from_numpy(state).float().to(device)
        else:
            state=state.to(device)

        if SAN==True:
            state=state.unsqueeze(0).repeat(self.T,1,1)
            x=self.L1(state)
            mu=self.L2(x)
            log_std=self.L3(x)
        else:
            x=F.relu(self.fc1(state))
            x=F.relu(self.fc2(x))
            mu=self.mu(x)
            log_std=self.log_std(x)
        log_std=torch.clamp(log_std,self.min_log_std,self.max_log_std)
        return mu,log_std
    def action(self,state,det=False,SAN=None):
        if SAN is None:
            SAN=self.SAN
        if isinstance(state,np.ndarray):
            state=torch.from_numpy(state).float().to(device)
        else:
            state=state.to(device)

        mu,log_std=self(state,SAN)
        if det:
            action=self.action_bound*torch.tanh(mu)
            return action.cpu().detach().numpy()
        std=torch.exp(log_std)
        normal=torch.distributions.Normal(mu,std)
        raw_action=normal.rsample()
        action=self.action_bound*torch.tanh(raw_action)
        log_probs=normal.log_prob(raw_action).sum(axis=-1,keepdim=True)
        transform=2*(np.log(2)-raw_action-self.sp(-2*raw_action)).sum(axis=-1,keepdim=True)
        log_probs-=transform
        if SAN==True:
            functional.reset_net(self)
        return action,log_probs

class Agent:
    def __init__(self,env,hidden_dim,hidden_dim2,seed=None,SAN=False,logger=None):
        self.env=env
        self.SAN=SAN
        self.logger=logger
        if self.SAN==True:
            self.arch="SANSAC"
        else:
            self.arch="SAC"
        if seed is not None:
            self.seed=seed
            self.set_seed(seed)
        self.state_dim=self.env.observation_space.shape[0]
        self.action_dim=self.env.action_space.shape[0]
        self.action_bound=self.env.action_space.high[0]
        self.buffer=ReplayBuffer(num_actions=4,num_states=24)
        self.learning_rate=3e-4
        self.tau=.005
        self.gamma=.975
        self.alpha=.2
        self.hidden_dim=hidden_dim
        self.hidden_dim2=hidden_dim2
        self.surrogate_gradient='Sigmoid'
        self.actor=Actor(self.state_dim,self.action_dim,self.action_bound,hidden_dim,hidden_dim2,self.surrogate_gradient,SAN=self.SAN).to(device)
        self.critic=Critic(self.state_dim,self.action_dim,hidden_dim,hidden_dim2).to(device)
        self.target_critic=deepcopy(self.critic).to(device)
        self.actor_optimizer=optim.Adam(self.actor.parameters(),lr=self.learning_rate)
        self.critic_optimizer=optim.Adam(self.critic.parameters(),lr=self.learning_rate)
        self.critic2=Critic(self.state_dim,self.action_dim,hidden_dim,hidden_dim2).to(device)
        self.target_critic2=deepcopy(self.critic2).to(device)
        self.critic_optimizer2=optim.Adam(self.critic2.parameters(),lr=self.learning_rate)
        self.file_path='samplefile'
        self.totalsteps=0
        self.hip_front=[]
        self.hip_back=[]
        self.front_contact=[]
        self.back_contact=[]
        self.reward_history=[]
        self.best_reward=-float('inf')
        self.best_reward_episode=0
        self.use_max_reward=True
        
    def soft_update(self):
        for target_param,param in zip(self.target_critic.parameters(),self.critic.parameters()):
            target_param.data.copy_(self.tau*param.data+(1-self.tau)*target_param.data)
        for target_param,param in zip(self.target_critic2.parameters(),self.critic2.parameters()):
            target_param.data.copy_(self.tau*param.data+(1-self.tau)*target_param.data)

    def select_action(self,state):
        action,_=self.actor.action(state,SAN=self.SAN)
        action=action.cpu().detach().numpy()
        return action
    def set_seed(self,seed_value):
        random.seed(seed_value)
        np.random.seed(seed_value)
        torch.manual_seed(seed_value)
        torch.cuda.manual_seed(seed_value)
        self.env.action_space.seed(seed_value)
        torch.backends.cudnn.deterministic=True
        torch.backends.cudnn.benchmark=False
        self.seed=seed_value
    def train(self,max_episodes):
        episode_rewards=[]
        for episode in tqdm(range(max_episodes)):
            state,_=self.env.reset(seed=self.seed)
            episode_reward=0
            done=False
            trunc=False
            episode_step_count=0
            while not (done or trunc):
                self.totalsteps+=1
                episode_step_count+=1
                if self.totalsteps<1000:
                    action=self.env.action_space.sample()
                else:
                    action=self.select_action(state)
                    if self.SAN==True:
                        action = action.squeeze(0)
                total_r=0
                for _ in range(1):
                    next_state,reward,done,trunc,info=self.env.step(action)
                    total_r+=reward
                    if done:
                        if self.use_max_reward:
                            total_r=max(total_r,0.0)
                        break
                total_r*=5

                self.buffer.store(state,action,total_r,next_state,done)

                self._train_step()
                episode_reward+=total_r
                state=next_state
                if done or trunc:
                    break
            episode_rewards.append(episode_reward)
            self.reward_history.append(episode_reward)
            if episode_reward>self.best_reward:
                self.best_reward=episode_reward
                self.best_reward_episode=episode
            if self.logger:
                self.logger.log_metrics({
                    'episode_reward':episode_reward,
                    'episode_length':episode_step_count,
                    'total_steps':self.totalsteps,
                    'best_reward so far':self.best_reward,
                },step=self.totalsteps)
                if episode %50==0:
                    self.save_checkpoint(episode)
        return episode_rewards
    def _train_step(self):
        states,actions,rewards,next_states,dones=self.buffer.sample()
        rewards=torch.unsqueeze(rewards,1)
        dones=torch.unsqueeze(dones,1)
        q1=self.critic(states,actions)
        q2=self.critic2(states,actions)
        with torch.no_grad():
            next_action,next_log_probs=self.actor.action(next_states,SAN=self.SAN)
            q1_next_target=self.target_critic(next_states,next_action)
            q2_next_target=self.target_critic2(next_states,next_action)
            q_next_target=torch.min(q1_next_target,q2_next_target)
            value_target=rewards+(1-dones)*self.gamma*(q_next_target-self.alpha*next_log_probs)
        q1_loss=((q1-value_target)**2).mean()
        q2_loss=((q2-value_target)**2).mean()
        self.critic_optimizer.zero_grad()
        self.critic_optimizer2.zero_grad()
        q1_loss.backward()
        q2_loss.backward()
        self.critic_optimizer.step()
        self.critic_optimizer2.step()
        self.actor_optimizer.zero_grad()
        actions_pred,log_pred=self.actor.action(states,SAN=self.SAN)
        q1_pred=self.critic(states,actions_pred)
        q2_pred=self.critic2(states,actions_pred)
        q_pred=torch.min(q1_pred,q2_pred)
        actor_loss=(self.alpha*log_pred-q_pred).mean()
        actor_loss.backward()
        self.actor_optimizer.step()
        self.soft_update()
        if self.logger and self.totalsteps % 100==0:
            self.logger.log_metrics({
                'reward':rewards.mean().item(),
                'critic1_loss':q1_loss.item(),
                'critic2_loss':q2_loss.item(),
                'actor_loss':actor_loss.item(),
                'q_value_mean':q_pred.mean().item(),
                'log_prob_mean':log_pred.mean().item()
            },step=self.totalsteps)
    def test(self):
        state,_=self.env.reset()
        total_reward=0
        start_time=time.time()
        done=False
        trunc=False
        step=0
        while not (done or trunc):
            action=self.actor.action(state,det=True,SAN=self.SAN)
            if self.SAN==True:
                action=action.squeeze(0)
            next_state,reward,done,trunc,_=self.env.step(action)
            state=next_state
            total_reward+=reward
            step+=1
            if done or trunc:
                break
        elapsed_t =elapsed_time(start_time)
        return total_reward,step+1,elapsed_t
    def eval_agent(self,num_tests):
        total_reward=0 
        for test in range(num_tests):
            tot_r, tot_step,tot_time=self.test()
            total_reward+=tot_r
            if self.logger:
                self.logger.log_metrics({
                    "test_reward":tot_r, 
                })
      
        avg_reward=total_reward/num_tests
        if self.logger:
            self.logger.log_metrics({"average_reward":avg_reward,})
        return avg_reward
    def _get_model_dir(self):
        model_dir=os.path.join("experiments",self.file_path,"Models")
        if not os.path.exists(model_dir):
            os.makedirs(model_dir)
        return model_dir
    def save_model(self):
        model_dir=self._get_model_dir()
        actor_path=os.path.join(model_dir,f'{self.seed}_actor.pth')
        critic1_path=os.path.join(model_dir,f'{self.seed}_critic1.pth')
        critic2_path=os.path.join(model_dir,f'{self.seed}_critic2.pth')
        torch.save(self.actor.state_dict(),actor_path,_use_new_zipfile_serialization=True)
        torch.save(self.critic.state_dict(),critic1_path,_use_new_zipfile_serialization=True)
        torch.save(self.critic2.state_dict(),critic2_path,_use_new_zipfile_serialization=True)
        print("Model Saved")
    def load_model(self):
        model_dir=self._get_model_dir()
        actor_path=os.path.join(model_dir,f'{self.seed}_actor.pth')
        critic1_path=os.path.join(model_dir,f'{self.seed}_critic1.pth')
        critic2_path=os.path.join(model_dir,f'{self.seed}_critic2.pth')
        self.actor.load_state_dict(torch.load(actor_path))
        self.critic.load_state_dict(torch.load(critic1_path))
        self.critic2.load_state_dict(torch.load(critic2_path))
        print("Model Loaded")
    def save_checkpoint(self,episode):
        model_dir=os.path.join("experiments",self.file_path,"Checkpoints")
        if not os.path.exists(model_dir):
            os.makedirs(model_dir)
        checkpoint_path=os.path.join(model_dir,f'checkpoint_{episode}.pth')
        torch.save({
            'actor_state_dict':self.actor.state_dict(),
            'critic_state_dict':self.critic.state_dict(),
            'critic2_state_dict':self.critic2.state_dict(),
            'actor_optimizer_state_dict':self.actor_optimizer.state_dict(),
            'critic_optimizer_state_dict':self.critic_optimizer.state_dict(),
            'critic_optimizer2_state_dict':self.critic_optimizer2.state_dict(),
        },checkpoint_path)
        print("Checkpoint saved")
    def load_checkpoint(self):
        checkpoint_dir=os.path.join("experiments",self.file_path,"Checkpoints")
        checkpoint=torch.load(checkpoint_dir)
        self.actor.load_state_dict(checkpoint['actor_state_dict'])
        self.critic.load_state_dict(checkpoint['critic_state_dict'])
        self.critic2.load_state_dict(checkpoint['critic2_state_dict'])
        self.actor_optimizer.load_state_dict(checkpoint['actor_state_dict'])
        self.critic_optimizer.load_state_dict(checkpoint['critic_optimizer_state_dict'])
        self.critic_optimizer2.load_state_dict(checkpoint['critic_optimizer2_state_dict']) 
    def get_surrogate_info(self):
        surrogate_info={}
        if self.SAN:
            for name,module in self.actor.L1.named_modules():
                if hasattr(module,'surrogate_function'):
                    surrogate_info[f'{name}']=type(module.surrogate_function).__name__
        return surrogate_info