import torch import numpy as np import torch.nn as nn import gym import os import random from torch.utils.data import Dataset, DataLoader import time from skimage.util.shape import view_as_windows from diy_gym.utils import flatten, unflatten class eval_mode(object): def __init__(self, *models): self.models = models def __enter__(self): self.prev_states = [] for model in self.models: self.prev_states.append(model.training) model.train(False) def __exit__(self, *args): for model, state in zip(self.models, self.prev_states): model.train(state) return False def soft_update_params(net, target_net, tau): for param, target_param in zip(net.parameters(), target_net.parameters()): target_param.data.copy_( tau * param.data + (1 - tau) * target_param.data ) def set_seed_everywhere(seed): torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) np.random.seed(seed) random.seed(seed) def module_hash(module): result = 0 for tensor in module.state_dict().values(): result += tensor.sum().item() return result def make_dir(dir_path): os.makedirs(dir_path, exist_ok=True) return dir_path def preprocess_obs(obs, bits=5): """Preprocessing image, see https://arxiv.org/abs/1807.03039.""" bins = 2**bits assert obs.dtype == torch.float32 if bits < 8: obs = torch.floor(obs / 2**(8 - bits)) obs = obs / bins obs = obs + torch.rand_like(obs) / bins obs = obs - 0.5 return obs class ReplayBuffer(Dataset): """Buffer to store environment transitions.""" def __init__(self, obs_space, action_space, capacity, batch_size, device, image_size=84, transform=None): obs_shape = flatten(obs_space.sample()).shape action_shape = action_space.shape self.obs_space = obs_space self.capacity = capacity self.batch_size = batch_size self.device = device self.image_size = image_size self.transform = transform # the proprioceptive obs is stored as float32, mixeds obs as uint8 obs_dtype = np.float16 if len(obs_shape) == 1 else np.uint8 self.obses = np.empty((capacity, *obs_shape), dtype=obs_dtype) self.next_obses = np.empty((capacity, *obs_shape), dtype=obs_dtype) self.actions = np.empty((capacity, *action_shape), dtype=np.float32) self.rewards = np.empty((capacity, 1), dtype=np.float32) self.not_dones = np.empty((capacity, 1), dtype=np.float32) self.idx = 0 self.last_save = 0 self.full = False def add(self, obs, action, reward, next_obs, done): obs = flatten(obs) next_obs = flatten(next_obs) np.copyto(self.obses[self.idx], obs) np.copyto(self.actions[self.idx], action) np.copyto(self.rewards[self.idx], reward) np.copyto(self.next_obses[self.idx], next_obs) np.copyto(self.not_dones[self.idx], not done) self.idx = (self.idx + 1) % self.capacity self.full = self.full or self.idx == 0 def unflatten_obs(self, obs): obs = [unflatten(o, self.obs_space) for o in obs] obs = { k: np.stack([o[k] for o in obs]) for k in obs[0].keys() } return obs def sample_proprio(self): idxs = np.random.randint( 0, self.capacity if self.full else self.idx, size=self.batch_size ) obses = self.as_tensor_obs(self.unflatten_obs(self.obses[idxs])) next_obses = self.as_tensor_obs(self.unflatten_obs(self.next_obses[idxs])) actions = torch.as_tensor(self.actions[idxs], device=self.device) rewards = torch.as_tensor(self.rewards[idxs], device=self.device) not_dones = torch.as_tensor(self.not_dones[idxs], device=self.device) return obses, actions, rewards, next_obses, not_dones def sample_cpc(self): start = time.time() idxs = np.random.randint( 0, self.capacity if self.full else self.idx, size=self.batch_size ) obses_raw = self.unflatten_obs(self.obses[idxs]) next_obses_raw = self.unflatten_obs(self.next_obses[idxs]) # Split mixed obs into image and state state, obses = split_obs(obses_raw) next_state, next_obses = split_obs(next_obses_raw) pos = obses.copy() # Crop obses = random_crop(obses, self.image_size) next_obses = random_crop(next_obses, self.image_size) pos = random_crop(pos, self.image_size) # Recombine obses = self.as_tensor_obs(combine_obs(state, obses)) next_obses = self.as_tensor_obs(combine_obs(next_state, next_obses)) pos = self.as_tensor_obs(combine_obs(state, pos)) actions = torch.as_tensor(self.actions[idxs], device=self.device) rewards = torch.as_tensor(self.rewards[idxs], device=self.device) not_dones = torch.as_tensor(self.not_dones[idxs], device=self.device) cpc_kwargs = dict(obs_anchor=obses, obs_pos=pos, time_anchor=None, time_pos=None) return obses, actions, rewards, next_obses, not_dones, cpc_kwargs def as_tensor_obs(self, obses): obses['img'] = torch.as_tensor(obses['img'], device=self.device).float() obses['state'] = torch.as_tensor(obses['state'], device=self.device).float() return obses def save(self, save_dir): if self.idx == self.last_save: return path = os.path.join(save_dir, '%d_%d.pt' % (self.last_save, self.idx)) payload = [ self.obses[self.last_save:self.idx], self.next_obses[self.last_save:self.idx], self.actions[self.last_save:self.idx], self.rewards[self.last_save:self.idx], self.not_dones[self.last_save:self.idx] ] self.last_save = self.idx torch.save(payload, path) def load(self, save_dir): chunks = os.listdir(save_dir) chucks = sorted(chunks, key=lambda x: int(x.split('_')[0])) for chunk in chucks: start, end = [int(x) for x in chunk.split('.')[0].split('_')] path = os.path.join(save_dir, chunk) payload = torch.load(path) assert self.idx == start self.obses[start:end] = payload[0] self.next_obses[start:end] = payload[1] self.actions[start:end] = payload[2] self.rewards[start:end] = payload[3] self.not_dones[start:end] = payload[4] self.idx = end def __getitem__(self, idx): idx = np.random.randint( 0, self.capacity if self.full else self.idx, size=1 ) idx = idx[0] obs = self.unflatten_obs(self.obses[idx]) action = self.actions[idx] reward = self.rewards[idx] next_obs = self.unflatten_obs(self.next_obses[idx]) not_done = self.not_dones[idx] if self.transform: obs = self.transform(obs) next_obs = self.transform(next_obs) return obs, action, reward, next_obs, not_done def __len__(self): return self.capacity def random_crop(imgs, output_size): """ Vectorized way to do random crop using sliding windows and picking out random ones args: imgs, batch images with shape (B,C,H,W) """ # batch size n = imgs.shape[0] img_size = imgs.shape[-1] crop_max = img_size - output_size imgs = np.transpose(imgs, (0, 2, 3, 1)) w1 = np.random.randint(0, crop_max, n) h1 = np.random.randint(0, crop_max, n) # creates all sliding windows combinations of size (output_size) windows = view_as_windows( imgs, (1, output_size, output_size, 1))[..., 0,:,:, 0] # selects a random window for each batch element cropped_imgs = windows[np.arange(n), w1, h1] return cropped_imgs def center_crop_image(image, output_size): h, w = image.shape[1:] new_h, new_w = output_size, output_size top = (h - new_h)//2 left = (w - new_w)//2 image = image[:, top:top + new_h, left:left + new_w] return image def split_obs(obs): """Split a dict obs into state and images.""" return obs['state'], obs['img'] def combine_obs(state, img): return dict(state=state, img=img) def split_obs_shape(obs_shape): obs = np.zeros(obs_shape) state, img = split_obs(obs) return state.shape, img.shape