From b1eb2efb770accd587c8fd6f55dcd47b1e418321 Mon Sep 17 00:00:00 2001 From: wassname Date: Thu, 18 Feb 2021 11:57:29 +0800 Subject: [PATCH] working with obs=dict(state=state, img=img) --- README.md | 31 +++++++ conda_env.yml | 2 +- curl_sac.py | 26 +++--- encoder.py | 16 +++- logger.py | 1 + .../args.json | 44 ++++++++++ ...ut.tfevents.1613620610.mjcdesktop.457460.0 | Bin 0 -> 143 bytes scripts/run.sh | 14 ++- train.py | 74 ++++++++-------- utils.py | 80 +++++++++--------- 10 files changed, 192 insertions(+), 96 deletions(-) create mode 100644 runs/ApplePick-v0-02-18-im84-b32-s1-mixed/args.json create mode 100644 runs/ApplePick-v0-02-18-im84-b32-s1-mixed/tb/events.out.tfevents.1613620610.mjcdesktop.457460.0 diff --git a/README.md b/README.md index fd0cf5d..b4ee31a 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,33 @@ +Adapt to Apple gym + + +TODO: +- apple + - make env that returns seperate obs dict + - make a wrapper that flattens act, and normalises act + - make a wrapper that flattens obs... + - mkae a wrapper that does grconv + - norm + - [x] space should be acc, but it's not fix in diy gym + - [ ] a wrapper that will normaction + - [ ] a wrapper to norm dict obs + - then norm based on space? +- this + - make env handle it + - make replay buffer handle it + +- we take in flat array which contains [...state, ...img_flat] +- we reshape in encoder +- no resize img +- always use pixels of course, otherwise no curl +- [ ] replay buffer store flat... even tho pixel mode +- [ ] custom encoder? + + +`./scripts/run.sh` + +------------------------- + # CURL: Contrastive Unsupervised Representation Learning for Sample-Efficient Reinforcement Learning This repository is the official implementation of [CURL](https://mishalaskin.github.io/curl/) for the DeepMind control experiments. Atari experiments were done in a separate codebase available [here](https://github.com/aravindsrinivas/curl_rainbow). Our implementation of SAC is based on [SAC+AE](https://github.com/denisyarats/pytorch_sac_ae) by Denis Yarats. @@ -28,6 +58,7 @@ In your console, you should see printouts that look like: ``` | train | E: 221 | S: 28000 | D: 18.1 s | R: 785.2634 | BR: 3.8815 | A_LOSS: -305.7328 | CR_LOSS: 190.9854 | CU_LOSS: 0.0000 + | train | E: 225 | S: 28500 | D: 18.6 s | R: 832.4937 | BR: 3.9644 | A_LOSS: -308.7789 | CR_LOSS: 126.0638 | CU_LOSS: 0.0000 | train | E: 229 | S: 29000 | D: 18.8 s | R: 683.6702 | BR: 3.7384 | A_LOSS: -311.3941 | CR_LOSS: 140.2573 | CU_LOSS: 0.0000 | train | E: 233 | S: 29500 | D: 19.6 s | R: 838.0947 | BR: 3.7254 | A_LOSS: -316.9415 | CR_LOSS: 136.5304 | CU_LOSS: 0.0000 diff --git a/conda_env.yml b/conda_env.yml index fffc1a5..df9cda9 100644 --- a/conda_env.yml +++ b/conda_env.yml @@ -17,4 +17,4 @@ dependencies: - imageio - imageio-ffmpeg - torchvision - - scikit-image \ No newline at end of file + - scikit-image diff --git a/curl_sac.py b/curl_sac.py index 5836371..002b69b 100644 --- a/curl_sac.py +++ b/curl_sac.py @@ -248,7 +248,7 @@ class CurlSacAgent(object): critic_beta=0.9, critic_tau=0.005, critic_target_update_freq=2, - encoder_type='pixel', + encoder_type='mixed', encoder_feature_dim=50, encoder_lr=1e-3, encoder_tau=0.005, @@ -267,7 +267,7 @@ class CurlSacAgent(object): self.critic_target_update_freq = critic_target_update_freq self.cpc_update_freq = cpc_update_freq self.log_interval = log_interval - self.image_size = obs_shape[-1] + self.image_size = obs_shape['img'][-1] self.curl_latent_dim = curl_latent_dim self.detach_encoder = detach_encoder self.encoder_type = encoder_type @@ -311,7 +311,7 @@ class CurlSacAgent(object): [self.log_alpha], lr=alpha_lr, betas=(alpha_beta, 0.999) ) - if self.encoder_type == 'pixel': + if self.encoder_type == 'mixed': # create CURL encoder (the 128 batch size is probably unnecessary) self.CURL = CURL(obs_shape, encoder_feature_dim, self.curl_latent_dim, self.critic,self.critic_target, output_type='continuous').to(self.device) @@ -333,7 +333,7 @@ class CurlSacAgent(object): self.training = training self.actor.train(training) self.critic.train(training) - if self.encoder_type == 'pixel': + if self.encoder_type == 'mixed': self.CURL.train(training) @property @@ -342,20 +342,22 @@ class CurlSacAgent(object): def select_action(self, obs): with torch.no_grad(): - obs = torch.FloatTensor(obs).to(self.device) - obs = obs.unsqueeze(0) + obs['img'] = torch.FloatTensor(obs['img']).to(self.device).unsqueeze(0) + obs['state'] = torch.FloatTensor(obs['state']).to(self.device).unsqueeze(0) mu, _, _, _ = self.actor( obs, compute_pi=False, compute_log_pi=False ) return mu.cpu().data.numpy().flatten() def sample_action(self, obs): - if obs.shape[-1] != self.image_size: - obs = utils.center_crop_image(obs, self.image_size) + if obs['img'].shape[-1] != self.image_size: + state, img = utils.split_obs(obs) + img = utils.center_crop_image(img, self.image_size) + obs = utils.combine_obs(state, img) with torch.no_grad(): - obs = torch.FloatTensor(obs).to(self.device) - obs = obs.unsqueeze(0) + obs['img'] = torch.FloatTensor(obs['img']).to(self.device).unsqueeze(0) + obs['state'] = torch.FloatTensor(obs['state']).to(self.device).unsqueeze(0) mu, pi, _, _ = self.actor(obs, compute_log_pi=False) return pi.cpu().data.numpy().flatten() @@ -435,7 +437,7 @@ class CurlSacAgent(object): def update(self, replay_buffer, L, step): - if self.encoder_type == 'pixel': + if self.encoder_type == 'mixed': obs, action, reward, next_obs, not_done, cpc_kwargs = replay_buffer.sample_cpc() else: obs, action, reward, next_obs, not_done = replay_buffer.sample_proprio() @@ -460,7 +462,7 @@ class CurlSacAgent(object): self.encoder_tau ) - if step % self.cpc_update_freq == 0 and self.encoder_type == 'pixel': + if step % self.cpc_update_freq == 0 and self.encoder_type == 'mixed': obs_anchor, obs_pos = cpc_kwargs["obs_anchor"], cpc_kwargs["obs_pos"] self.update_cpc(obs_anchor, obs_pos,cpc_kwargs, L, step) diff --git a/encoder.py b/encoder.py index 9da499f..db56c6b 100644 --- a/encoder.py +++ b/encoder.py @@ -15,7 +15,7 @@ OUT_DIM_64 = {2: 29, 4: 25, 6: 21} class PixelEncoder(nn.Module): - """Convolutional encoder of pixels observations.""" + """Convolutional encoder of mixeds observations.""" def __init__(self, obs_shape, feature_dim, num_layers=2, num_filters=32,output_logits=False): super().__init__() @@ -114,7 +114,19 @@ class IdentityEncoder(nn.Module): pass -_AVAILABLE_ENCODERS = {'pixel': PixelEncoder, 'identity': IdentityEncoder} + +class MixedEncoder(PixelEncoder): + def __init__(self, obs_shape, feature_dim, num_layers=2, num_filters=32, output_logits=False): + img_shape = obs_shape['img'] + super().__init__(img_shape, feature_dim, num_layers, num_filters, output_logits) + self.feature_dim = feature_dim + obs_shape['state'][0] + def forward(self, obs, detach=False): + h = super().forward(obs['img'], detach) + return torch.cat([obs['state'], h], 1) + + + +_AVAILABLE_ENCODERS = {'pixel': PixelEncoder, 'identity': IdentityEncoder, 'mixed': MixedEncoder} def make_encoder( diff --git a/logger.py b/logger.py index e35d6a0..18fc30b 100644 --- a/logger.py +++ b/logger.py @@ -7,6 +7,7 @@ import torch import torchvision import numpy as np from termcolor import colored +from rich import print FORMAT_CONFIG = { 'rl': { diff --git a/runs/ApplePick-v0-02-18-im84-b32-s1-mixed/args.json b/runs/ApplePick-v0-02-18-im84-b32-s1-mixed/args.json new file mode 100644 index 0000000..20023f5 --- /dev/null +++ b/runs/ApplePick-v0-02-18-im84-b32-s1-mixed/args.json @@ -0,0 +1,44 @@ +{ + "action_repeat": 1, + "actor_beta": 0.9, + "actor_log_std_max": 2, + "actor_log_std_min": -10, + "actor_lr": 0.001, + "actor_update_freq": 2, + "agent": "curl_sac", + "alpha_beta": 0.5, + "alpha_lr": 0.0001, + "batch_size": 32, + "critic_beta": 0.9, + "critic_lr": 0.001, + "critic_target_update_freq": 2, + "critic_tau": 0.01, + "curl_latent_dim": 128, + "detach_encoder": false, + "discount": 0.99, + "domain_name": "ApplePick-v0", + "encoder_feature_dim": 50, + "encoder_lr": 0.001, + "encoder_tau": 0.05, + "encoder_type": "mixed", + "eval_freq": 1000, + "frame_stack": 3, + "hidden_dim": 1024, + "image_size": 84, + "init_steps": 1000, + "init_temperature": 0.1, + "log_interval": 100, + "num_eval_episodes": 10, + "num_filters": 32, + "num_layers": 4, + "num_train_steps": 1000000, + "pre_transform_image_size": 124, + "render": false, + "replay_buffer_capacity": 10000, + "save_buffer": false, + "save_model": false, + "save_tb": true, + "save_video": false, + "seed": 1, + "work_dir": "./runs/ApplePick-v0-02-18-im84-b32-s1-mixed" +} \ No newline at end of file diff --git a/runs/ApplePick-v0-02-18-im84-b32-s1-mixed/tb/events.out.tfevents.1613620610.mjcdesktop.457460.0 b/runs/ApplePick-v0-02-18-im84-b32-s1-mixed/tb/events.out.tfevents.1613620610.mjcdesktop.457460.0 new file mode 100644 index 0000000000000000000000000000000000000000..223248349026f86efb8c324fd20bfffebe40b857 GIT binary patch literal 143 zcmb1OfPlsI-b$Q|#~KnUxor7Pik3Wj(%!EW^sN>swha=)w;eX>QME&A8vHx{4Dx49b~SS43`v_FtYmiqSW%l Qq7>1v=}%l%Ot= args.init_steps: diff --git a/utils.py b/utils.py index 6603caf..31ecb42 100644 --- a/utils.py +++ b/utils.py @@ -3,11 +3,11 @@ import numpy as np import torch.nn as nn import gym import os -from collections import deque 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): @@ -48,10 +48,7 @@ def module_hash(module): def make_dir(dir_path): - try: - os.mkdir(dir_path) - except OSError: - pass + os.makedirs(dir_path, exist_ok=True) return dir_path @@ -69,13 +66,16 @@ def preprocess_obs(obs, bits=5): class ReplayBuffer(Dataset): """Buffer to store environment transitions.""" - def __init__(self, obs_shape, action_shape, capacity, batch_size, device,image_size=84,transform=None): + 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, pixels obs as uint8 + # the proprioceptive obs is stored as float32, mixeds obs as uint8 obs_dtype = np.float32 if len(obs_shape) == 1 else np.uint8 self.obses = np.empty((capacity, *obs_shape), dtype=obs_dtype) @@ -92,7 +92,7 @@ class ReplayBuffer(Dataset): def add(self, obs, action, reward, next_obs, done): - + obs = flatten(obs) np.copyto(self.obses[self.idx], obs) np.copyto(self.actions[self.idx], action) np.copyto(self.rewards[self.idx], reward) @@ -111,6 +111,9 @@ class ReplayBuffer(Dataset): obses = self.obses[idxs] next_obses = self.next_obses[idxs] + obses = unflatten(obses, self.obs_space) + next_obses = unflatten(next_obses, self.obs_space) + obses = torch.as_tensor(obses, device=self.device).float() actions = torch.as_tensor(self.actions[idxs], device=self.device) rewards = torch.as_tensor(self.rewards[idxs], device=self.device) @@ -127,13 +130,27 @@ class ReplayBuffer(Dataset): 0, self.capacity if self.full else self.idx, size=self.batch_size ) - obses = self.obses[idxs] - next_obses = self.next_obses[idxs] + obses_raw = self.obses[idxs] + next_obses_raw = self.next_obses[idxs] + + obses_raw = unflatten(obses_raw, self.obs_space) + next_obses_raw = unflatten(next_obses_raw, self.obs_space) + + # 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 = combine_obs(state, obses) + next_obses = combine_obs(next_state, next_obses) + pos = combine_obs(state, pos) obses = torch.as_tensor(obses, device=self.device).float() next_obses = torch.as_tensor( @@ -189,6 +206,9 @@ class ReplayBuffer(Dataset): next_obs = self.next_obses[idx] not_done = self.not_dones[idx] + obs = unflatten(obs, self.obs_space) + next_obs = unflatten(next_obs, self.obs_space) + if self.transform: obs = self.transform(obs) next_obs = self.transform(next_obs) @@ -198,35 +218,6 @@ class ReplayBuffer(Dataset): def __len__(self): return self.capacity -class FrameStack(gym.Wrapper): - def __init__(self, env, k): - gym.Wrapper.__init__(self, env) - self._k = k - self._frames = deque([], maxlen=k) - shp = env.observation_space.shape - self.observation_space = gym.spaces.Box( - low=0, - high=1, - shape=((shp[0] * k,) + shp[1:]), - dtype=env.observation_space.dtype - ) - self._max_episode_steps = env._max_episode_steps - - def reset(self): - obs = self.env.reset() - for _ in range(self._k): - self._frames.append(obs) - return self._get_obs() - - def step(self, action): - obs, reward, done, info = self.env.step(action) - self._frames.append(obs) - return self._get_obs(), reward, done, info - - def _get_obs(self): - assert len(self._frames) == self._k - return np.concatenate(list(self._frames), axis=0) - def random_crop(imgs, output_size): """ @@ -260,5 +251,16 @@ def center_crop_image(image, output_size): 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