[TNNLS] PGDQN: A generalized and efficient preference-guided epsilon-greedy policy equipped DQN for Atari and Autonomous Driving
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Sep 16 19:20:43 2023

@author: oscar
"""

import random
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.nn.utils import clip_grad_value_
from torch.distributions.categorical import Categorical

from Network import Net
from Network import device
from Prioritized_memory import PER

class ReplayMemory(object):

    def __init__(self, capacity, dimension, seed):
        
        self.capacity = capacity
        self.frame_buffer = None
        self.action_buffer = None
        self.reward_buffer = None
        self.memory_counter = 0
        self.seed = seed
        
        torch.manual_seed(self.seed)
        torch.cuda.manual_seed(self.seed)
        random.seed(self.seed)
        np.random.seed(self.seed)
        torch.manual_seed(self.seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    def store_frame(self, frame):
        if len(frame.shape) > 1:
            frame = frame.transpose(2, 0, 1)
        
        if self.frame_buffer is None:
            self.frame_buffer = np.zeros([self.capacity] + list(frame.shape), dtype = np.uint8)
            self.action_buffer = np.zeros([self.capacity], dtype = np.int64)
            self.reward_buffer = np.zeros([self.capacity], dtype = np.float32)

        index = self.memory_counter % self.capacity
        self.frame_buffer[index, :] = frame
        self.memory_counter += 1
        
        return index

    def store_state(self, idx, action, reward):
        self.action_buffer[idx] = action
        self.reward_buffer[idx] = reward

    def push(self, transition):
        """Save a transition"""
        index = self.memory_counter % self.capacity
        self.memory[index,:] = transition
        self.memory_counter += 1

    def sample(self, batch_size):
        up_bound = min(self.memory_counter - 1, self.capacity - 1)
        sample_index = np.random.choice(up_bound, batch_size)
        sample_index_next = sample_index + 1
        sample_index_next[sample_index_next > up_bound] = up_bound
        return self.frame_buffer[sample_index, :], self.action_buffer[sample_index],\
               self.reward_buffer[sample_index], self.frame_buffer[sample_index_next, :]

    def __len__(self):
        return self.memory_counter

class DQN():
    
    def __init__(self,
                 height,
                 width,
                 channel,
                 n_obs,
                 n_actions,
                 DOUBLE,
                 DUELING,
                 PRIORITY,
                 IMPORTANTSAMPLING,
                 PREFERENCE,
                 ENTROPY,
                 JOINT,
                 BATCH_SIZE,
                 GAMMA,
                 EPS_START,
                 EPS_END,
                 EPS_FRAME,
                 MEMORY_CAPACITY,
                 FRAME_HISTORY_LEN,
                 seed,
                 ):
        
        self.height = height
        self.width = width
        self.channel = channel
        self.n_obs = n_obs
        self.n_actions = n_actions
        self.batch_size = BATCH_SIZE
        self.gamma = GAMMA
        self.eps_start = EPS_START
        self.eps_end = EPS_END
        self.eps_frame = EPS_FRAME
        self.memory_capacity = MEMORY_CAPACITY
        self.frame_history_len = FRAME_HISTORY_LEN
        self.double = DOUBLE
        self.dueling = DUELING
        self.per = PRIORITY
        self.pref = PREFERENCE
        self.auto_entropy = ENTROPY
        self.imsamp = IMPORTANTSAMPLING
        self.joint = JOINT
        self.seed = seed
        
        self.lr = 0.00025
        self.lr_p = 0.0001
        self.lr_temp = 0.0001
        self.alpha = 0.95
        self.eps = 0.01
        self.threshold = 50000

        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        
        self.policy_net = Net(height, width, channel, n_actions, DUELING, PREFERENCE, seed)#.to(device)
        self.target_net = Net(height, width, channel, n_actions, DUELING, PREFERENCE, seed)#.to(device)

        if torch.cuda.device_count() > 8:
            print("Let's use", torch.cuda.device_count(), "GPUs!")
            # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
            self.policy_net = nn.DataParallel(self.policy_net)
            self.target_net = nn.DataParallel(self.target_net)
            self.policy_net.to(device)
            self.target_net.to(device)
        else:
            self.policy_net.to(device)
            self.target_net.to(device)

        self.target_net.load_state_dict(self.policy_net.state_dict())
        self.target_net.eval()

        self.optimizer = optim.RMSprop(self.policy_net.parameters(),lr=self.lr,
                                       alpha=self.alpha, eps=self.eps)
        self.optimizer_p = optim.RMSprop(self.policy_net.parameters(),lr=self.
                                         lr_p, alpha=self.alpha, eps=self.eps)

        self.scheduler = lr_scheduler.ExponentialLR(self.optimizer, gamma=1.0)
        self.scheduler_p = lr_scheduler.ExponentialLR(self.optimizer_p, gamma=1.0)
        
        self.steps_done = 0
        self.loss_critic = 0.0
        self.loss_actor = 0.0
        self.loss_entropy = 0.0
        self.eps_threshold = 0.0
        self.action_distribution = 0.0
        self.q_policy = np.zeros(self.n_actions)
        self.q_target = np.zeros(self.n_actions)
        
        self.target_entropy_ratio = 0.98
        self.temperature_copy = 0.0
        
        if self.auto_entropy:
            self.target_entropy = \
                -np.log(1.0 / self.n_actions) * self.target_entropy_ratio
            self.log_temp = torch.zeros(1, requires_grad=True, device=device)
            self.temperature = self.log_temp.exp()
            self.temp_optim = optim.Adam([self.log_temp], lr=self.lr_temp)
        else:
            self.temperature = 0.01 #1.0 / max(1, self.n_actions)
        
        if (PRIORITY):
            self.memory = PER(self.memory_capacity, self.seed)
        else:
            self.memory = ReplayMemory(self.memory_capacity, self.n_obs, self.seed)

    def select_action_deterministic(self, x):
        x = torch.unsqueeze(torch.FloatTensor(x), 0)
        sample = random.random()
        eps_threshold = 0.0
        self.steps_done += 1
        if sample > eps_threshold:
            with torch.no_grad():
                if (self.pref):
                    action_distribution, action = self.policy_net.forward(x)
                    action_idx = action.max(1)[1].view(1, 1)
                    self.action_distribution = action_distribution.cpu().numpy()
                    return action_idx
                
                return self.policy_net.forward(x).max(1)[1].view(1, 1)
        else:
            return torch.tensor([[random.randrange(self.n_actions)]], 
                                device=device, dtype=torch.long) 

    def select_action_ucb(self, action, episode, number_selected_action, gain):
        return np.argmax(action + gain * np.sqrt(np.log(episode + 0.1)/(number_selected_action + 0.1)))

    def select_action(self, x, i_steps):
        sample = random.random()
        
        self.eps_threshold = self.calc_eps_threshold(i_steps)
        with torch.no_grad():
            x = torch.FloatTensor(x._force().transpose(2,0,1)[None]) / 255.0
            if sample > self.eps_threshold:
                if (self.pref):
                    action_distribution, q = self.policy_net.forward(x)
                    action_idx = q.max(1)[1].view(1, 1)
                    self.action_distribution = action_distribution.cpu().numpy()
                    return action_idx
                return self.policy_net.forward(x).max(1)[1].view(1, 1)
            else:
                if (self.pref):
                    action_distribution, q = self.policy_net.forward(x)
                    action_distribution = action_distribution.squeeze(0).cpu().numpy()
                    action_distribution /= action_distribution.sum().tolist()
                    return torch.tensor([[np.random.choice(np.arange(0, self.n_actions),
                                        p=action_distribution)]], device=device, dtype=torch.long)
                return torch.tensor([[random.randrange(self.n_actions)]], 
                                    device=device, dtype=torch.long)

    def calc_eps_threshold(self, i_steps):
        if (i_steps <= self.threshold):
            return self.eps_start
        else:
            fraction = min((i_steps - self.threshold) / self.eps_frame, 1.0)
            return self.eps_start + fraction * (self.eps_end - self.eps_start)

    def append_sample(self, state, action, reward, next_state):
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
        next_state_tensor = torch.FloatTensor(next_state).unsqueeze(0).to(device)
        
        if (self.pref):
            _, q_policy_temp = self.policy_net.forward(state_tensor)
            q_policy_temp = q_policy_temp[action]
            _, next_q_target_temp = self.target_net.forward(next_state_tensor)
        else:
            q_policy_temp = self.policy_net.forward(state_tensor).squeeze(0)[action]
            next_q_target_temp = self.target_net.forward(next_state_tensor)
            
        loss_temp = torch.abs(torch.from_numpy(np.array(reward)).to(device) + 
                              torch.max(next_q_target_temp) * self.gamma - q_policy_temp)
        state_roll, next_state_roll = self.preprocess(state, next_state)
        transition = np.hstack((state_roll, action, reward, next_state_roll))
        self.memory.add(loss_temp.detach().cpu().numpy(), transition)
        
        if (len(self.memory) > 50000):
            print(len(self.memory))
        if (len(self.memory) == self.memory.capacity):
            print('Memory pool is full filled.')

    def preprocess(self, state, next_state):
        state = state.reshape((-1))
        next_state = next_state.reshape((-1))
        return state, next_state

    def optimize_q_branch(self, state, action, reward, next_state):
        if (self.per):
            transitions, idxs, is_weight = self.memory.sample(self.batch_size)
        else:
            state_batch, action_batch, reward_batch, next_state_batch = self.memory.sample(self.batch_size)

        state_batch = torch.FloatTensor(state_batch).to(device)
        action_batch = torch.LongTensor(action_batch).unsqueeze(1).to(device)
        reward_batch = torch.FloatTensor(reward_batch).to(device)
        next_state_batch = torch.FloatTensor(next_state_batch).to(device)

        state_batch = state_batch.view(-1, self.channel, self.height, self.width) /255
        next_state_batch = next_state_batch.view(-1, self.channel, self.height, self.width) /255

        if (not self.pref):
            q_policy = self.policy_net.forward(state_batch).gather(1, action_batch)
        else:
            action_distribution, q_policy = self.policy_net.forward(state_batch)
            q_policy = q_policy.gather(1, action_batch)

        next_q_target = torch.zeros(self.batch_size, device=device)
        if (self.double):
            if (self.pref):
                _, next_q_policy_temp = self.policy_net.forward(next_state_batch)
                _, next_q_target_temp = self.target_net.forward(next_state_batch)
                next_q_target_temp = next_q_target_temp.detach()
            else:
                next_q_policy_temp = self.policy_net.forward(next_state_batch)
                next_q_target_temp = self.target_net.forward(next_state_batch).detach()
                
            max_action_indices = torch.argmax(next_q_policy_temp, dim=1)
            indices_batch = torch.LongTensor(np.arange(self.batch_size))#torch.fromnumpy((np.arange(self.batch_size))
            next_q_target = next_q_target_temp[indices_batch, max_action_indices]
        elif (self.pref):
            _, next_q_target = self.target_net.forward(next_state_batch)
            next_q_target = next_q_target.max(1)[0].detach()
        else:
            next_q_target = self.target_net.forward(next_state_batch).max(1)[0].detach()

        # Compute the expected Q values
        q_target = (next_q_target * self.gamma) + reward_batch

        # Compute Huber loss
        criterion = nn.SmoothL1Loss()

        if (self.per):
            loss_batch = abs(q_policy.squeeze(1)  - q_target).detach().cpu().numpy()
            
            for i in range(self.batch_size):
                idx = idxs[i]
                self.memory.update(idx, loss_batch[i])

            if (self.imsamp):
                is_weight = torch.FloatTensor(is_weight.reshape(-1,1)).to(device)
            else:
                is_weight = 1.0

        else:
            is_weight = 1.0

        #critic loss
        loss_critic = criterion(q_policy * is_weight, q_target.unsqueeze(1) * is_weight)
        if torch.isnan(loss_critic):
            print('q loss is nan.')
        self.loss_critic = loss_critic.detach().cpu().numpy()

        self.optimizer.zero_grad()
        if (self.pref and self.joint):
            loss_policy = self.policy_gradient(state, action, reward, next_state)
            loss_total = loss_critic + loss_policy
            loss_total.backward()
        elif (not self.pref and self.joint):
            print('Mentor is false, but together is true.')
            loss_critic.backward()
        else:
            loss_critic.backward()

        clip_grad_value_(self.policy_net.parameters(), 1)
        self.optimizer.step()

    def policy_gradient(self, state, action, reward, next_state):
        state_tensor = torch.FloatTensor(state._force().transpose(2,0,1)[None]).to(device) / 255.0
        next_state_tensor = torch.FloatTensor(state._force().transpose(2,0,1)[None]).to(device) / 255.0
        action_distribution_policy, q_policy = self.policy_net.forward(state_tensor)
        action_distribution_target, _ = self.target_net.forward(state_tensor)
        _, next_q_target_temp = self.target_net.forward(next_state_tensor)

        q_target = torch.from_numpy(np.array(reward)).to(device) +\
                   next_q_target_temp * self.gamma

        self.ac_dis_policy = action_distribution_policy.squeeze(0).cpu().detach().numpy()
        self.ac_dis_target = action_distribution_target.squeeze(0).cpu().detach().numpy()
        self.q_policy = q_policy.squeeze(0).cpu().detach().numpy()
        self.q_target = q_target.squeeze(0).cpu().detach().numpy()

        action_distribution_policy = action_distribution_policy.squeeze(0)
        action_distribution_target = action_distribution_target.squeeze(0)
        action_prob_policy = Categorical(action_distribution_policy)

        q_policy = q_policy.squeeze(0)
        q_target = q_target.squeeze(0)

        state_value = torch.matmul(action_distribution_target, q_target)
        advantage_function = (q_target - state_value).detach()

        loss_policy = - torch.matmul(action_prob_policy.probs, advantage_function)
        if torch.isnan(loss_policy):
            print('policy loss is nan.')
        loss_entropy =  - action_prob_policy.entropy().mean()
        if torch.isnan(loss_entropy):
            print('entropy loss is nan.')

        self.loss_policy = loss_policy.detach().cpu().numpy()
        self.loss_entropy = loss_entropy.detach().cpu().numpy()

        if (self.temperature > 0):
            loss_policy = loss_policy + loss_entropy * self.temperature
        else:
            loss_policy = loss_policy
        
        if self.auto_entropy:
            self.optimize_entropy_parameter(loss_entropy.detach())
        
        if self.auto_entropy:
            self.temperature_copy = self.temperature.detach().squeeze().cpu().numpy()
        else:
            self.temperature_copy = self.temperature

        return loss_policy
        
    def optimize_preference_branch(self, state, action, reward, next_state):
        if (not self.pref or self.joint):
            return
        loss_policy = self.policy_gradient(state, action, reward, next_state)
        self.optimizer_p.zero_grad()
        loss_policy.backward()
        clip_grad_value_(self.policy_net.parameters(), 1)
        self.optimizer_p.step()
        
    def optimize_entropy_parameter(self, entropy):
        temp_loss = -torch.mean(self.log_temp * (self.target_entropy + entropy))
        self.temp_optim.zero_grad()
        temp_loss.backward()
        self.temp_optim.step()
        self.temperature = self.log_temp.detach().exp()        
        

DRL.py
OscarHuangWind/Preference-Guided-DQN-Atari
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Sep 16 19:21:55 2023

@author: oscar
"""

