141 lines
5.8 KiB
Python
141 lines
5.8 KiB
Python
import gym
|
|
from gym import spaces
|
|
from sklearn.neighbors import KDTree
|
|
from scipy.stats import norm
|
|
import numpy as np
|
|
import warnings
|
|
from abc import ABC, abstractmethod
|
|
from typing import Dict, Generator, Optional, Union
|
|
import torch as th
|
|
|
|
try:
|
|
# Check memory used by replay buffer when possible
|
|
import psutil
|
|
except ImportError:
|
|
psutil = None
|
|
|
|
from stable_baselines3.common.preprocessing import get_action_dim, get_obs_shape
|
|
from stable_baselines3.common.type_aliases import ReplayBufferSamples, RolloutBufferSamples
|
|
from stable_baselines3.common.vec_env import VecNormalize
|
|
from stable_baselines3.common.buffers import ReplayBuffer
|
|
|
|
|
|
class AdRILWrapper(gym.Env):
|
|
metadata = {'render.modes': ['human']}
|
|
|
|
def __init__(self, base_env):
|
|
super(AdRILWrapper, self).__init__()
|
|
self.base_env = base_env
|
|
self.iter = 0
|
|
self.observation_space = self.base_env.observation_space
|
|
self.action_space = self.base_env.action_space
|
|
self.trajs = list()
|
|
self.num_trajs = 0
|
|
self.curr_state = None
|
|
def step(self, action):
|
|
next_obs, _, done, info = self.base_env.step(action)
|
|
reward = self.iter # Transformed by replay buffer
|
|
self.trajs.append((self.curr_state, action, next_obs, done))
|
|
if done:
|
|
self.num_trajs += 1
|
|
self.curr_state = next_obs
|
|
return next_obs, reward, done, info
|
|
def reset(self):
|
|
obs = self.base_env.reset()
|
|
self.curr_state = obs
|
|
return obs
|
|
def render(self, mode='human'):
|
|
self.base_env.render(mode=mode)
|
|
def close (self):
|
|
self.base_env.close()
|
|
def get_learner_trajs(self):
|
|
return self.trajs
|
|
def set_iter(self, k):
|
|
self.iter = k
|
|
|
|
class AdRILReplayBuffer(ReplayBuffer):
|
|
def __init__(
|
|
self,
|
|
buffer_size: int,
|
|
observation_space: spaces.Space,
|
|
action_space: spaces.Space,
|
|
device: Union[th.device, str] = "cpu",
|
|
n_envs: int = 1,
|
|
optimize_memory_usage: bool = False,
|
|
expert_data: dict = dict(),
|
|
N_expert: int = 0,
|
|
balanced: bool = True,
|
|
):
|
|
super(AdRILReplayBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs, optimize_memory_usage=optimize_memory_usage)
|
|
|
|
self.expert_states = expert_data['obs']
|
|
self.expert_actions = expert_data['acts']
|
|
self.expert_next_states = expert_data['next_obs']
|
|
self.expert_dones = expert_data['dones']
|
|
n_expert = len(expert_data["obs"])
|
|
self.iter = 0
|
|
self.N_expert = N_expert
|
|
self.N_learner = 0
|
|
self.normalizer = 1
|
|
self.balanced = balanced
|
|
|
|
def set_iter(self, k):
|
|
self.iter = k
|
|
normalizer = 0
|
|
for i in range(0, k):
|
|
normalizer += 1 ** (-i) # written to support decaying learning rate
|
|
self.normalizer = normalizer
|
|
|
|
def set_n_learner(self, n):
|
|
self.N_learner = n
|
|
|
|
def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> ReplayBufferSamples:
|
|
num_samples = len(batch_inds)
|
|
if self.balanced:
|
|
num_expert_samples = int(num_samples / 2)
|
|
batch_inds = batch_inds[:num_expert_samples]
|
|
expert_inds = np.random.randint(0, len(self.expert_states), size=num_expert_samples)
|
|
# balanced sampling
|
|
if self.optimize_memory_usage:
|
|
next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, 0, :], env)
|
|
else:
|
|
next_obs = self._normalize_obs(self.next_observations[batch_inds, 0, :], env)
|
|
next_obs = np.concatenate((next_obs, self._normalize_obs(self.expert_next_states[expert_inds], env)), axis=0)
|
|
obs = self._normalize_obs(self.observations[batch_inds, 0, :], env)
|
|
obs = np.concatenate((obs, self._normalize_obs(self.expert_states[expert_inds], env)), axis=0)
|
|
actions = self.actions[batch_inds, 0, :]
|
|
actions = np.concatenate((actions, self.expert_actions[expert_inds].reshape(num_expert_samples, -1)), axis=0)
|
|
dones = self.dones[batch_inds]
|
|
dones = np.concatenate((dones, self.expert_dones[expert_inds].reshape(num_expert_samples, -1)), axis=0)
|
|
# AdRIL Rewards (indicator kernel)
|
|
mask1 = (self.rewards[batch_inds] >= 0).astype(np.float32)
|
|
mask2 = (self.rewards[batch_inds] < self.iter).astype(np.float32)
|
|
r1 = - (1. ** (-self.rewards[batch_inds])) * mask1 * mask2 # Past iter
|
|
r2 = np.zeros_like(self.rewards[batch_inds]) * mask1 * (1 - mask2) # current iter
|
|
r3 = -self.rewards[batch_inds] * (1 - mask1) # Expert
|
|
if self.iter > 0:
|
|
rewards = (r1 / self.N_learner) + r2 + r3
|
|
else:
|
|
rewards = r1 + r2 + r3
|
|
rewards = np.concatenate((rewards, np.ones_like(rewards) / self.N_expert), axis=0)
|
|
else:
|
|
if self.optimize_memory_usage:
|
|
next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, 0, :], env)
|
|
else:
|
|
next_obs = self._normalize_obs(self.next_observations[batch_inds, 0, :], env)
|
|
obs = self._normalize_obs(self.observations[batch_inds, 0, :], env)
|
|
actions = self.actions[batch_inds, 0, :]
|
|
dones = self.dones[batch_inds]
|
|
# AdRIL Rewards (indicator kernel)
|
|
mask1 = (self.rewards[batch_inds] >= 0).astype(np.float32)
|
|
mask2 = (self.rewards[batch_inds] < self.iter).astype(np.float32)
|
|
r1 = - (1. ** (-self.rewards[batch_inds])) * mask1 * mask2 # Past iter
|
|
r2 = np.zeros_like(self.rewards[batch_inds]) * mask1 * (1 - mask2) # current iter
|
|
r3 = -self.rewards[batch_inds] * (1 - mask1) / self.N_expert # Expert
|
|
if self.iter > 0:
|
|
rewards = (r1 * 1. / self.N_learner) + r2 + r3
|
|
else:
|
|
rewards = r1 + r2 + r3
|
|
data = (obs, actions, next_obs, dones, rewards)
|
|
return ReplayBufferSamples(*tuple(map(self.to_torch, data)))
|