AdVIL tests

scratch/etienne/pillbox/learners/adril.py

@@ -0,0 +1,140 @@
+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)))