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

fix

Browse Source
This commit is contained in:
wassname
2021-02-18 13:04:44 +08:00
parent 5ac6c81256
commit fadb3d7d35
5 changed files with 35 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
curl_sac.py
+2 -2
View File
@@ -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()
encoder.py
+3 -1
View File
@@ -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)
scripts/run.sh
+1 -1
View File
@@ -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
train.py
+3 -3
View File
@@ -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:
utils.py
+26 -30
View File
@@ -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):