wassname/curl
1
0
Fork 0
mirror of https://github.com/wassname/curl.git synced 2026-09-09 11:20:39 +08:00
Code Issues Packages Projects Releases Wiki Activity

working with obs=dict(state=state, img=img)

Browse Source
This commit is contained in:
wassname
2021-02-18 11:57:29 +08:00
parent 8416d6e386
commit b1eb2efb77
10 changed files with 192 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+31
View File
@@ -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
conda_env.yml
+1 -1
View File
@@ -17,4 +17,4 @@ dependencies:
- imageio
- imageio-ffmpeg
- torchvision
- scikit-image
- scikit-image
curl_sac.py
+14 -12
View File
@@ -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)
encoder.py
+14 -2
View File
@@ -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(
logger.py
+1
View File
@@ -7,6 +7,7 @@ import torch
import torchvision
import numpy as np
from termcolor import colored
from rich import print
FORMAT_CONFIG = {
'rl': {
runs/ApplePick-v0-02-18-im84-b32-s1-mixed/args.json
+44
View File
@@ -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"
}
runs/ApplePick-v0-02-18-im84-b32-s1-mixed/tb/events.out.tfevents.1613620610.mjcdesktop.457460.0
BIN
View File
Binary file not shown.
scripts/run.sh
+5 -9
View File
@@ -1,9 +1,5 @@
CUDA_VISIBLE_DEVICES=5 python train.py \
--domain_name cartpole \
--task_name swingup \
--encoder_type pixel \
--action_repeat 8 \
--save_tb --pre_transform_image_size 100 --image_size 84 \
--work_dir ./tmp/cartpole \
--agent curl_sac --frame_stack 3 \
--seed -1 --critic_lr 1e-3 --actor_lr 1e-3 --eval_freq 10000 --batch_size 128 --num_train_steps 1000000
#!/bin/bash
CUDA_VISIBLE_DEVICES=1 python \
-m pdb -c continue \
train.py \
--save_tb
train.py
+41 -33
View File
@@ -8,7 +8,7 @@ import sys
import random
import time
import json
import dmc2gym
# import dmc2gym
import copy
import utils
@@ -17,20 +17,21 @@ from video import VideoRecorder
from curl_sac import CurlSacAgent
from torchvision import transforms
import apple_gym.env
from diy_gym.utils import flatten, unflatten
def parse_args():
parser = argparse.ArgumentParser()
# environment
parser.add_argument('--domain_name', default='cheetah')
parser.add_argument('--task_name', default='run')
parser.add_argument('--pre_transform_image_size', default=100, type=int)
parser.add_argument('--domain_name', default='ApplePick-v0')
parser.add_argument('--pre_transform_image_size', default=124, type=int)
parser.add_argument('--image_size', default=84, type=int)
parser.add_argument('--action_repeat', default=1, type=int)
parser.add_argument('--frame_stack', default=3, type=int)
parser.add_argument('--render', action='store_true')
# replay buffer
parser.add_argument('--replay_buffer_capacity', default=100000, type=int)
parser.add_argument('--replay_buffer_capacity', default=10000, type=int)
# train
parser.add_argument('--agent', default='curl_sac', type=str)
parser.add_argument('--init_steps', default=1000, type=int)
@@ -52,7 +53,7 @@ def parse_args():
parser.add_argument('--actor_log_std_max', default=2, type=float)
parser.add_argument('--actor_update_freq', default=2, type=int)
# encoder
parser.add_argument('--encoder_type', default='pixel', type=str)
parser.add_argument('--encoder_type', default='mixed', type=str)
parser.add_argument('--encoder_feature_dim', default=50, type=int)
parser.add_argument('--encoder_lr', default=1e-3, type=float)
parser.add_argument('--encoder_tau', default=0.05, type=float)
@@ -66,7 +67,7 @@ def parse_args():
parser.add_argument('--alpha_beta', default=0.5, type=float)
# misc
parser.add_argument('--seed', default=1, type=int)
parser.add_argument('--work_dir', default='.', type=str)
parser.add_argument('--work_dir', default='./runs', type=str)
parser.add_argument('--save_tb', default=False, action='store_true')
parser.add_argument('--save_buffer', default=False, action='store_true')
parser.add_argument('--save_video', default=False, action='store_true')
@@ -91,8 +92,10 @@ def evaluate(env, agent, video, num_episodes, L, step, args):
episode_reward = 0
while not done:
# center crop image
if args.encoder_type == 'pixel':
obs = utils.center_crop_image(obs,args.image_size)
if args.encoder_type == 'mixed':
state, img = utils.split_obs(obs)
img = utils.center_crop_image(img, args.image_size)
obs = utils.combine_obs(state, img)
with utils.eval_mode(agent):
if sample_stochastically:
action = agent.sample_action(obs)
@@ -106,7 +109,7 @@ def evaluate(env, agent, video, num_episodes, L, step, args):
L.log('eval/' + prefix + 'episode_reward', episode_reward, step)
all_ep_rewards.append(episode_reward)
L.log('eval/' + prefix + 'eval_time', time.time()-start_time , step)
L.log('eval/' + prefix + 'eval_time', time.time() - start_time , step)
mean_ep_reward = np.mean(all_ep_rewards)
best_ep_reward = np.max(all_ep_rewards)
L.log('eval/' + prefix + 'mean_episode_reward', mean_ep_reward, step)
@@ -155,27 +158,31 @@ def main():
if args.seed == -1:
args.__dict__["seed"] = np.random.randint(1,1000000)
utils.set_seed_everywhere(args.seed)
env = dmc2gym.make(
domain_name=args.domain_name,
task_name=args.task_name,
seed=args.seed,
visualize_reward=False,
from_pixels=(args.encoder_type == 'pixel'),
height=args.pre_transform_image_size,
width=args.pre_transform_image_size,
frame_skip=args.action_repeat
)
env = gym.make(args.domain_name, render=args.render)
# TODO action repeat wrapper?
# env = dmc2gym.make(
# domain_name=args.domain_name,
# task_name=args.task_name,
# seed=args.seed,
# visualize_reward=False,
# from_mixeds=(args.encoder_type == 'mixed'),
# height=args.pre_transform_image_size,
# width=args.pre_transform_image_size,
# frame_skip=args.action_repeat
# )
env.seed(args.seed)
# stack several consecutive frames together
if args.encoder_type == 'pixel':
env = utils.FrameStack(env, k=args.frame_stack)
# # stack several consecutive frames together
if args.encoder_type == 'mixed':
from apple_gym.env.wrappers import FrameStack, ImageState, PermuteImages
env = FrameStack(PermuteImages(ImageState(env), keys=['img']), n=args.frame_stack, keys=['img'])
# make directory
ts = time.gmtime()
ts = time.strftime("%m-%d", ts)
env_name = args.domain_name + '-' + args.task_name
env_name = args.domain_name
exp_name = env_name + '-' + ts + '-im' + str(args.image_size) +'-b' \
+ str(args.batch_size) + '-s' + str(args.seed) + '-' + args.encoder_type
args.work_dir = args.work_dir + '/' + exp_name
@@ -194,16 +201,15 @@ def main():
action_shape = env.action_space.shape
if args.encoder_type == 'pixel':
obs_shape = (3*args.frame_stack, args.image_size, args.image_size)
pre_aug_obs_shape = (3*args.frame_stack,args.pre_transform_image_size,args.pre_transform_image_size)
else:
obs_shape = env.observation_space.shape
pre_aug_obs_shape = obs_shape
# TODO, I need cropped aug shape now...
# TODO I need to make split obs and combine obs?
img = env.observation_space.sample()['img']
img_aug = utils.center_crop_image(img, args.image_size)
obs_shape = {'img': img_aug.shape, 'state': env.observation_space['state'].shape}
replay_buffer = utils.ReplayBuffer(
obs_shape=pre_aug_obs_shape,
action_shape=action_shape,
obs_space=env.observation_space,
action_space=env.action_space,
capacity=args.replay_buffer_capacity,
batch_size=args.batch_size,
device=device,
@@ -243,6 +249,7 @@ def main():
L.log('train/episode_reward', episode_reward, step)
obs = env.reset()
assert env.observation_space.contains(obs), f'obs should be in space. ob={obs} space={env.observation_space}'
done = False
episode_reward = 0
episode_step = 0
@@ -256,6 +263,7 @@ def main():
else:
with utils.eval_mode(agent):
action = agent.sample_action(obs)
assert env.action_space.contains(action), f'obs should be in space. ob={action} space={env.action_space}'
# run training update
if step >= args.init_steps:
utils.py
+41 -39
View File
@@ -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