From fadb3d7d35e5bec242e7b66681f42e007f9055d7 Mon Sep 17 00:00:00 2001 From: wassname Date: Thu, 18 Feb 2021 13:04:44 +0800 Subject: [PATCH] fix --- curl_sac.py | 4 ++-- encoder.py | 4 +++- scripts/run.sh | 2 +- train.py | 6 +++--- utils.py | 56 +++++++++++++++++++++++--------------------------- 5 files changed, 35 insertions(+), 37 deletions(-) diff --git a/curl_sac.py b/curl_sac.py index 002b69b..cdedcab 100644 --- a/curl_sac.py +++ b/curl_sac.py @@ -206,9 +206,9 @@ class CURL(nn.Module): """ if ema: with torch.no_grad(): - z_out = self.encoder_target(x) + z_out = self.encoder_target(x, with_state=False) else: - z_out = self.encoder(x) + z_out = self.encoder(x, with_state=False) if detach: z_out = z_out.detach() diff --git a/encoder.py b/encoder.py index db56c6b..876f951 100644 --- a/encoder.py +++ b/encoder.py @@ -120,8 +120,10 @@ class MixedEncoder(PixelEncoder): img_shape = obs_shape['img'] super().__init__(img_shape, feature_dim, num_layers, num_filters, output_logits) self.feature_dim = feature_dim + obs_shape['state'][0] - def forward(self, obs, detach=False): + def forward(self, obs, detach=False, with_state=True): h = super().forward(obs['img'], detach) + if not with_state: + return h return torch.cat([obs['state'], h], 1) diff --git a/scripts/run.sh b/scripts/run.sh index e52bb69..d78b28e 100755 --- a/scripts/run.sh +++ b/scripts/run.sh @@ -3,4 +3,4 @@ set -e CUDA_VISIBLE_DEVICES=1 /home/wassname/anaconda/envs/diygym4/bin/python \ -m pdb -c continue \ train.py \ - --save_tb + --save_tb --init_steps 10 diff --git a/train.py b/train.py index 3b67fb9..7ed5782 100644 --- a/train.py +++ b/train.py @@ -39,7 +39,7 @@ def parse_args(): # 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("--num_train_steps", default=3000000, type=int) parser.add_argument("--batch_size", default=32, type=int) parser.add_argument("--hidden_dim", default=1024, type=int) # eval @@ -261,10 +261,10 @@ def main(): episode, episode_reward, done = 0, 0, True start_time = time.time() - for step in tqdm(range(args.num_train_steps), desc="train", unit="step"): + for step in tqdm(range(args.num_train_steps), desc="train", unit="step", mininterval=360): # evaluate agent periodically - if step % args.eval_freq == 400: + if (step % args.eval_freq == 0) and (step > args.eval_freq): L.log("eval/episode", episode, step) evaluate(env, agent, video, args.num_eval_episodes, L, step, args) if args.save_model: diff --git a/utils.py b/utils.py index 2ddb2ac..a70642b 100644 --- a/utils.py +++ b/utils.py @@ -103,24 +103,26 @@ class ReplayBuffer(Dataset): self.idx = (self.idx + 1) % self.capacity self.full = self.full or self.idx == 0 + def unflatten_obs(self, obs): + obs = [unflatten(o, self.obs_space) for o in obs] + obs = { + k: np.stack([o[k] for o in obs]) + for k in obs[0].keys() + } + return obs + + def sample_proprio(self): 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] + obses = self.as_tensor_obs(self.unflatten_obs(self.obses[idxs])) + next_obses = self.as_tensor_obs(self.unflatten_obs(self.next_obses[idxs])) - obses = unflatten(obses, self.obs_space) - next_obses = unflatten(next_obses, self.obs_space) - - obses = torch.as_tensor(obses, device=self.device).float() actions = torch.as_tensor(self.actions[idxs], device=self.device) rewards = torch.as_tensor(self.rewards[idxs], device=self.device) - next_obses = torch.as_tensor( - next_obses, device=self.device - ).float() not_dones = torch.as_tensor(self.not_dones[idxs], device=self.device) return obses, actions, rewards, next_obses, not_dones @@ -130,12 +132,9 @@ class ReplayBuffer(Dataset): idxs = np.random.randint( 0, self.capacity if self.full else self.idx, size=self.batch_size ) - - obses_raw = self.obses[idxs] - next_obses_raw = self.next_obses[idxs] - - obses_raw = unflatten(obses_raw, self.obs_space) - next_obses_raw = unflatten(next_obses_raw, self.obs_space) + + obses_raw = self.unflatten_obs(self.obses[idxs]) + next_obses_raw = self.unflatten_obs(self.next_obses[idxs]) # Split mixed obs into image and state state, obses = split_obs(obses_raw) @@ -149,24 +148,24 @@ class ReplayBuffer(Dataset): pos = random_crop(pos, self.image_size) # Recombine - obses = combine_obs(state, obses) - next_obses = combine_obs(next_state, next_obses) - pos = combine_obs(state, pos) - - obses = torch.as_tensor(obses, device=self.device).float() - next_obses = torch.as_tensor( - next_obses, device=self.device - ).float() + obses = self.as_tensor_obs(combine_obs(state, obses)) + next_obses = self.as_tensor_obs(combine_obs(next_state, next_obses)) + pos = self.as_tensor_obs(combine_obs(state, pos)) + 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() 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 as_tensor_obs(self, obses): + obses['img'] = torch.as_tensor(obses['img'], device=self.device).float() + obses['state'] = torch.as_tensor(obses['state'], device=self.device).float() + return obses + def save(self, save_dir): if self.idx == self.last_save: return @@ -201,15 +200,12 @@ class ReplayBuffer(Dataset): 0, self.capacity if self.full else self.idx, size=1 ) idx = idx[0] - obs = self.obses[idx] + obs = self.unflatten_obs(self.obses[idx]) action = self.actions[idx] reward = self.rewards[idx] - next_obs = self.next_obses[idx] + next_obs = self.unflatten_obs(self.next_obses[idx]) not_done = self.not_dones[idx] - obs = unflatten(obs, self.obs_space) - next_obs = unflatten(next_obs, self.obs_space) - if self.transform: obs = self.transform(obs) next_obs = self.transform(next_obs) @@ -217,7 +213,7 @@ class ReplayBuffer(Dataset): return obs, action, reward, next_obs, not_done def __len__(self): - return self.capacity + return self.capacity def random_crop(imgs, output_size):