import torch
import torch.nn as nn
import numpy as np
import random

if torch.cuda.is_available():
    device = torch.device("cuda", 0 if torch.cuda.is_available() else "cpu")
else:
    device = torch.device("cpu")
print('Use:', device)

##############################################################################
class Net(nn.Module):

    def __init__(self, height, width, channel, num_outputs, dueling, preference, seed):
        super(Net, self).__init__()
        self.dueling = dueling
        self.preference = preference
        self.height = height
        self.width = width
        self.feature_dim = 512
        linear_input_size = self.linear_size_input()
        
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        random.seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        
        self.features = nn.Sequential(
            nn.Conv2d(channel, 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU(),
        )
        
        if (self.dueling):
            self.advantage_func = nn.Sequential(
                nn.Linear(linear_input_size, self.feature_dim),
                nn.ReLU(),
                nn.Linear(self.feature_dim, num_outputs)
                )
            
            self.state_value_func = nn.Sequential(
                nn.Linear(linear_input_size, self.feature_dim ),
                nn.ReLU(),
                nn.Linear(self.feature_dim , 1)
                )
        elif (self.preference):
            self.actor_func = nn.Sequential(
                nn.Linear(linear_input_size, self.feature_dim),
                nn.ReLU(),
                nn.Linear(self.feature_dim, num_outputs),
                )
            
            self.q_func = nn.Sequential(
                nn.Linear(linear_input_size, self.feature_dim),
                nn.ReLU(),
                nn.Linear(self.feature_dim, num_outputs),
                )
        else:
            self.fc = nn.Sequential(
                nn.Linear(linear_input_size, self.feature_dim),
                nn.ReLU(),
                nn.Linear(self.feature_dim, num_outputs)
                )

    def linear_size_input(self):    
        convw = self.conv2d_size_3rd(self.conv2d_size_2nd(self.conv2d_size_1st(self.width)))
        convh = self.conv2d_size_3rd(self.conv2d_size_2nd(self.conv2d_size_1st(self.height)))
        return convw * convh * 64

    def conv2d_size_1st(self, size, kernel_size = 8, stride = 4):
        return (size - (kernel_size - 1) - 1) // stride  + 1
        
    def conv2d_size_2nd(self, size, kernel_size = 4, stride = 2):
        return (size - (kernel_size - 1) - 1) // stride  + 1
   
    def conv2d_size_3rd(self, size, kernel_size = 3, stride = 1):
        return (size - (kernel_size - 1) - 1) // stride  + 1

    def forward(self, x):
        x = x.to(device)
        x = self.features(x)
        x = x.contiguous().view(-1, self.linear_size_input())
        if (self.dueling):
            advantage_vec = self.advantage_func(x)
            value_scalar = self.state_value_func(x)
            x = value_scalar + advantage_vec - advantage_vec.mean()
            return x
        elif (self.preference):
            q_value = self.q_func(x)
            action_distribution = self.actor_func(x)
            normalize = nn.Softmax(dim=0)
            action_distribution = normalize(action_distribution.T).T
            return action_distribution, q_value
        else:
            x = self.fc(x)
            return x
        
##############################################################################

Network.py
OscarHuangWind/Preference-Guided-DQN-Atari
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Sep 16 19:25:44 2023

@author: oscar
"""

import random
import numpy as np
from Sumtree import SumTree

class PER:  # stored as ( s, a, r, s_ ) in SumTree
    epsilon = 0.01
    alpha = 0.6
    beta = 0.4
    beta_increment_per_sampling = 0.001


    def __init__(self, capacity, seed):
        self.tree = SumTree(capacity)
        self.capacity = capacity
        self.seed = seed
        random.seed(self.seed)
        np.random.seed(self.seed)

    def __len__(self):
        return self.tree.n_entries

    def _get_priority(self, error):
        return (np.abs(error) + self.epsilon) ** self.alpha

    def add(self, error, sample):
        p = self._get_priority(error)
        self.tree.add(p, sample)

    def sample(self, n):
        batch = np.empty((n, self.tree.data[0].size), dtype=float)
        idxs = []
        segment = self.tree.total() / n
        priorities = []

        self.beta = np.min([1., self.beta + self.beta_increment_per_sampling])

        for i in range(n):
            a = segment * i
            b = segment * (i + 1)

            s = random.uniform(a, b)
            (idx, p, data) = self.tree.get(s)
            priorities.append(max(p, self.epsilon))
            batch[i, :] = data
            idxs.append(idx)

        sampling_probabilities = priorities / self.tree.total()
        is_weight = np.power(self.tree.n_entries * sampling_probabilities, -self.beta)
        is_weight /= is_weight.max()

        return batch, idxs, is_weight

    def update(self, idx, error):
        p = self._get_priority(error)
        self.tree.update(idx, p)
Prioritized_memory.py
OscarHuangWind/Preference-Guided-DQN-Atari
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Sep 16 19:26:18 2023

@author: oscar
"""

import numpy


# SumTree
# a binary tree data structure where the parent’s value is the sum of its children
class SumTree:
    write = 0

    def __init__(self, capacity):
        self.capacity = capacity # amount of leaf node
        self.tree = numpy.zeros(2 * capacity - 1) # amount of whole node
        self.data = numpy.zeros(capacity, dtype=object)
        self.n_entries = 0

    # update to the root node
    def _propagate(self, idx, change):
        parent = (idx - 1) // 2

        self.tree[parent] += change

        if parent != 0:
            self._propagate(parent, change) #recursive

    # find sample on leaf node
    def _retrieve(self, idx, s):
        left = 2 * idx + 1
        right = left + 1

        if left >= len(self.tree):
            return idx

        if s <= self.tree[left]:
            return self._retrieve(left, s)
        else:
            return self._retrieve(right, s - self.tree[left])

    def total(self):
        return self.tree[0]

    # store priority and sample
    def add(self, p, data):
        idx = self.write + self.capacity - 1

        self.data[self.write] = data
        self.update(idx, p)

        self.write += 1
        if self.write >= self.capacity:
            self.write = 0

        if self.n_entries < self.capacity:
            self.n_entries += 1

    # update priority
    def update(self, idx, p):
        change = p - self.tree[idx]

        self.tree[idx] = p
        self._propagate(idx, change)

    # get priority and sample
    def get(self, s):
        idx = self._retrieve(0, s)
        dataIdx = idx - self.capacity + 1

        return (idx, self.tree[idx], self.data[dataIdx])
Sumtree.py
OscarHuangWind/Preference-Guided-DQN-Atari
import numpy as np
import os
os.environ.setdefault('PATH', '')
from collections import deque
import gym
from gym import spaces
import cv2
cv2.ocl.setUseOpenCL(False)

"""
the wrapper is taken from the openai baselines

"""

class NoopResetEnv(gym.Wrapper):
    def __init__(self, env, noop_max=30):
        """Sample initial states by taking random number of no-ops on reset.
        No-op is assumed to be action 0.
        """
        gym.Wrapper.__init__(self, env)
        self.noop_max = noop_max
        self.override_num_noops = None
        self.noop_action = 0
        assert env.unwrapped.get_action_meanings()[0] == 'NOOP'

    def reset(self, **kwargs):
        """ Do no-op action for a number of steps in [1, noop_max]."""
        self.env.reset(**kwargs)
        if self.override_num_noops is not None:
            noops = self.override_num_noops
        else:
            noops = self.unwrapped.np_random.randint(1, self.noop_max + 1) #self.unwrapped.np_random.integers(1, self.noop_max + 1) #pylint: disable=E1101
        assert noops > 0
        obs = None
        for _ in range(noops):
            obs, _, done, _ = self.env.step(self.noop_action)
            if done:
                obs = self.env.reset(**kwargs)
        return obs

    def step(self, ac):
        return self.env.step(ac)

class FireResetEnv(gym.Wrapper):
    def __init__(self, env):
        """Take action on reset for environments that are fixed until firing."""
        gym.Wrapper.__init__(self, env)
        assert env.unwrapped.get_action_meanings()[1] == 'FIRE'
        assert len(env.unwrapped.get_action_meanings()) >= 3

    def reset(self, **kwargs):
        self.env.reset(**kwargs)
        obs, _, done, _ = self.env.step(1)
        if done:
            self.env.reset(**kwargs)
        obs, _, done, _ = self.env.step(2)
        if done:
            self.env.reset(**kwargs)
        return obs

    def step(self, ac):
        return self.env.step(ac)

class EpisodicLifeEnv(gym.Wrapper):
    def __init__(self, env):
        """Make end-of-life == end-of-episode, but only reset on true game over.
        Done by DeepMind for the DQN and co. since it helps value estimation.
        """
        gym.Wrapper.__init__(self, env)
        self.lives = 0
        self.was_real_done  = True

    def step(self, action):
        obs, reward, done, info = self.env.step(action)
        self.was_real_done = done
        # check current lives, make loss of life terminal,
        # then update lives to handle bonus lives
        lives = self.env.unwrapped.ale.lives()
        if lives < self.lives and lives > 0:
            # for Qbert sometimes we stay in lives == 0 condition for a few frames
            # so it's important to keep lives > 0, so that we only reset once
            # the environment advertises done.
            done = True
        self.lives = lives
        return obs, reward, done, info

    def reset(self, **kwargs):
        """Reset only when lives are exhausted.
        This way all states are still reachable even though lives are episodic,
        and the learner need not know about any of this behind-the-scenes.
        """
        if self.was_real_done:
            obs = self.env.reset(**kwargs)
        else:
            # no-op step to advance from terminal/lost life state
            obs, _, _, _ = self.env.step(0)
        self.lives = self.env.unwrapped.ale.lives()
        return obs

class MaxAndSkipEnv(gym.Wrapper):
    def __init__(self, env, skip=4):
        """Return only every `skip`-th frame"""
        gym.Wrapper.__init__(self, env)
        # most recent raw observations (for max pooling across time steps)
        self._obs_buffer = np.zeros((2,)+env.observation_space.shape, dtype=np.uint8)
        self._skip       = skip

    def step(self, action):
        """Repeat action, sum reward, and max over last observations."""
        total_reward = 0.0
        done = None
        for i in range(self._skip):
            obs, reward, done, info = self.env.step(action)
            if i == self._skip - 2: self._obs_buffer[0] = obs
            if i == self._skip - 1: self._obs_buffer[1] = obs
            total_reward += reward
            if done:
                break
        # Note that the observation on the done=True frame
        # doesn't matter
        max_frame = self._obs_buffer.max(axis=0)

        return max_frame, total_reward, done, info

    def reset(self, **kwargs):
        return self.env.reset(**kwargs)

class ClipRewardEnv(gym.RewardWrapper):
    def __init__(self, env):
        gym.RewardWrapper.__init__(self, env)

    def reward(self, reward):
        """Bin reward to {+1, 0, -1} by its sign."""
        return np.sign(reward)


class WarpFrame(gym.ObservationWrapper):
    def __init__(self, env, width=84, height=84, grayscale=True, dict_space_key=None):
        """
        Warp frames to 84x84 as done in the Nature paper and later work.

        If the environment uses dictionary observations, `dict_space_key` can be specified which indicates which
        observation should be warped.
        """
        super().__init__(env)
        self._width = width
        self._height = height
        self._grayscale = grayscale
        self._key = dict_space_key
        if self._grayscale:
            num_colors = 1
        else:
            num_colors = 3

        new_space = gym.spaces.Box(
            low=0,
            high=255,
            shape=(self._height, self._width, num_colors),
            dtype=np.uint8,
        )
        if self._key is None:
            original_space = self.observation_space
            self.observation_space = new_space
        else:
            original_space = self.observation_space.spaces[self._key]
            self.observation_space.spaces[self._key] = new_space
        assert original_space.dtype == np.uint8 and len(original_space.shape) == 3

    def observation(self, obs):
        if self._key is None:
            frame = obs
        else:
            frame = obs[self._key]

        if self._grayscale:
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
        frame = cv2.resize(
            frame, (self._width, self._height), interpolation=cv2.INTER_AREA
        )
        if self._grayscale:
            frame = np.expand_dims(frame, -1)

        if self._key is None:
            obs = frame
        else:
            obs = obs.copy()
            obs[self._key] = frame
        return obs


class FrameStack(gym.Wrapper):
    def __init__(self, env, k):
        """Stack k last frames.

        Returns lazy array, which is much more memory efficient.

        See Also
        --------
        baselines.common.atari_wrappers.LazyFrames
        """
        gym.Wrapper.__init__(self, env)
        self.k = k
        self.frames = deque([], maxlen=k)
        shp = env.observation_space.shape
        self.observation_space = spaces.Box(low=0, high=255, shape=(shp[:-1] + (shp[-1] * k,)), dtype=env.observation_space.dtype)

    def reset(self):
        ob = self.env.reset()
        for _ in range(self.k):
            self.frames.append(ob)
        return self._get_ob()

    def step(self, action):
        ob, reward, done, info = self.env.step(action)
        self.frames.append(ob)
        return self._get_ob(), reward, done, info

    def _get_ob(self):
        assert len(self.frames) == self.k
        return LazyFrames(list(self.frames))

class ScaledFloatFrame(gym.ObservationWrapper):
    def __init__(self, env):
        gym.ObservationWrapper.__init__(self, env)
        self.observation_space = gym.spaces.Box(low=0, high=1, shape=env.observation_space.shape, dtype=np.float32)

    def observation(self, observation):
        # careful! This undoes the memory optimization, use
        # with smaller replay buffers only.
        return np.array(observation).astype(np.float32) / 255.0

class LazyFrames(object):
    def __init__(self, frames):
        """This object ensures that common frames between the observations are only stored once.
        It exists purely to optimize memory usage which can be huge for DQN's 1M frames replay
        buffers.

        This object should only be converted to numpy array before being passed to the model.

        You'd not believe how complex the previous solution was."""
        self._frames = frames
        self._out = None

    def _force(self):
        if self._out is None:
            self._out = np.concatenate(self._frames, axis=-1)
            self._frames = None
        return self._out

    def __array__(self, dtype=None):
        out = self._force()
        if dtype is not None:
            out = out.astype(dtype)
        return out

    def __len__(self):
        return len(self._force())

    def __getitem__(self, i):
        return self._force()[i]

    def count(self):
        frames = self._force()
        return frames.shape[frames.ndim - 1]

    def frame(self, i):
        return self._force()[..., i]

def make_atari(env_id, max_episode_steps=None):
    env = gym.make(env_id)
    assert 'NoFrameskip' in env.spec.id
    env = NoopResetEnv(env, noop_max=30)
    env = MaxAndSkipEnv(env, skip=4)
    if max_episode_steps is not None:
        env = TimeLimit(env, max_episode_steps=max_episode_steps)
    return env

def wrap_deepmind(env, episode_life=True, clip_rewards=True, frame_stack=False, scale=False):
    """Configure environment for DeepMind-style Atari.
    """
    if episode_life:
        env = EpisodicLifeEnv(env)
    if 'FIRE' in env.unwrapped.get_action_meanings():
        env = FireResetEnv(env)
    env = WarpFrame(env)
    if scale:
        env = ScaledFloatFrame(env)
    if clip_rewards:
        env = ClipRewardEnv(env)
    if frame_stack:
        env = FrameStack(env, 4)
    return env

# time limit
class TimeLimit(gym.Wrapper):
    def __init__(self, env, max_episode_steps=None):
        super(TimeLimit, self).__init__(env)
        self._max_episode_steps = max_episode_steps
        self._elapsed_steps = 0

    def step(self, ac):
        observation, reward, done, info = self.env.step(ac)
        self._elapsed_steps += 1
        if self._elapsed_steps >= self._max_episode_steps:
            done = True
            info['TimeLimit.truncated'] = True
        return observation, reward, done, info

    def reset(self, **kwargs):
        self._elapsed_steps = 0
        return self.env.reset(**kwargs)

env_wrapper/atari_wrapper.py
OscarHuangWind/Preference-Guided-DQN-Atari
from rl_utils.env_wrapper.atari_wrapper import make_atari, wrap_deepmind
from rl_utils.env_wrapper.multi_envs_wrapper import SubprocVecEnv
from rl_utils.env_wrapper.frame_stack import VecFrameStack
from rl_utils.logger import logger, bench
import os
import gym

"""
this functions is to create the environments

"""

def create_single_env(args, rank=0):
    # setup the log files
    if rank == 0:
        if not os.path.exists(args.log_dir):
            os.mkdir(args.log_dir)
        log_path = args.log_dir + '/{}/'.format(args.env_name)
        logger.configure(log_path)
    # start to create environment
    if args.env_type == 'atari':
        # create the environment
        env = make_atari(args.env_name)
        # the monitor
        env = bench.Monitor(env, logger.get_dir())
        # use the deepmind environment wrapper
        env = wrap_deepmind(env, clip_rewards=False, frame_stack=True)
    else:
        env = gym.make(args.env_name)
        # add log information
        env = bench.Monitor(env, logger.get_dir(), allow_early_resets=True)
    # set seeds to the environment to make sure the reproducebility
    env.seed(args.seed + rank)
    return env

# create multiple environments - for multiple
def create_multiple_envs(args):
    # now only support the atari games
    if args.env_type == 'atari':
        def make_env(rank):
            def _thunk():
                if not os.path.exists(args.log_dir):
                    os.mkdir(args.log_dir)
                log_path = args.log_dir + '/{}/'.format(args.env_name)
                logger.configure(log_path)
                env = make_atari(args.env_name)
                # set the seed for the environment
                env.seed(args.seed + rank)
                # set loggler
                env = bench.Monitor(env, logger.get_dir() and os.path.join(logger.get_dir(), str(rank)))
                # use the deepmind environment wrapper
                env = wrap_deepmind(env)
                return env
            return _thunk
            # put into sub processing 
        envs = SubprocVecEnv([make_env(i) for i in range(args.num_workers)])
        # then, frame stack
        envs = VecFrameStack(envs, 4)
    else:
        raise NotImplementedError
    return envs


env_wrapper/create_env.py
OscarHuangWind/Preference-Guided-DQN-Atari
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Sep 16 18:10:36 2023

@author: oscar
"""
import os
import yaml
import numpy as np
from tqdm import tqdm 
import matplotlib.pyplot as plt

import torch

from DRL import DQN
from env_wrapper.atari_wrapper import make_atari, wrap_deepmind

def plot_animation_figure():
    plt.figure()
    plt.clf() 
    plt.title(env_name + ' ' + algo + ' ' + str(agent.lr) + '_' + str(agent.lr_p) + '_' + str(agent.lr_temp))
    plt.xlabel('Episode')
    plt.ylabel('Reward')
    plt.plot(train_durations, reward_list)
    plt.plot(train_durations, reward_mean_list)

    plt.pause(0.001) 

def smooth(train_duration_list, reward_list, smooth_horizon):
    mean1 = np.mean(train_duration_list[-min(len(train_duration_list), smooth_horizon):])
    train_durations_mean_list.append(mean1)
    mean3 = np.mean(reward_list[-min(len(reward_list), smooth_horizon):])
    reward_mean_list.append(mean3)

def preprocess(s, a, r, s_):
    state = np.squeeze(s._force().transpose(2,0,1)[None])
    state = state.reshape((-1))
    next_state = np.squeeze(s_._force().transpose(2,0,1)[None])
    next_state = next_state.reshape((-1))
    action = a.cpu().numpy().squeeze().astype(np.int32)
    reward = np.float32(r)
    
    return state, action, reward, next_state

def interaction():
    reward_total = 0.0 
    error = 0.0 
    num_steps = 0
    train_steps = 0
    epoc = 1
    cumulate_flag = False

    s = env.reset()

    if UCB:
        ucb_gain = 2.0
        number_selected_action = np.zeros(n_actions)

    for t in tqdm(range(1, MAX_NUM_STEPS+1), ascii=True):
        
        if VISUALIZATION:
            env.render()

        if num_steps > MAX_NUM_STEPS:
            print('Done')
            break

        # Select and perform an action
        if UCB:
            q = agent.policy_net.forward(torch.FloatTensor(s._force().transpose(2,0,1)[None]) / 255.0).cpu().detach().numpy()
            a = agent.select_action_ucb(q, epoc, number_selected_action, ucb_gain)
            number_selected_action[a] += 1
            a = torch.tensor(a).reshape(1,1)
        else:
            a = agent.select_action(s, num_steps)

        s_, r, done, _ = env.step(a.item())

        #preprocess the state
        state, action, reward, next_state = preprocess(s, a, r, s_)

        # Store the transition in memory
        if (PER):
            agent.append_sample(state, action, reward, next_state) 
        else:
            index = agent.memory.store_frame(state)
            agent.memory.store_state(index, action, reward)

        if (num_steps > THRESHOLD and num_steps % Q_FREQ == 0):
            agent.optimize_q_branch(s, a, reward, s_)
        if (agent.pref and not agent.joint and num_steps % PREFERENCE_FREQ == 0):
            agent.optimize_preference_branch(s, a, reward, s_)

        if (num_steps > THRESHOLD):
            reward_total += reward
            error += agent.loss_critic
            train_steps += 1

            if num_steps % TARGET_UPDATE == 0:
                agent.target_net.load_state_dict(agent.policy_net.state_dict())

            if env.was_real_done:
                
                if (cumulate_flag):
                    reward_list.append(reward_total)
                    
                agent.scheduler.step()
                
                if (agent.pref and not agent.joint):
                    agent.scheduler_p.step()

                if num_steps % PRINT_INTERVAL == 0:
                    print('\n',
                          '\n',
                          'Epoc:', epoc,
                          'Step:', num_steps,
                          'ExpR:', agent.eps_threshold,
                          'R:', reward_total,
                          'Freq:', PREFERENCE_FREQ,
                          'Lr:', agent.scheduler.get_last_lr(),
                          'Lr_p:', agent.scheduler_p.get_last_lr(),
                          'Lr_temp:', agent.lr_temp,
                          'current_temp:', agent.temperature_copy,
                          'Algo:', algo,
                          'seed:', seed,
                          'Env:', env_name,
                          '\n')

                if (cumulate_flag):
                    reward_mean_list.append(np.mean(reward_list[-100:]))
                    train_durations.append(num_steps)
                    
                reward_total = 0
                error = 0
                epoc += 1
                cumulate_flag = True
                
        # Move to the next state and accumulate the reward along trajectory
        s = s_
        num_steps += 1

        if done:
            s = env.reset()

        if (num_steps % SAVE_INTERVAL == 0):
            q_policy_list.append(agent.q_policy)
            q_target_list.append(agent.q_target)
            state_list.append(s._force())

        if (num_steps % SAVE_INTERVAL == 0):
            np.save(os.path.join('data', 'reward_memo'+str(MEMORY_CAPACITY)+
                                 '_step'+str(MAX_NUM_STEPS)+'_seed'+ str(seed) +\
                                 '_' + str(agent.lr)+'_'+env_name+'_' + algo),
                                 [reward_mean_list], allow_pickle=True, fix_imports=True)
            
            np.save(os.path.join('data', 'steps_memo'+str(MEMORY_CAPACITY)+
                                 '_step'+str(MAX_NUM_STEPS)+'_seed'+ str(seed) +\
                                 '_' + str(agent.lr)+'_'+env_name+'_' + algo),
                                 [train_durations], allow_pickle=True, fix_imports=True)        
            
            torch.save(agent.policy_net.state_dict(), os.path.join('network',
                      algo +'_memo' + str(MEMORY_CAPACITY) + '_step' + str(MAX_NUM_STEPS)+
                      '_seed' + str(seed) + '_' + str(agent.lr) + '_' + env_name + '_policynet.pkl'))
            
            plot_animation_figure()


    print('Complete')
    env.close()
    plt.ioff()
    plt.show()

if __name__ == "__main__":

    plt.ion()
    
    # Create folders
    if not os.path.exists("./data"):
        os.makedirs("./data")
        
    if not os.path.exists("./network"):
        os.makedirs("./network")
    
    # Load the config
    path = os.getcwd()
    yaml_path = os.path.join(path, 'config.yaml')
    with open(yaml_path) as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    
    # DQN params
    algo = 'PGDQN' # pls check config.yaml file to check other algos
    algo_param = config[algo]
    DOUBLE = algo_param['DOUBLE']
    DUELING = algo_param['DUELING']
    UCB = algo_param['UCB']
    PER = algo_param['PER']
    IMPORTANT_SAMPLING = algo_param['IMPORTANT_SAMPLING'] 
    PREFERENCE = algo_param['PREFERENCE']
    ENTROPY = algo_param['ENTROPY']
    JOINT_OPT = algo_param['JOINT_OPT']
    
    # Training Params
    THRESHOLD = config['THRESHOLD']
    TARGET_UPDATE = config['TARGET_UPDATE']
    BATCH_SIZE = config['BATCH_SIZE']
    GAMMA = config['GAMMA']
    MEMORY_CAPACITY = config['MEMORY_CAPACITY']
    Q_FREQ = config['Q_FREQ']
    PREFERENCE_FREQ = config['PREFERENCE_FREQ']
    EPS_START = config['EPS_START']
    EPS_END = config['EPS_END']
    EPS_DECAY = config['EPS_DECAY']
    FRAME_HISTORY_LEN = config['FRAME_HISTORY_LEN']
    MAX_NUM_STEPS = config['MAX_NUM_STEPS']
    PRINT_INTERVAL = config['PRINT_INTERVAL']
    SAVE_INTERVAL = config['SAVE_INTERVAL']
    VISUALIZATION = config['VISUALIZATION']
    
    # Environment
    env = make_atari('PongNoFrameskip-v4')
    env = wrap_deepmind(env, clip_rewards=False, frame_stack=True)
    env_name = 'PongNoFrameskip-v4'
    
    # State and Action Space
    obs = env.reset()
    img_h, img_w, channel = env.observation_space.shape
    n_obs = img_h * img_w * channel
    n_actions = env.action_space.n

    train_durations_list = []
    step_list_list = []
    reward_list_list = []
    legend_bar = []
    
    for seed in range(525,526):
        env.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        
        legend_bar.append('seed'+str(seed))

        agent = DQN(img_h, img_w, channel, n_obs, n_actions,
                    DOUBLE, DUELING, PER, IMPORTANT_SAMPLING, PREFERENCE,
                    ENTROPY, JOINT_OPT, BATCH_SIZE, GAMMA, EPS_START, EPS_END,
                    EPS_DECAY, MEMORY_CAPACITY, FRAME_HISTORY_LEN, seed)
        
        train_durations = []
        train_durations_mean_list = []
        reward_list = []
        reward_mean_list = []
        q_policy_list = []
        q_target_list = []
        state_list = []
        
        print('The object is:', algo, '|Preference:', agent.pref,
              '|Auto Entropy:', agent.auto_entropy, '|Priority:', agent.per,
              '|Double:', agent.double, '|Dueling:', agent.dueling, '|UCB', UCB,
              '|Joint Optimization:', agent.joint, '|Seed:', seed)

        interaction()
        train_durations_list.append(train_durations)
        reward_list_list.append(reward_mean_list)

        np.save(os.path.join('data', 'reward_memo'+str(MEMORY_CAPACITY)+
                             '_step'+str(MAX_NUM_STEPS)+'_seed'+ str(seed) +\
                             '_' + str(agent.lr)+'_'+env_name+'_' + algo),
                             [reward_mean_list], allow_pickle=True, fix_imports=True)
        
        np.save(os.path.join('data', 'steps_memo'+str(MEMORY_CAPACITY)+
                             '_step'+str(MAX_NUM_STEPS)+'_seed'+ str(seed) +\
                             '_' + str(agent.lr)+'_'+env_name+'_' + algo),
                             [train_durations], allow_pickle=True, fix_imports=True)        
        
        torch.save(agent.policy_net.state_dict(), os.path.join('network',
                  algo +'_memo' + str(MEMORY_CAPACITY) + '_step' + str(MAX_NUM_STEPS)+
                  '_seed' + str(seed) + '_' + str(agent.lr) + '_' + env_name + '_policynet.pkl'))

main.py
OscarHuangWind/Preference-Guided-DQN-Atari