curl

.gitignore

@@ -0,0 +1,3 @@
+tmp/
+notebooks
+__pycache__

.ipynb_checkpoints/Plots-checkpoint.ipynb

@@ -0,0 +1,6 @@
+{
+ "cells": [],
+ "metadata": {},
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

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).

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

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))
+        )
+ 

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
+    )

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')

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 

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 

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()

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))
+
+

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)

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