From fe75a963ea491fd9dec8a9852f2350b3a426340d Mon Sep 17 00:00:00 2001 From: MishaLaskin Date: Sat, 8 Feb 2020 16:36:30 -0800 Subject: [PATCH] curl --- .gitignore | 3 + .ipynb_checkpoints/Plots-checkpoint.ipynb | 6 + README.md | 60 +++ conda_env.yml | 17 + curl_sac.py | 490 +++++++++++++++++++++ encoder.py | 124 ++++++ logger.py | 163 +++++++ scripts/long_experiments.sh | 0 scripts/run.sh | 66 +++ scripts/run_cheetah.sh | 65 +++ scripts/short_experiments.sh | 0 train.py | 273 ++++++++++++ utils.py | 497 ++++++++++++++++++++++ video.py | 38 ++ wrappers.py | 293 +++++++++++++ 15 files changed, 2095 insertions(+) create mode 100644 .gitignore create mode 100644 .ipynb_checkpoints/Plots-checkpoint.ipynb create mode 100644 README.md create mode 100644 conda_env.yml create mode 100644 curl_sac.py create mode 100644 encoder.py create mode 100644 logger.py create mode 100755 scripts/long_experiments.sh create mode 100755 scripts/run.sh create mode 100644 scripts/run_cheetah.sh create mode 100755 scripts/short_experiments.sh create mode 100644 train.py create mode 100644 utils.py create mode 100644 video.py create mode 100644 wrappers.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..13ee690 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +tmp/ +notebooks +__pycache__ diff --git a/.ipynb_checkpoints/Plots-checkpoint.ipynb b/.ipynb_checkpoints/Plots-checkpoint.ipynb new file mode 100644 index 0000000..2fd6442 --- /dev/null +++ b/.ipynb_checkpoints/Plots-checkpoint.ipynb @@ -0,0 +1,6 @@ +{ + "cells": [], + "metadata": {}, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/README.md b/README.md new file mode 100644 index 0000000..85dffa1 --- /dev/null +++ b/README.md @@ -0,0 +1,60 @@ +# SAC+CPC implementaiton in PyTorch + +# + +## Instructions +To train an SAC+CPC agent on the `cheetah run` task from image-based observations run: +``` +CUDA_VISIBLE_DEVICES=4 python train_cpc.py \ + --dmc2gym \ + --domain_name reacher \ + --task_name easy \ + --encoder_type pixel \ + --decoder_type identity \ + --action_repeat 4 --frame_stack 1 \ + --save_tb --save_video --num_train_steps 1000000 \ + --work_dir ./tmp/dmc/reacher_easy_cpc \ + --agent sac_cpc \ + --seed 2 --critic_lr 1e-3 --actor_lr 1e-3 --eval_freq 10000 + +CUDA_VISIBLE_DEVICES=7 python train.py \ + --domain_name walker \ + --task_name walk --dmc2gym \ + --encoder_type pixel \ + --decoder_type identity \ + --action_repeat 4 \ + --save_tb --pre_transform_image_size 84 --image_size 84 \ + --work_dir ./tmp/icml/vanilla_sac/ML0107walker_vanilla_sac \ + --agent sac_ae --frame_stack 3 \ + --seed -1 --critic_lr 1e-3 --actor_lr 1e-3 --eval_freq 10000 --batch_size 128 --num_train_steps 1000000 +``` +Try - reducing log std actor max from 2->1 or increase 2->3 + + +This will produce 'log' folder, where all the outputs are going to be stored including train/eval logs, tensorboard blobs, and evaluation episode videos. One can attacha tensorboard to monitor training by running: +``` +tensorboard --logdir log +``` +and opening up tensorboad in your browser. + +The console output is also available in a form: +``` +| train | E: 1 | S: 1000 | D: 0.8 s | R: 0.0000 | BR: 0.0000 | ALOSS: 0.0000 | CLOSS: 0.0000 | RLOSS: 0.0000 +``` +a training entry decodes as: +``` +train - training episode +E - total number of episodes +S - total number of environment steps +D - duration in seconds to train 1 episode +R - episode reward +BR - average reward of sampled batch +ALOSS - average loss of actor +CLOSS - average loss of critic +RLOSS - average reconstruction loss (only if is trained from pixels and decoder) +``` +while an evaluation entry: +``` +| eval | S: 0 | ER: 21.1676 +``` +which just tells the expected reward `ER` evaluating current policy after `S` steps. Note that `ER` is average evaluation performance over `num_eval_episodes` episodes (usually 10). diff --git a/conda_env.yml b/conda_env.yml new file mode 100644 index 0000000..f47a2b9 --- /dev/null +++ b/conda_env.yml @@ -0,0 +1,17 @@ +name: pytorch_sac_ae +channels: + - defaults +dependencies: + - python=3.6 + - pytorch + - torchvision + - cudatoolkit=9.2 + - absl-py + - pyparsing + - pip: + - termcolor + - git+git://github.com/deepmind/dm_control.git + - git+git://github.com/1nadequacy/dmc2gym.git + - tb-nightly + - imageio + - imageio-ffmpeg \ No newline at end of file diff --git a/curl_sac.py b/curl_sac.py new file mode 100644 index 0000000..99c0f2d --- /dev/null +++ b/curl_sac.py @@ -0,0 +1,490 @@ +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +import copy +import math + +import utils +from encoder import make_encoder + +LOG_FREQ = 10000 + + +def gaussian_logprob(noise, log_std): + """Compute Gaussian log probability.""" + residual = (-0.5 * noise.pow(2) - log_std).sum(-1, keepdim=True) + return residual - 0.5 * np.log(2 * np.pi) * noise.size(-1) + + +def squash(mu, pi, log_pi): + """Apply squashing function. + See appendix C from https://arxiv.org/pdf/1812.05905.pdf. + """ + mu = torch.tanh(mu) + if pi is not None: + pi = torch.tanh(pi) + if log_pi is not None: + log_pi -= torch.log(F.relu(1 - pi.pow(2)) + 1e-6).sum(-1, keepdim=True) + return mu, pi, log_pi + + +def weight_init(m): + """Custom weight init for Conv2D and Linear layers.""" + if isinstance(m, nn.Linear): + nn.init.orthogonal_(m.weight.data) + m.bias.data.fill_(0.0) + elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d): + # delta-orthogonal init from https://arxiv.org/pdf/1806.05393.pdf + assert m.weight.size(2) == m.weight.size(3) + m.weight.data.fill_(0.0) + m.bias.data.fill_(0.0) + mid = m.weight.size(2) // 2 + gain = nn.init.calculate_gain('relu') + nn.init.orthogonal_(m.weight.data[:, :, mid, mid], gain) + + +class Actor(nn.Module): + """MLP actor network.""" + def __init__( + self, obs_shape, action_shape, hidden_dim, encoder_type, + encoder_feature_dim, log_std_min, log_std_max, num_layers, num_filters + ): + super().__init__() + + self.encoder = make_encoder( + encoder_type, obs_shape, encoder_feature_dim, num_layers, + num_filters, output_logits=True + ) + + self.log_std_min = log_std_min + self.log_std_max = log_std_max + + self.trunk = nn.Sequential( + nn.Linear(self.encoder.feature_dim, hidden_dim), nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), + nn.Linear(hidden_dim, 2 * action_shape[0]) + ) + + self.outputs = dict() + self.apply(weight_init) + + def forward( + self, obs, compute_pi=True, compute_log_pi=True, detach_encoder=False + ): + obs = self.encoder(obs, detach=detach_encoder) + + mu, log_std = self.trunk(obs).chunk(2, dim=-1) + + # constrain log_std inside [log_std_min, log_std_max] + log_std = torch.tanh(log_std) + log_std = self.log_std_min + 0.5 * ( + self.log_std_max - self.log_std_min + ) * (log_std + 1) + + self.outputs['mu'] = mu + self.outputs['std'] = log_std.exp() + + if compute_pi: + std = log_std.exp() + noise = torch.randn_like(mu) + pi = mu + noise * std + else: + pi = None + entropy = None + + if compute_log_pi: + log_pi = gaussian_logprob(noise, log_std) + else: + log_pi = None + + mu, pi, log_pi = squash(mu, pi, log_pi) + + return mu, pi, log_pi, log_std + + def log(self, L, step, log_freq=LOG_FREQ): + if step % log_freq != 0: + return + + for k, v in self.outputs.items(): + L.log_histogram('train_actor/%s_hist' % k, v, step) + + L.log_param('train_actor/fc1', self.trunk[0], step) + L.log_param('train_actor/fc2', self.trunk[2], step) + L.log_param('train_actor/fc3', self.trunk[4], step) + + +class QFunction(nn.Module): + """MLP for q-function.""" + def __init__(self, obs_dim, action_dim, hidden_dim): + super().__init__() + + self.trunk = nn.Sequential( + nn.Linear(obs_dim + action_dim, hidden_dim), nn.ReLU(), + nn.Linear(hidden_dim, hidden_dim), nn.ReLU(), + nn.Linear(hidden_dim, 1) + ) + + def forward(self, obs, action): + assert obs.size(0) == action.size(0) + + obs_action = torch.cat([obs, action], dim=1) + return self.trunk(obs_action) + + +class Critic(nn.Module): + """Critic network, employes two q-functions.""" + def __init__( + self, obs_shape, action_shape, hidden_dim, encoder_type, + encoder_feature_dim, num_layers, num_filters + ): + super().__init__() + + + self.encoder = make_encoder( + encoder_type, obs_shape, encoder_feature_dim, num_layers, + num_filters, output_logits=True + ) + + self.Q1 = QFunction( + self.encoder.feature_dim, action_shape[0], hidden_dim + ) + self.Q2 = QFunction( + self.encoder.feature_dim, action_shape[0], hidden_dim + ) + + self.outputs = dict() + self.apply(weight_init) + + def forward(self, obs, action, detach_encoder=False): + # detach_encoder allows to stop gradient propogation to encoder + obs = self.encoder(obs, detach=detach_encoder) + + q1 = self.Q1(obs, action) + q2 = self.Q2(obs, action) + + self.outputs['q1'] = q1 + self.outputs['q2'] = q2 + + return q1, q2 + + def log(self, L, step, log_freq=LOG_FREQ): + if step % log_freq != 0: + return + + self.encoder.log(L, step, log_freq) + + for k, v in self.outputs.items(): + L.log_histogram('train_critic/%s_hist' % k, v, step) + + for i in range(3): + L.log_param('train_critic/q1_fc%d' % i, self.Q1.trunk[i * 2], step) + L.log_param('train_critic/q2_fc%d' % i, self.Q2.trunk[i * 2], step) + + +class CURL(nn.Module): + """ + CURL + """ + + def __init__(self, obs_shape, z_dim, batch_size, critic, critic_target, output_type="continuous"): + super(CURL, self).__init__() + self.batch_size = batch_size + + self.encoder = critic.encoder + #PixelEncoder(obs_shape, z_dim, num_layers=2, num_filters=32) + + self.encoder_target = critic_target.encoder + #PixelEncoder(obs_shape, z_dim, num_layers=2, num_filters=32) + + #self.encoder_target.load_state_dict(self.encoder.state_dict()) + + self.W = nn.Parameter(torch.rand(z_dim, z_dim)) + self.output_type = output_type + + def encode(self, x, detach=False, ema=False): + """ + Encoder: z_t = e(x_t) + :param x: x_t, x y coordinates + :return: z_t, value in r2 + """ + if ema: + with torch.no_grad(): + z_out = self.encoder_target(x) + else: + z_out = self.encoder(x) + + if detach: + z_out = z_out.detach() + return z_out + + #def update_target(self): + # utils.soft_update_params(self.encoder, self.encoder_target, 0.05) + + def compute_logits(self, z_a, z_pos): + """ + Uses logits trick for CURL: + - compute (B,B) matrix z_a (W z_pos.T) + - positives are all diagonal elements + - negatives are all other elements + - to compute loss use multiclass cross entropy with identity matrix for labels + """ + Wz = torch.matmul(self.W, z_pos.T) # (z_dim,B) + logits = torch.matmul(z_a, Wz) # (B,B) + logits = logits - torch.max(logits, 1)[0][:, None] + return logits + +class CurlSacAgent(object): + """CURL representation learning with SAC.""" + def __init__( + self, + obs_shape, + action_shape, + device, + hidden_dim=256, + discount=0.99, + init_temperature=0.01, + alpha_lr=1e-3, + alpha_beta=0.9, + actor_lr=1e-3, + actor_beta=0.9, + actor_log_std_min=-10, + actor_log_std_max=2, + actor_update_freq=2, + critic_lr=1e-3, + critic_beta=0.9, + critic_tau=0.005, + critic_target_update_freq=2, + encoder_type='pixel', + encoder_feature_dim=50, + encoder_lr=1e-3, + encoder_tau=0.005, + num_layers=4, + num_filters=32, + cpc_update_freq=1, + log_interval=100, + detach_encoder=False, + curl_latent_dim=128 + ): + self.device = device + self.discount = discount + self.critic_tau = critic_tau + self.encoder_tau = encoder_tau + self.actor_update_freq = actor_update_freq + 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.curl_latent_dim = curl_latent_dim + self.detach_encoder = detach_encoder + + self.actor = Actor( + obs_shape, action_shape, hidden_dim, encoder_type, + encoder_feature_dim, actor_log_std_min, actor_log_std_max, + num_layers, num_filters + ).to(device) + + self.critic = Critic( + obs_shape, action_shape, hidden_dim, encoder_type, + encoder_feature_dim, num_layers, num_filters + ).to(device) + + self.critic_target = Critic( + obs_shape, action_shape, hidden_dim, encoder_type, + encoder_feature_dim, num_layers, num_filters + ).to(device) + + self.critic_target.load_state_dict(self.critic.state_dict()) + + # 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) + + # tie encoders between actor and critic, and CURL and critic + self.actor.encoder.copy_conv_weights_from(self.critic.encoder) + + self.log_alpha = torch.tensor(np.log(init_temperature)).to(device) + self.log_alpha.requires_grad = True + # set target entropy to -|A| + self.target_entropy = -np.prod(action_shape) + + # optimizers + self.actor_optimizer = torch.optim.Adam( + self.actor.parameters(), lr=actor_lr, betas=(actor_beta, 0.999) + ) + + self.critic_optimizer = torch.optim.Adam( + self.critic.parameters(), lr=critic_lr, betas=(critic_beta, 0.999) + ) + + self.log_alpha_optimizer = torch.optim.Adam( + [self.log_alpha], lr=alpha_lr, betas=(alpha_beta, 0.999) + ) + + # optimizer for critic encoder for reconstruction loss + self.encoder_optimizer = torch.optim.Adam( + self.critic.encoder.parameters(), lr=encoder_lr + ) + + self.cpc_optimizer = torch.optim.Adam( + self.CURL.parameters(), lr=encoder_lr + ) + self.cross_entropy_loss = nn.CrossEntropyLoss() + + self.train() + self.critic_target.train() + + def train(self, training=True): + self.training = training + self.actor.train(training) + self.critic.train(training) + self.CURL.train(training) + + @property + def alpha(self): + return self.log_alpha.exp() + + def select_action(self, obs): + with torch.no_grad(): + obs = torch.FloatTensor(obs).to(self.device) + obs = obs.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) + + with torch.no_grad(): + obs = torch.FloatTensor(obs).to(self.device) + obs = obs.unsqueeze(0) + mu, pi, _, _ = self.actor(obs, compute_log_pi=False) + return pi.cpu().data.numpy().flatten() + + def update_critic(self, obs, action, reward, next_obs, not_done, L, step): + with torch.no_grad(): + _, policy_action, log_pi, _ = self.actor(next_obs) + target_Q1, target_Q2 = self.critic_target(next_obs, policy_action) + target_V = torch.min(target_Q1, + target_Q2) - self.alpha.detach() * log_pi + target_Q = reward + (not_done * self.discount * target_V) + + # get current Q estimates + current_Q1, current_Q2 = self.critic( + obs, action, detach_encoder=self.detach_encoder) + critic_loss = F.mse_loss(current_Q1, + target_Q) + F.mse_loss(current_Q2, target_Q) + if step % self.log_interval == 0: + L.log('train_critic/loss', critic_loss, step) + + + # Optimize the critic + self.critic_optimizer.zero_grad() + critic_loss.backward() + self.critic_optimizer.step() + + self.critic.log(L, step) + + def update_actor_and_alpha(self, obs, L, step): + # detach encoder, so we don't update it with the actor loss + _, pi, log_pi, log_std = self.actor(obs, detach_encoder=True) + actor_Q1, actor_Q2 = self.critic(obs, pi, detach_encoder=True) + + actor_Q = torch.min(actor_Q1, actor_Q2) + actor_loss = (self.alpha.detach() * log_pi - actor_Q).mean() + + if step % self.log_interval == 0: + L.log('train_actor/loss', actor_loss, step) + L.log('train_actor/target_entropy', self.target_entropy, step) + entropy = 0.5 * log_std.shape[1] * \ + (1.0 + np.log(2 * np.pi)) + log_std.sum(dim=-1) + if step % self.log_interval == 0: + L.log('train_actor/entropy', entropy.mean(), step) + + # optimize the actor + self.actor_optimizer.zero_grad() + actor_loss.backward() + self.actor_optimizer.step() + + self.actor.log(L, step) + + self.log_alpha_optimizer.zero_grad() + alpha_loss = (self.alpha * + (-log_pi - self.target_entropy).detach()).mean() + if step % self.log_interval == 0: + L.log('train_alpha/loss', alpha_loss, step) + L.log('train_alpha/value', self.alpha, step) + alpha_loss.backward() + self.log_alpha_optimizer.step() + + def update_cpc(self, obs_anchor, obs_pos, cpc_kwargs, L, step): + + # time flips + """ + time_pos = cpc_kwargs["time_pos"] + time_anchor= cpc_kwargs["time_anchor"] + obs_anchor = torch.cat((obs_anchor, time_anchor), 0) + obs_pos = torch.cat((obs_anchor, time_pos), 0) + """ + z_a = self.CURL.encode(obs_anchor) + z_pos = self.CURL.encode(obs_pos, ema=True) + + logits = self.CURL.compute_logits(z_a, z_pos) + labels = torch.arange(logits.shape[0]).long().to(self.device) + loss = self.cross_entropy_loss(logits, labels) + + self.encoder_optimizer.zero_grad() + self.cpc_optimizer.zero_grad() + loss.backward() + + self.encoder_optimizer.step() + self.cpc_optimizer.step() + if step % self.log_interval == 0: + L.log('train_cpc/cpc_loss', loss, step) + + + def update(self, replay_buffer, L, step): + obs, action, reward, next_obs, not_done, cpc_kwargs = replay_buffer.sample_cpc() + + if step % self.log_interval == 0: + L.log('train/batch_reward', reward.mean(), step) + + self.update_critic(obs, action, reward, next_obs, not_done, L, step) + + if step % self.actor_update_freq == 0: + self.update_actor_and_alpha(obs, L, step) + + if step % self.critic_target_update_freq == 0: + utils.soft_update_params( + self.critic.Q1, self.critic_target.Q1, self.critic_tau + ) + utils.soft_update_params( + self.critic.Q2, self.critic_target.Q2, self.critic_tau + ) + utils.soft_update_params( + self.critic.encoder, self.critic_target.encoder, + self.encoder_tau + ) + + obs_anchor, obs_pos = cpc_kwargs["obs_anchor"], cpc_kwargs["obs_pos"] + if step % self.cpc_update_freq == 0: + self.update_cpc(obs_anchor, obs_pos,cpc_kwargs, L, step) + + def save(self, model_dir, step): + torch.save( + self.actor.state_dict(), '%s/actor_%s.pt' % (model_dir, step) + ) + torch.save( + self.critic.state_dict(), '%s/critic_%s.pt' % (model_dir, step) + ) + + def load(self, model_dir, step): + self.actor.load_state_dict( + torch.load('%s/actor_%s.pt' % (model_dir, step)) + ) + self.critic.load_state_dict( + torch.load('%s/critic_%s.pt' % (model_dir, step)) + ) + \ No newline at end of file diff --git a/encoder.py b/encoder.py new file mode 100644 index 0000000..b90e070 --- /dev/null +++ b/encoder.py @@ -0,0 +1,124 @@ +import torch +import torch.nn as nn + + +def tie_weights(src, trg): + assert type(src) == type(trg) + trg.weight = src.weight + trg.bias = src.bias + + +OUT_DIM = {2: 39, 4: 35, 6: 31} +OUT_DIM_64 = {2: 29, 4: 25, 6: 21} + + +class PixelEncoder(nn.Module): + """Convolutional encoder of pixels observations.""" + def __init__(self, obs_shape, feature_dim, num_layers=2, num_filters=32,output_logits=False): + super().__init__() + + assert len(obs_shape) == 3 + self.obs_shape = obs_shape + self.feature_dim = feature_dim + self.num_layers = num_layers + # try 2 5x5s with strides 2x2. with samep adding, it should reduce 84 to 21, so with valid, it should be even smaller than 21. + self.convs = nn.ModuleList( + [nn.Conv2d(obs_shape[0], num_filters, 3, stride=2)] + ) + for i in range(num_layers - 1): + self.convs.append(nn.Conv2d(num_filters, num_filters, 3, stride=1)) + + out_dim = OUT_DIM_64[num_layers] if obs_shape[-1] == 64 else OUT_DIM[num_layers] + self.fc = nn.Linear(num_filters * out_dim * out_dim, self.feature_dim) + self.ln = nn.LayerNorm(self.feature_dim) + + self.outputs = dict() + self.output_logits = output_logits + + def reparameterize(self, mu, logstd): + std = torch.exp(logstd) + eps = torch.randn_like(std) + return mu + eps * std + + def forward_conv(self, obs): + obs = obs / 255. + self.outputs['obs'] = obs + + conv = torch.relu(self.convs[0](obs)) + self.outputs['conv1'] = conv + + for i in range(1, self.num_layers): + conv = torch.relu(self.convs[i](conv)) + self.outputs['conv%s' % (i + 1)] = conv + + h = conv.view(conv.size(0), -1) + return h + + def forward(self, obs, detach=False): + h = self.forward_conv(obs) + + if detach: + h = h.detach() + + h_fc = self.fc(h) + self.outputs['fc'] = h_fc + + h_norm = self.ln(h_fc) + self.outputs['ln'] = h_norm + + if self.output_logits: + out = h_norm + else: + out = torch.tanh(h_norm) + self.outputs['tanh'] = out + + return out + + def copy_conv_weights_from(self, source): + """Tie convolutional layers""" + # only tie conv layers + for i in range(self.num_layers): + tie_weights(src=source.convs[i], trg=self.convs[i]) + + def log(self, L, step, log_freq): + if step % log_freq != 0: + return + + for k, v in self.outputs.items(): + L.log_histogram('train_encoder/%s_hist' % k, v, step) + if len(v.shape) > 2: + L.log_image('train_encoder/%s_img' % k, v[0], step) + + for i in range(self.num_layers): + L.log_param('train_encoder/conv%s' % (i + 1), self.convs[i], step) + L.log_param('train_encoder/fc', self.fc, step) + L.log_param('train_encoder/ln', self.ln, step) + + +class IdentityEncoder(nn.Module): + def __init__(self, obs_shape, feature_dim, num_layers, num_filters,*args): + super().__init__() + + assert len(obs_shape) == 1 + self.feature_dim = obs_shape[0] + + def forward(self, obs, detach=False): + return obs + + def copy_conv_weights_from(self, source): + pass + + def log(self, L, step, log_freq): + pass + + +_AVAILABLE_ENCODERS = {'pixel': PixelEncoder, 'identity': IdentityEncoder} + + +def make_encoder( + encoder_type, obs_shape, feature_dim, num_layers, num_filters, output_logits=False +): + assert encoder_type in _AVAILABLE_ENCODERS + return _AVAILABLE_ENCODERS[encoder_type]( + obs_shape, feature_dim, num_layers, num_filters, output_logits + ) diff --git a/logger.py b/logger.py new file mode 100644 index 0000000..8e31fd4 --- /dev/null +++ b/logger.py @@ -0,0 +1,163 @@ +from torch.utils.tensorboard import SummaryWriter +from collections import defaultdict +import json +import os +import shutil +import torch +import torchvision +import numpy as np +from termcolor import colored + +FORMAT_CONFIG = { + 'rl': { + 'train': [ + ('episode', 'E', 'int'), ('step', 'S', 'int'), + ('duration', 'D', 'time'), ('episode_reward', 'R', 'float'), + ('batch_reward', 'BR', 'float'), ('actor_loss', 'ALOSS', 'float'), + ('critic_loss', 'CLOSS', 'float'), ('ae_loss', 'RLOSS', 'float') + ], + 'eval': [('step', 'S', 'int'), ('episode_reward', 'ER', 'float')] + } +} + + +class AverageMeter(object): + def __init__(self): + self._sum = 0 + self._count = 0 + + def update(self, value, n=1): + self._sum += value + self._count += n + + def value(self): + return self._sum / max(1, self._count) + + +class MetersGroup(object): + def __init__(self, file_name, formating): + self._file_name = file_name + if os.path.exists(file_name): + os.remove(file_name) + self._formating = formating + self._meters = defaultdict(AverageMeter) + + def log(self, key, value, n=1): + self._meters[key].update(value, n) + + def _prime_meters(self): + data = dict() + for key, meter in self._meters.items(): + if key.startswith('train'): + key = key[len('train') + 1:] + else: + key = key[len('eval') + 1:] + key = key.replace('/', '_') + data[key] = meter.value() + return data + + def _dump_to_file(self, data): + with open(self._file_name, 'a') as f: + f.write(json.dumps(data) + '\n') + + def _format(self, key, value, ty): + template = '%s: ' + if ty == 'int': + template += '%d' + elif ty == 'float': + template += '%.04f' + elif ty == 'time': + template += '%.01f s' + else: + raise 'invalid format type: %s' % ty + return template % (key, value) + + def _dump_to_console(self, data, prefix): + prefix = colored(prefix, 'yellow' if prefix == 'train' else 'green') + pieces = ['{:5}'.format(prefix)] + for key, disp_key, ty in self._formating: + value = data.get(key, 0) + pieces.append(self._format(disp_key, value, ty)) + print('| %s' % (' | '.join(pieces))) + + def dump(self, step, prefix): + if len(self._meters) == 0: + return + data = self._prime_meters() + data['step'] = step + self._dump_to_file(data) + self._dump_to_console(data, prefix) + self._meters.clear() + + +class Logger(object): + def __init__(self, log_dir, use_tb=True, config='rl'): + self._log_dir = log_dir + if use_tb: + tb_dir = os.path.join(log_dir, 'tb') + if os.path.exists(tb_dir): + shutil.rmtree(tb_dir) + self._sw = SummaryWriter(tb_dir) + else: + self._sw = None + self._train_mg = MetersGroup( + os.path.join(log_dir, 'train.log'), + formating=FORMAT_CONFIG[config]['train'] + ) + self._eval_mg = MetersGroup( + os.path.join(log_dir, 'eval.log'), + formating=FORMAT_CONFIG[config]['eval'] + ) + + def _try_sw_log(self, key, value, step): + if self._sw is not None: + self._sw.add_scalar(key, value, step) + + def _try_sw_log_image(self, key, image, step): + if self._sw is not None: + assert image.dim() == 3 + grid = torchvision.utils.make_grid(image.unsqueeze(1)) + self._sw.add_image(key, grid, step) + + def _try_sw_log_video(self, key, frames, step): + if self._sw is not None: + frames = torch.from_numpy(np.array(frames)) + frames = frames.unsqueeze(0) + self._sw.add_video(key, frames, step, fps=30) + + def _try_sw_log_histogram(self, key, histogram, step): + if self._sw is not None: + self._sw.add_histogram(key, histogram, step) + + def log(self, key, value, step, n=1): + assert key.startswith('train') or key.startswith('eval') + if type(value) == torch.Tensor: + value = value.item() + self._try_sw_log(key, value / n, step) + mg = self._train_mg if key.startswith('train') else self._eval_mg + mg.log(key, value, n) + + def log_param(self, key, param, step): + self.log_histogram(key + '_w', param.weight.data, step) + if hasattr(param.weight, 'grad') and param.weight.grad is not None: + self.log_histogram(key + '_w_g', param.weight.grad.data, step) + if hasattr(param, 'bias'): + self.log_histogram(key + '_b', param.bias.data, step) + if hasattr(param.bias, 'grad') and param.bias.grad is not None: + self.log_histogram(key + '_b_g', param.bias.grad.data, step) + + def log_image(self, key, image, step): + assert key.startswith('train') or key.startswith('eval') + self._try_sw_log_image(key, image, step) + + def log_video(self, key, frames, step): + assert key.startswith('train') or key.startswith('eval') + self._try_sw_log_video(key, frames, step) + + def log_histogram(self, key, histogram, step): + assert key.startswith('train') or key.startswith('eval') + self._try_sw_log_histogram(key, histogram, step) + + def dump(self, step): + self._train_mg.dump(step, 'train') + self._eval_mg.dump(step, 'eval') diff --git a/scripts/long_experiments.sh b/scripts/long_experiments.sh new file mode 100755 index 0000000..e69de29 diff --git a/scripts/run.sh b/scripts/run.sh new file mode 100755 index 0000000..90016db --- /dev/null +++ b/scripts/run.sh @@ -0,0 +1,66 @@ + +# curl cheetah, crop 76 > 64, grayscale + random crop, deep stack +# batch size = 256 instead of 128. maybe 256 makes 64x64 work. And, try using 512 with 64x64. +# then try adam LR (smaller) - 3e-4... You can try 2e-4 and 5e-4. +# try using the stochastic policy for eval. (you can do later.. for now the important thing is to run ablations.) +# try bigger frame stack, maybe 8. +# parser.add_argument('--critic_tau', default=0.01, type=float) # try 0.05 or 0.1 +# run 1: batch 256, first try with 84 +# run 2: batch 512, first try with 84, then try 64, then try their encoder +# run 3: batch 256, first try with 84 +# run 4: batch 512, first try with 84, then try 64, then try their encoder +# run 5: 2e-4 lr for all +# run 6: 5e-4 for all +# run 7: critic higher tau, 0.05 + +# try stochastic critic eval +CUDA_VISIBLE_DEVICES=1 python train_cpc.py \ + --domain_name cheetah \ + --task_name run --dmc2gym \ + --encoder_type pixel \ + --decoder_type identity \ + --action_repeat 4 --batch_size 256 \ + --save_tb --work_dir ./tmp/icml/feb2cheetah/curl_cheetah_b256_84 \ + --agent sac_cpc --frame_stack 3 --pre_transform_image_size 100 --image_size 84 \ + --seed 23 --critic_lr 1e-3 --actor_lr 1e-3 --eval_freq 10000 --batch_size 128 --num_train_steps 3000000 & + +CUDA_VISIBLE_DEVICES=2 python train_cpc.py \ + --domain_name cheetah \ + --task_name run --dmc2gym \ + --encoder_type pixel \ + --decoder_type identity \ + --action_repeat 4 \ + --save_tb --work_dir ./tmp/icml/feb2cheetah/curl_cheetah_b512_84 \ + --agent sac_cpc --frame_stack 3 --pre_transform_image_size 100 --image_size 84 \ + --seed 23 --critic_lr 1e-3 --actor_lr 1e-3 --eval_freq 20000 --batch_size 512 --num_train_steps 3000000 & + +CUDA_VISIBLE_DEVICES=3 python train.py \ + --domain_name cheetah \ + --task_name run --dmc2gym \ + --encoder_type pixel \ + --decoder_type identity \ + --action_repeat 4 --batch_size 256 \ + --save_tb --work_dir ./tmp/icml/feb2cheetah/rad_cheetah_b256_84 \ + --agent sac_ae --frame_stack 3 --pre_transform_image_size 100 --image_size 84 \ + --seed 23 --critic_lr 1e-3 --actor_lr 1e-3 --eval_freq 20000 --batch_size 128 --num_train_steps 3000000 & + +CUDA_VISIBLE_DEVICES=4 python train.py \ + --domain_name cheetah \ + --task_name run --dmc2gym \ + --encoder_type pixel \ + --decoder_type identity \ + --action_repeat 4 \ + --save_tb --work_dir ./tmp/icml/feb2cheetah/rad_cheetah_b512_84 \ + --agent sac_ae --frame_stack 3 --pre_transform_image_size 100 --image_size 84 \ + --seed 23 --critic_lr 1e-3 --actor_lr 1e-3 --eval_freq 20000 --batch_size 512 --num_train_steps 3000000 & + +CUDA_VISIBLE_DEVICES=7 python train_cpc.py \ + --domain_name cheetah \ + --task_name run --dmc2gym \ + --encoder_type pixel \ + --decoder_type identity \ + --action_repeat 4 \ + --save_tb --work_dir ./tmp/icml/feb2cheetah/curl_cheetah_b256_84_lr3e4 \ + --agent sac_cpc --frame_stack 3 --pre_transform_image_size 100 --image_size 84 \ + --seed 23 --encoder_lr 3e-4 --critic_lr 3e-4 --actor_lr 3e-4 \ + --eval_freq 20000 --batch_size 256 --num_train_steps 3000000 \ No newline at end of file diff --git a/scripts/run_cheetah.sh b/scripts/run_cheetah.sh new file mode 100644 index 0000000..11762a3 --- /dev/null +++ b/scripts/run_cheetah.sh @@ -0,0 +1,65 @@ +#!/bin/bash +# 256 +# 2e4 lr 2x +# 5e4 lr 2x +# 1e3 lr 2x + +CUDA_VISIBLE_DEVICES=1 python train_cpc.py \ + --domain_name cheetah \ + --task_name run --dmc2gym \ + --encoder_type pixel \ + --decoder_type identity \ + --action_repeat 4 --batch_size 256 \ + --save_tb --work_dir ./tmp/icml/feb3cheetah/curl_cheetah_b256_84_lr2e4_a \ + --agent sac_cpc --frame_stack 3 --pre_transform_image_size 100 --image_size 84 \ + --seed 23 --encoder_lr 2e-4 --critic_lr 2e-4 --actor_lr 2e-4 --eval_freq 20000 --batch_size 128 --num_train_steps 3000000 & + +CUDA_VISIBLE_DEVICES=2 python train_cpc.py \ + --domain_name cheetah \ + --task_name run --dmc2gym \ + --encoder_type pixel \ + --decoder_type identity \ + --action_repeat 4 --batch_size 256 \ + --save_tb --work_dir ./tmp/icml/feb3cheetah/curl_cheetah_b256_84_lr2e4_b \ + --agent sac_cpc --frame_stack 3 --pre_transform_image_size 100 --image_size 84 \ + --seed -1 --encoder_lr 2e-4 --critic_lr 2e-4 --actor_lr 2e-4 --eval_freq 20000 --batch_size 128 --num_train_steps 3000000 & + +CUDA_VISIBLE_DEVICES=3 python train_cpc.py \ + --domain_name cheetah \ + --task_name run --dmc2gym \ + --encoder_type pixel \ + --decoder_type identity \ + --action_repeat 4 --batch_size 256 \ + --save_tb --work_dir ./tmp/icml/feb3cheetah/curl_cheetah_b256_84_lr5e4_a \ + --agent sac_cpc --frame_stack 3 --pre_transform_image_size 100 --image_size 84 \ + --seed 23 --encoder_lr 5e-4 --critic_lr 5e-4 --actor_lr 5e-4 --eval_freq 20000 --batch_size 128 --num_train_steps 3000000 & + +CUDA_VISIBLE_DEVICES=4 python train_cpc.py \ + --domain_name cheetah \ + --task_name run --dmc2gym \ + --encoder_type pixel \ + --decoder_type identity \ + --action_repeat 4 --batch_size 256 \ + --save_tb --work_dir ./tmp/icml/feb3cheetah/curl_cheetah_b256_84_lr5e4_b \ + --agent sac_cpc --frame_stack 3 --pre_transform_image_size 100 --image_size 84 \ + --seed -1 --encoder_lr 5e-4 --critic_lr 5e-4 --actor_lr 5e-4 --eval_freq 20000 --batch_size 128 --num_train_steps 3000000 & + +CUDA_VISIBLE_DEVICES=5 python train_cpc.py \ + --domain_name cheetah \ + --task_name run --dmc2gym \ + --encoder_type pixel \ + --decoder_type identity \ + --action_repeat 4 --batch_size 256 \ + --save_tb --work_dir ./tmp/icml/feb3cheetah/curl_cheetah_b256_84_lr1e3_a \ + --agent sac_cpc --frame_stack 3 --pre_transform_image_size 100 --image_size 84 \ + --seed 23 --encoder_lr 1e-3 --critic_lr 1e-3 --actor_lr 1e-3 --eval_freq 20000 --batch_size 128 --num_train_steps 3000000 & + +CUDA_VISIBLE_DEVICES=6 python train_cpc.py \ + --domain_name cheetah \ + --task_name run --dmc2gym \ + --encoder_type pixel \ + --decoder_type identity \ + --action_repeat 4 --batch_size 256 \ + --save_tb --work_dir ./tmp/icml/feb3cheetah/curl_cheetah_b256_84_lr1e3_b \ + --agent sac_cpc --frame_stack 3 --pre_transform_image_size 100 --image_size 84 \ + --seed -1 --encoder_lr 1e-3 --critic_lr 1e-3 --actor_lr 1e-3 --eval_freq 20000 --batch_size 128 --num_train_steps 3000000 \ No newline at end of file diff --git a/scripts/short_experiments.sh b/scripts/short_experiments.sh new file mode 100755 index 0000000..e69de29 diff --git a/train.py b/train.py new file mode 100644 index 0000000..a3178b0 --- /dev/null +++ b/train.py @@ -0,0 +1,273 @@ +import numpy as np +import torch +import argparse +import os +import math +import gym +import sys +import random +import time +import json +import dmc2gym +import copy + +import utils +from logger import Logger +from video import VideoRecorder + +from curl_sac import CurlSacAgent +from torchvision import transforms + + +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('--image_size', default=84, type=int) + parser.add_argument('--action_repeat', default=1, type=int) + parser.add_argument('--frame_stack', default=3, type=int) + # replay buffer + parser.add_argument('--replay_buffer_capacity', default=100000, type=int) + # train + parser.add_argument('--agent', default='curl_sac', type=str) + parser.add_argument('--init_steps', default=1000, type=int) + parser.add_argument('--num_train_steps', default=1000000, type=int) + parser.add_argument('--batch_size', default=32, type=int) + parser.add_argument('--hidden_dim', default=1024, type=int) + # eval + parser.add_argument('--eval_freq', default=1000, type=int) + parser.add_argument('--num_eval_episodes', default=10, type=int) + # critic + parser.add_argument('--critic_lr', default=1e-3, type=float) + parser.add_argument('--critic_beta', default=0.9, type=float) + parser.add_argument('--critic_tau', default=0.01, type=float) # try 0.05 or 0.1 + parser.add_argument('--critic_target_update_freq', default=2, type=int) # try to change it to 1 and retain 0.01 above + # actor + parser.add_argument('--actor_lr', default=1e-3, type=float) + parser.add_argument('--actor_beta', default=0.9, type=float) + parser.add_argument('--actor_log_std_min', default=-10, type=float) + 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_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) + parser.add_argument('--num_layers', default=4, type=int) + parser.add_argument('--num_filters', default=32, type=int) + parser.add_argument('--curl_latent_dim', default=128, type=int) + # sac + parser.add_argument('--discount', default=0.99, type=float) + parser.add_argument('--init_temperature', default=0.1, type=float) + parser.add_argument('--alpha_lr', default=1e-4, type=float) + 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('--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') + parser.add_argument('--save_model', default=False, action='store_true') + parser.add_argument('--detach_encoder', default=False, action='store_true') + + parser.add_argument('--log_interval', default=100, type=int) + args = parser.parse_args() + return args + + +def evaluate(env, agent, video, num_episodes, L, step, args): + all_ep_rewards = [] + + def run_eval_loop(sample_stochastically=True): + prefix = 'stochastic' if sample_stochastically else '' + for i in range(num_episodes): + obs = env.reset() + video.init(enabled=(i == 0)) + done = False + episode_reward = 0 + while not done: + # center crop image + obs = utils.center_crop_image(obs,args.image_size) + with utils.eval_mode(agent): + if sample_stochastically: + action = agent.sample_action(obs) + else: + action = agent.select_action(obs) + obs, reward, done, _ = env.step(action) + video.record(env) + episode_reward += reward + + video.save('%d.mp4' % step) + L.log('eval/' + prefix + '_episode_reward', episode_reward, step) + all_ep_rewards.append(episode_reward) + + 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) + L.log('eval/' + prefix + 'best_episode_reward', best_ep_reward, step) + + run_eval_loop(sample_stochastically=True) + run_eval_loop(sample_stochastically=False) + L.dump(step) + + +def make_agent(obs_shape, action_shape, args, device): + if args.agent == 'curl_sac': + return CurlSacAgent( + obs_shape=obs_shape, + action_shape=action_shape, + device=device, + hidden_dim=args.hidden_dim, + discount=args.discount, + init_temperature=args.init_temperature, + alpha_lr=args.alpha_lr, + alpha_beta=args.alpha_beta, + actor_lr=args.actor_lr, + actor_beta=args.actor_beta, + actor_log_std_min=args.actor_log_std_min, + actor_log_std_max=args.actor_log_std_max, + actor_update_freq=args.actor_update_freq, + critic_lr=args.critic_lr, + critic_beta=args.critic_beta, + critic_tau=args.critic_tau, + critic_target_update_freq=args.critic_target_update_freq, + encoder_type=args.encoder_type, + encoder_feature_dim=args.encoder_feature_dim, + encoder_lr=args.encoder_lr, + encoder_tau=args.encoder_tau, + num_layers=args.num_layers, + num_filters=args.num_filters, + log_interval=args.log_interval, + detach_encoder=args.detach_encoder, + curl_latent_dim=args.curl_latent_dim + + ) + else: + assert 'agent is not supported: %s' % args.agent + +def main(): + args = parse_args() + if args.seed == -1: + args.__dict__["seed"] = np.random.randint(1,1000) + 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.seed(args.seed) + + # stack several consecutive frames together + if args.encoder_type == 'pixel': + env = utils.FrameStack(env, k=args.frame_stack) + + utils.make_dir(args.work_dir) + video_dir = utils.make_dir(os.path.join(args.work_dir, 'video')) + model_dir = utils.make_dir(os.path.join(args.work_dir, 'model')) + buffer_dir = utils.make_dir(os.path.join(args.work_dir, 'buffer')) + + video = VideoRecorder(video_dir if args.save_video else None) + + with open(os.path.join(args.work_dir, 'args.json'), 'w') as f: + json.dump(vars(args), f, sort_keys=True, indent=4) + + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + 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 + + replay_buffer = utils.ReplayBuffer( + obs_shape=pre_aug_obs_shape, + action_shape=action_shape, + capacity=args.replay_buffer_capacity, + batch_size=args.batch_size, + device=device, + image_size=args.image_size, + ) + + agent = make_agent( + obs_shape=obs_shape, + action_shape=action_shape, + args=args, + device=device + ) + + L = Logger(args.work_dir, use_tb=args.save_tb) + + episode, episode_reward, done = 0, 0, True + start_time = time.time() + + for step in range(args.num_train_steps): + # evaluate agent periodically + + if step % args.eval_freq == 0: + L.log('eval/episode', episode, step) + evaluate(env, agent, video, args.num_eval_episodes, L, step,args) + if args.save_model: + agent.save(model_dir, step) + if args.save_buffer: + replay_buffer.save(buffer_dir) + + if done: + if step > 0: + if step % args.log_interval == 0: + L.log('train/duration', time.time() - start_time, step) + L.dump(step) + start_time = time.time() + if step % args.log_interval == 0: + L.log('train/episode_reward', episode_reward, step) + + obs = env.reset() + done = False + episode_reward = 0 + episode_step = 0 + episode += 1 + if step % args.log_interval == 0: + L.log('train/episode', episode, step) + + # sample action for data collection + if step < args.init_steps: + action = env.action_space.sample() + else: + with utils.eval_mode(agent): + action = agent.sample_action(obs) + + # run training update + if step >= args.init_steps: + num_updates = 1 #args.init_steps if step == args.init_steps else 1 + for _ in range(num_updates): + agent.update(replay_buffer, L, step) + + next_obs, reward, done, _ = env.step(action) + + # allow infinit bootstrap + done_bool = 0 if episode_step + 1 == env._max_episode_steps else float( + done + ) + episode_reward += reward + #action = np.array([action], dtype="float32") + replay_buffer.add(obs, action, reward, next_obs, done_bool) + + obs = next_obs + episode_step += 1 + + +if __name__ == '__main__': + torch.multiprocessing.set_start_method('spawn') + + main() diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..71ced8d --- /dev/null +++ b/utils.py @@ -0,0 +1,497 @@ +import torch +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 +from torchvision import transforms +import time +from imgaug import augmenters as iaa +from skimage.util.shape import view_as_windows + +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): + try: + os.mkdir(dir_path) + except OSError: + pass + 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_shape, action_shape, capacity, batch_size, device,image_size=84,transform=None): + 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 + obs_dtype = np.float32 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 + + self.seq = iaa.Sequential([ + # crop images from each side by 0 to 16px (randomly chosen) + iaa.Crop(px=(0, 20)), + #iaa.Fliplr(0.5), + # iaa.Affine( + #scale={"x": (0.8, 1.2), "y": (0.8, 1.2)}, + #translate_percent={"x": (-0.2, 0.2), "y": (-0.2, 0.2)}, + #rotate=(-20, 20), + #shear=(-8, 8) + #) + + ], random_order=True) + + + + def add(self, obs, action, reward, next_obs, done): + + 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 sample(self,data_aug=False): + + start = time.time() + idxs = np.random.randint( + 0, self.capacity if self.full else self.idx, size=self.batch_size + ) + + obses = self.obses[idxs] + next_obses = self.next_obses[idxs] + + """ + Try to do a batch crop and see result + """ + + # crop batch + obses = fast_random_crop(obses, self.image_size) + next_obses = fast_random_crop(next_obses, self.image_size) + + + 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) + next_obses = torch.as_tensor( + next_obses, device=self.device + ).float() + not_dones = torch.as_tensor(self.not_dones[idxs], device=self.device) + #obses = random_grayscale_stack(obses,self.device,p=0.2) + #next_obses = random_grayscale_stack(next_obses,self.device,p=0.2) + 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 = self.obses[idxs] + next_obses = self.next_obses[idxs] + pos = obses.copy() + + #obses = self.seq(images=np.transpose(obses, (0, 2, 3, 1))) + #pos = self.seq(images=np.transpose(pos, (0, 2, 3, 1))) + #neg = self.seq(images=np.transpose(neg, (0, 2, 3, 1))) + + #next_obses = self.seq(images=np.transpose(next_obses, (0, 2, 3, 1))) + + #obses = np.transpose(obses, (0, 3, 1, 2)) + #next_obses = np.transpose(next_obses, (0, 3, 1, 2)) + #pos = np.transpose(pos, (0, 3, 1, 2)) + #neg = np.transpose(neg, (0, 3, 1, 2)) + + # random crop + # time flip + #time_flip_obses = obses[:, ::-1, ...].copy() + #time_pos = time_flip_obses.copy() + + obses = fast_random_crop(obses, self.image_size) + next_obses = fast_random_crop(next_obses, self.image_size) + pos = fast_random_crop(pos, self.image_size) + #time_flip_obses = fast_random_crop(time_flip_obses, 84) + #time_pos = fast_random_crop(time_pos, 84) + + # random flip + #obses = random_flip(obses,.2) + #next_obses = random_flip(next_obses, .2) + #pos = random_flip(pos, .2) + + obses = torch.as_tensor(obses, device=self.device).float() + #time_flip_obses = torch.as_tensor(time_flip_obses, device=self.device).float() + #time_pos = torch.as_tensor(time_pos, device=self.device).float() + next_obses = torch.as_tensor( + next_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) + not_dones = torch.as_tensor(self.not_dones[idxs], device=self.device) + + pos = torch.as_tensor(pos, device=self.device).float() + #obses = random_grayscale_stack(obses,self.device,p=.2) + #next_obses = random_grayscale_stack(obses,self.device,p=.2) + #pos = random_grayscale_stack(obses,self.device,p=.2) + # repeat negatives to make a batch of N + # most likely source of mistake is these negatives + # currently being stacked [[1,2,3],[1,2,3],...] + # but we may wany [[1,1,1],[2,2,2],[3,3,3]...] + #neg = neg.repeat(64,1,1,1) + #cpc_kwargs = dict(obs_anchor=obses,obs_pos=pos,time_anchor=time_flip_obses, time_pos=time_pos) + 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 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.obses[idx] + action = self.actions[idx] + reward = self.rewards[idx] + next_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 + +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) + +""" +Various transforms +""" + + +class RandomCrop(object): + """Crop randomly the image in a sample. + + Args: + output_size (tuple or int): Desired output size. If int, square crop + is made. + """ + + def __init__(self, output_size): + assert isinstance(output_size, (int, tuple)) + if isinstance(output_size, int): + self.output_size = (output_size, output_size) + else: + assert len(output_size) == 2 + self.output_size = output_size + + def __call__(self, image): + + h, w = image.shape[1:] + new_h, new_w = self.output_size + + top = np.random.randint(0, h - new_h) + left = np.random.randint(0, w - new_w) + + image = image[:, top: top + new_h, left: left + new_w] + + return image + +def random_crop(imgs,output_size): + h, w = imgs.shape[2:] + new_h, new_w = output_size, output_size + + if h > new_h: + top = np.random.randint(0, h - new_h) + left = np.random.randint(0, w - new_w) + + imgs = imgs[:,:, top: top + new_h, left: left + new_w] + + return imgs + + +def fast_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 random_flip(imgs, prob=0.2): + B = imgs.shape[0] + N = int(prob*B) + flipped_imgs = imgs[..., ::-1].copy() + idxs = np.random.choice(B, size=(N,), replace=False) + imgs[idxs] = flipped_imgs[idxs] + return imgs + +def time_flip(imgs,device): + + time_flipped_imgs = imgs[:,::-1, ...].copy() + all_imgs = np.concatenate((imgs, time_flipped_imgs), 0) + + return all_imgs + +def grayscale(imgs,device): + # imgs: b x c x h x w + b, c, h, w = imgs.shape + frames = c // 3 + + imgs = imgs.view([b,frames,3,h,w]) + imgs = imgs[:, :, 0, ...] * 0.2989 + imgs[:, :, 1, ...] * 0.587 + imgs[:, :, 2, ...] * 0.114 + + imgs = imgs.type(torch.uint8).float() + # assert len(imgs.shape) == 3, imgs.shape + imgs = imgs[:, :, None, :, :] + imgs = imgs * torch.ones([1, 1, 3, 1, 1], dtype=imgs.dtype).float().to(device) # broadcast tiling + return imgs + +def random_grayscale(images,device,p=1.): + # images: [B, C, H, W] + gray_images = grayscale(images,device) + rnd = np.random.uniform(0., 1., size=(images.shape[0],)) + mask = rnd <= p + mask = torch.from_numpy(mask) + frames = images.shape[1] // 3 + images = images.view(*gray_images.shape) + mask = mask[:, None] * torch.ones([1, frames]).type(mask.dtype) + mask = mask.type(images.dtype).to(device) + mask = mask[:, :, None, None, None] + return mask * gray_images + (1 - mask) * images + +def random_grayscale_stack(stack,device,p=0.5): + # stack: B X C x H x W, C = num_frames * 3. + bs, channels, h, w = stack.shape + num_frames = channels // 3 + #stack = stack.view([-1, 3, h, w]) + stack = random_grayscale(stack, device,p=p) + stack = stack.view([bs, -1, h, w]) + return stack + +def random_rotate(imgs): + k = np.random.randint(4) + imgs = np.ascontiguousarray(np.rot90(imgs,k=k,axes=(-2,-1))) + return 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 + +class CenterCrop(object): + """Center crop the image in a sample. + + Args: + output_size (tuple or int): Desired output size. If int, square crop + is made. + """ + + def __init__(self, output_size): + assert isinstance(output_size, (int,)) + self.output_size = (output_size, output_size) + + def __call__(self, image): + + h, w = image.shape[1:] + new_h, new_w = self.output_size + + top = (h - new_h)//2 + left = (w - new_w)//2 + + image = image[:, top: top + new_h, left: left + new_w] + + return image + +class ToTensor(object): + """Convert ndarrays in sample to Tensors.""" + + def __call__(self, image,device): + + # torch image: C X H X W + + + return torch.from_numpy(image,) + + +class Grayscale(object): + """Convert ndarrays in sample to grayscale randomly.""" + + def __init__(self, prob): + self.prob = prob + + def __call__(self, image): + + if self.prob > np.random.uniform(): + image = self.rgb2gray(image) + + return image + + def rgb2gray(self, rgb): + rgb = np.transpose(rgb, (1, 2, 0)) + rgb = np.expand_dims(np.dot(rgb[..., :3], [0.2989, 0.5870, 0.1140]), 0) + rgb = np.repeat(rgb, 3, 0) + return rgb.astype(np.uint8) + + +class Flip(object): + """Convert ndarrays in sample to flip randomly.""" + + def __init__(self, prob): + self.prob = prob + + def __call__(self, image): + + if self.prob > np.random.uniform(): + image = self.flip(image) + + return image + + def flip(self, img): + return np.transpose(img, (0, 2, 1)) + + diff --git a/video.py b/video.py new file mode 100644 index 0000000..0e319f8 --- /dev/null +++ b/video.py @@ -0,0 +1,38 @@ +import imageio +import os +import numpy as np + + +class VideoRecorder(object): + def __init__(self, dir_name, height=256, width=256, camera_id=0, fps=30): + self.dir_name = dir_name + self.height = height + self.width = width + self.camera_id = camera_id + self.fps = fps + self.frames = [] + + def init(self, enabled=True): + self.frames = [] + self.enabled = self.dir_name is not None and enabled + + def record(self, env): + if self.enabled: + try: + frame = env.render( + mode='rgb_array', + height=self.height, + width=self.width, + camera_id=self.camera_id + ) + except: + frame = env.render( + mode='rgb_array', + ) + + self.frames.append(frame) + + def save(self, file_name): + if self.enabled: + path = os.path.join(self.dir_name, file_name) + imageio.mimsave(path, self.frames, fps=self.fps) diff --git a/wrappers.py b/wrappers.py new file mode 100644 index 0000000..b01e601 --- /dev/null +++ b/wrappers.py @@ -0,0 +1,293 @@ +import gym +from gym import spaces +import numpy as np + +class GoalWrapper(gym.Wrapper): + + def __init__(self, env): + gym.Wrapper.__init__(self, env) + self.threshold = 0.05 + self._max_episode_steps = 200 + + def step(self, action): + obs, _, done, info = self.env.step(action) + reward, is_success = self.compute_goal_metrics(obs) + info['is_success'] = is_success + if not done: + done = is_success + return obs, reward, done, info + + def compute_goal_metrics(self, obs): + dist = np.linalg.norm(obs['achieved_goal'] - obs['desired_goal']) + is_success = dist < 0.05 + reward = is_success - 1 + return reward, is_success + + def compute_reward(self, obs): + reward, _, _ = self.compute_goal_metrics(obs) + return reward + + def reset(self): + obs = self.env.reset() + return obs + + +class ImageGoalDMCWrapper(gym.Wrapper): + def __init__(self, env): + gym.Wrapper.__init__(self, env) + + def step(self, action): + obs, r, done, info = env.step(action) + reward = 0 if r > 0 else - 1 + if reward == 0: + done = True + return dict(observation=obs, + achieved_goal=obs, + desired_goal=self.desired_goal) + + def reset(self): + self.desired_goal = self.env.reset() + obs = self.env.reset() + return dict(observation=obs, + achieved_goal=obs, + desired_goal=self.desired_goal) + + +class LatentGoalWrapper(gym.Wrapper): + + def __init__(self, env,encoder): + gym.Wrapper.__init__(self, env) + self.encoder = encoder + self.use_cuda = torch.cuda.is_available() + + def step(self, action): + obs_dict, r, done, info = env.step(action) + obs_dict = self.compute_latents(obs_dict) + reward = self.compute_reward(obs_dict) + if reward == 0: + done = True + return dict(observation=obs, + achieved_goal=obs, + desired_goal=self.desired_goal) + + def reset(self): + obs_dict = self.reset() + obs_dict = self.compute_latents(obs_dict) + + def compute_reward(self, obs_dict): + ag = obs_dict['latent_achieved_goal'] + dg = obs_dict['latent_desired_goal'] + dist = np.linalg.norm(ag - dg) + r = 0 if dist < 0.05 else -1 + + def compute_latents(self, obs_dict): + obs, ag, dg = obs_dict['observation'],obs_dict['achieved_goal'],obs_dict['desired_goal'] + obs = torch.tensor(obs) + ag = torch.tensor(ag) + dg = torch.tensor(dg) + if self.use_cuda: + obs = obs.cuda() + ag = ag.cuda() + dg = dg.cuda() + + if len(ag.shape) == 3: + obs = obs.unsqueeze(0) + ag = ag.unsqueeze(0) + dg = dg.unsqueeze(0) + # encode + stack = torch.cat((obs,ag,dg),dim=0) + encoded_latents = self.encoder (stack) + obs_latent = encoded_latents[0] + ag_latent = encoded_latents[1] + dg_latent = encoded_latents[2] + + obs_dict.update(dict(latent_observation=obs_latent, + latent_achieved_goal=ag_latent, + latent_desired_goal=dg_latent)) + return obs_dict + + +class FlatImageGoalDMCWrapper(gym.Wrapper): + def __init__(self, env): + gym.Wrapper.__init__(self, env) + + def step(self, action): + obs, r, done, info = env.step(action) + dist = self.dist_to_goal() + reward = 0 if dist < 0.1 else - 1 + if reward == 0: + done = True + return np.concatenate((obs, self.desired_goal), axis=0) + + def reset(self): + self.desired_goal = self.env.reset() + self.goal_xpos = self.env._env.physics.named.data.geom_xpos['finger'] + obs = self.env.reset() + return np.concatenate((obs, self.desired_goal), axis=0) + + @property + def current_xpos(self): + return self.env._env.physics.named.data.geom_xpos['finger'] + + def dist_to_goal(self): + return np.lingalg.norm(self.goal_xpos - self.current_xpos) + +class FlattenGoalWrapper(gym.Wrapper): + + def __init__(self, env): + gym.Wrapper.__init__(self, env) + self._max_episode_steps = 50 + + def step(self, action): + obs, _, done, info = self.env.step(action) + reward, is_success = self.compute_goal_metrics(obs) + info['is_success'] = is_success + if not done: + done = is_success + return obs, reward, done, info + + def compute_goal_metrics(self, obs): + dist = np.linalg.norm(obs['achieved_goal'] - obs['desired_goal']) + is_success = dist < 0.05 + reward = is_success - 1 + return reward, is_success + + def compute_reward(self, obs): + reward, _, _ = self.compute_goal_metrics(obs) + return reward + + def reset(self): + obs = self.env.reset() + return obs + +class MtnCarWrapper(gym.Wrapper): + + def __init__(self, env): + gym.Wrapper.__init__(self, env) + + obs = self.reset() + + self._max_episode_steps = 200 + + self.observation_space = gym.spaces.Dict( + desired_goal=spaces.Box(-np.inf, np.inf, + shape=(1,), dtype='float32'), + achieved_goal=spaces.Box(-np.inf, np.inf, + shape=(1,), dtype='float32'), + observation=spaces.Box(-np.inf, np.inf, + shape=obs['observation'].shape, dtype='float32'), + ) + + def step(self, action): + obs, reward, done, info = self.env.step(action) + obs_dict = dict( + observation=obs, desired_goal=self.desired_goal, achieved_goal=obs[0]) + return obs_dict, reward, done, info + + def reset(self): + obs = self.env.reset() + self.desired_goal = 0.05 + obs_dict = dict(observation=obs,desired_goal=self.desired_goal,achieved_goal=obs[0]) + return obs_dict + + +class MinigridCoordGoalWrapper(gym.core.Wrapper): + + def __init__(self, env): + super().__init__(env) + obs = self.reset() + obs_dim = 3 + obs_min = 0 + obs_max = self.width - 1 + dir_max = 3 + + self._max_episode_steps = 50 + + self.observation_space = gym.spaces.Dict( + desired_goal=spaces.Box(obs_min, obs_max, + shape=(obs_dim-1,), dtype='float32'), + achieved_goal=spaces.Box(obs_min, obs_max, + shape=(obs_dim-1,), dtype='float32'), + observation=spaces.Box(obs_min, obs_max, + shape=(obs_dim,), dtype='float32'), + ) + + def step(self, action): + self.steps +=1 + obs_max = self.width - 1 + dir_max = 3 + obs, reward, done, info = self.env.step(action) + pos = np.array(self.env.agent_pos) / obs_max + dir_ = np.array(self.env.agent_dir) / dir_max + obs = np.append(pos, dir_).astype(np.float32) + achieved_goal = pos + obs_dict = dict( + observation=obs, desired_goal=self.desired_goal, achieved_goal=achieved_goal) + reward = reward - 1 + info['is_success'] = reward > -1.0 + return obs_dict, reward, done, info + + def reset(self, **kwargs): + self.steps = 0 + obs_max = self.width - 1 + dir_max = 3 + self.desired_goal = np.array([7, 7]).astype(np.float32) / obs_max + self.env.reset() + pos = np.array(self.env.agent_pos) / obs_max + dir_ = np.array(self.env.agent_dir) / dir_max + obs = np.append(pos, dir_).astype(np.float32) + achieved_goal = pos + obs_dict = dict( + observation=obs, desired_goal=self.desired_goal, achieved_goal=achieved_goal) + return obs_dict + + +class MinigridImageGoalWrapper(gym.core.Wrapper): + + def __init__(self, env): + super().__init__(env) + obs = self.reset() + obs_dim = 3 + obs_min = 0 + obs_max = self.width - 1 + dir_max = 3 + + self._max_episode_steps = 50 + + self.observation_space = gym.spaces.Dict( + desired_goal=spaces.Box(obs_min, obs_max, + shape=(obs_dim-1,), dtype='float32'), + achieved_goal=spaces.Box(obs_min, obs_max, + shape=(obs_dim-1,), dtype='float32'), + observation=spaces.Box(obs_min, obs_max, + shape=(obs_dim,), dtype='float32'), + ) + + def step(self, action): + self.steps += 1 + obs_max = self.width - 1 + dir_max = 3 + obs, reward, done, info = self.env.step(action) + pos = np.array(self.env.agent_pos) / obs_max + dir_ = np.array(self.env.agent_dir) / dir_max + obs = np.append(pos, dir_).astype(np.float32) + achieved_goal = pos + obs_dict = dict( + observation=obs, desired_goal=self.desired_goal, achieved_goal=achieved_goal) + reward = reward - 1 + info['is_success'] = reward > -1.0 + return obs_dict, reward, done, info + + def reset(self,**kwargs): + self.steps = 0 + obs_max = self.width - 1 + dir_max = 3 + self.desired_goal = np.array([14, 14]).astype(np.float32) / obs_max + self.env.reset() + pos = np.array(self.env.agent_pos) / obs_max + dir_ = np.array(self.env.agent_dir) / dir_max + obs = np.append(pos, dir_).astype(np.float32) + achieved_goal = pos + obs_dict = dict( + observation=obs, desired_goal=self.desired_goal, achieved_goal=achieved_goal) + return obs_dict