diff --git a/README.md b/README.md index aa92a18..1b396a8 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,7 @@ # CURL: Contrastive Unsupervised Representation Learning for Sample-Efficient Reinforcement Learning +This repository is the official implementation of [CURL](https://www.mishalaskin.github.io/curl). + ## Installation All of the dependencies are in the `conda_env.yml` file. They can be installed manually or with the following command: @@ -31,16 +33,16 @@ In your console, you should see printouts that look like: | train | E: 233 | S: 29500 | D: 19.6 s | R: 838.0947 | BR: 3.7254 | A_LOSS: -316.9415 | CR_LOSS: 136.5304 | CU_LOSS: 0.0000 ``` -The maximum score for cartpole swing up is around 845 pts. Notice how CURL solves visual cartpole in 30k steps! This takes about and hour of training depending on your GPU. For reference, the state-state-of-the-art end-to-end method D4PG takes 50,000,000 timesteps to solve the same problem. CURL is ~1000x more efficient! +The maximum score for cartpole swing up is around 845 pts. Notice how CURL solves visual cartpole in <50k steps! This takes about and hour of training depending on your GPU. For reference, the state-state-of-the-art end-to-end method D4PG takes 50M timesteps to solve the same problem. -The above output decodes as: +Log abbreviation mapping: ``` 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 +R - mean episode reward BR - average reward of sampled batch A_LOSS - average loss of actor CR_LOSS - average loss of critic @@ -53,5 +55,6 @@ All data related to the run is stored in the specified `working_dir`. To enable tensorboard --logdir log --port 6006 ``` -and go to `localhost:6006` in your browser. If you're running headlessly, try port forwarding with ssh. +and go to `localhost:6006` in your browser. If you're running headlessly, try port forwarding with ssh. +For GPU accelerated rendering, make sure EGL is installed on your machine and set `export MUJOCO_GL=egl`. For environment troubleshooting issues, see the DeepMind control documentation. diff --git a/curl_sac.py b/curl_sac.py index 8dfc87a..5836371 100644 --- a/curl_sac.py +++ b/curl_sac.py @@ -214,9 +214,6 @@ class CURL(nn.Module): 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: @@ -420,13 +417,6 @@ class CurlSacAgent(object): 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) diff --git a/encoder.py b/encoder.py index b90e070..9da499f 100644 --- a/encoder.py +++ b/encoder.py @@ -8,7 +8,9 @@ def tie_weights(src, trg): trg.bias = src.bias +# for 84 x 84 inputs OUT_DIM = {2: 39, 4: 35, 6: 31} +# for 64 x 64 inputs OUT_DIM_64 = {2: 29, 4: 25, 6: 21} @@ -21,7 +23,7 @@ class PixelEncoder(nn.Module): 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)] ) diff --git a/train.py b/train.py index 294dfa4..3768b31 100644 --- a/train.py +++ b/train.py @@ -112,7 +112,6 @@ def evaluate(env, agent, video, num_episodes, L, step, args): 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) @@ -260,7 +259,7 @@ def main(): # run training update if step >= args.init_steps: - num_updates = 1 #args.init_steps if step == args.init_steps else 1 + num_updates = 1 for _ in range(num_updates): agent.update(replay_buffer, L, step) @@ -271,7 +270,6 @@ def main(): done ) episode_reward += reward - #action = np.array([action], dtype="float32") replay_buffer.add(obs, action, reward, next_obs, done_bool) obs = next_obs diff --git a/utils.py b/utils.py index ac2e150..6603caf 100644 --- a/utils.py +++ b/utils.py @@ -131,9 +131,9 @@ class ReplayBuffer(Dataset): next_obses = self.next_obses[idxs] pos = 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) + obses = random_crop(obses, self.image_size) + next_obses = random_crop(next_obses, self.image_size) + pos = random_crop(pos, self.image_size) obses = torch.as_tensor(obses, device=self.device).float() next_obses = torch.as_tensor( @@ -227,53 +227,8 @@ class FrameStack(gym.Wrapper): 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): +def random_crop(imgs, output_size): """ Vectorized way to do random crop using sliding windows and picking out random ones @@ -295,63 +250,6 @@ def fast_random_crop(imgs, output_size): 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 @@ -362,75 +260,5 @@ def center_crop_image(image, output_size): 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))