VedantDave/Curiosity
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

New trained model

Browse Source
This commit is contained in:
Vedant Dave 2023-04-18 16:47:30 +02:00
parent 21cefbab48
commit 3fa5e8e74a
1 changed files with 208 additions and 139 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

295
DPI/train.py
View File

@ -44,39 +44,41 @@ def parse_args():
parser.add_argument('--high_noise', action='store_true') parser.add_argument('--high_noise', action='store_true')
# replay buffer # replay buffer
parser.add_argument('--replay_buffer_capacity', default=50000, type=int) #50000 parser.add_argument('--replay_buffer_capacity', default=50000, type=int) #50000
parser.add_argument('--episode_length', default=51, type=int) parser.add_argument('--episode_length', default=21, type=int)
# train # train
parser.add_argument('--agent', default='dpi', type=str, choices=['baseline', 'bisim', 'deepmdp', 'db', 'dpi', 'rpc']) parser.add_argument('--agent', default='dpi', type=str, choices=['baseline', 'bisim', 'deepmdp', 'db', 'dpi', 'rpc'])
parser.add_argument('--init_steps', default=10000, type=int) parser.add_argument('--init_steps', default=10000, type=int)
parser.add_argument('--num_train_steps', default=100000, type=int) parser.add_argument('--num_train_steps', default=100000, type=int)
parser.add_argument('--batch_size', default=50, type=int) #512 parser.add_argument('--batch_size', default=128, type=int) #512
parser.add_argument('--state_size', default=512, type=int) parser.add_argument('--state_size', default=30, type=int)
parser.add_argument('--hidden_size', default=256, type=int) parser.add_argument('--hidden_size', default=256, type=int)
parser.add_argument('--history_size', default=256, type=int) parser.add_argument('--history_size', default=128, type=int)
parser.add_argument('--episode_collection', default=5, type=int)
parser.add_argument('--episodes_buffer', default=20, type=int)
parser.add_argument('--num-units', type=int, default=50, help='num hidden units for reward/value/discount models') parser.add_argument('--num-units', type=int, default=50, help='num hidden units for reward/value/discount models')
parser.add_argument('--load_encoder', default=None, type=str) parser.add_argument('--load_encoder', default=None, type=str)
parser.add_argument('--imagine_horizon', default=15, type=str) parser.add_argument('--imagine_horizon', default=10, type=str)
parser.add_argument('--grad_clip_norm', type=float, default=100.0, help='Gradient clipping norm') parser.add_argument('--grad_clip_norm', type=float, default=100.0, help='Gradient clipping norm')
# eval # eval
parser.add_argument('--eval_freq', default=10, type=int) # TODO: master had 10000 parser.add_argument('--eval_freq', default=10, type=int) # TODO: master had 10000
parser.add_argument('--num_eval_episodes', default=20, type=int) parser.add_argument('--num_eval_episodes', default=20, type=int)
parser.add_argument('--evaluation_interval', default=10000, type=int) # TODO: master had 10000 parser.add_argument('--evaluation_interval', default=10000, type=int) # TODO: master had 10000
# value # value
parser.add_argument('--value_lr', default=1e-3, type=float) parser.add_argument('--value_lr', default=8e-6, type=float)
parser.add_argument('--value_beta', default=0.9, type=float) parser.add_argument('--value_beta', default=0.9, type=float)
parser.add_argument('--value_tau', default=0.005, type=float) parser.add_argument('--value_tau', default=0.005, type=float)
parser.add_argument('--value_target_update_freq', default=100, type=int) parser.add_argument('--value_target_update_freq', default=100, type=int)
parser.add_argument('--td_lambda', default=0.95, type=int) parser.add_argument('--td_lambda', default=0.95, type=int)
# actor # actor
parser.add_argument('--actor_lr', default=1e-3, type=float) parser.add_argument('--actor_lr', default=8e-6, type=float)
parser.add_argument('--actor_beta', default=0.9, 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_min', default=-10, type=float)
parser.add_argument('--actor_log_std_max', default=2, type=float) parser.add_argument('--actor_log_std_max', default=2, type=float)
parser.add_argument('--actor_update_freq', default=2, type=int) parser.add_argument('--actor_update_freq', default=2, type=int)
# world/encoder/decoder # world/encoder/decoder
parser.add_argument('--encoder_type', default='pixel', type=str, choices=['pixel', 'pixelCarla096', 'pixelCarla098', 'identity']) parser.add_argument('--encoder_type', default='pixel', type=str, choices=['pixel', 'pixelCarla096', 'pixelCarla098', 'identity'])
parser.add_argument('--world_model_lr', default=1e-3, type=float) parser.add_argument('--world_model_lr', default=1e-5, type=float)
parser.add_argument('--encoder_tau', default=0.001, type=float) parser.add_argument('--encoder_tau', default=0.005, type=float)
parser.add_argument('--decoder_type', default='pixel', type=str, choices=['pixel', 'identity', 'contrastive', 'reward', 'inverse', 'reconstruction']) parser.add_argument('--decoder_type', default='pixel', type=str, choices=['pixel', 'identity', 'contrastive', 'reward', 'inverse', 'reconstruction'])
parser.add_argument('--num_layers', default=4, type=int) parser.add_argument('--num_layers', default=4, type=int)
parser.add_argument('--num_filters', default=32, type=int) parser.add_argument('--num_filters', default=32, type=int)
@ -117,6 +119,7 @@ class DPI:
self.env = make_env(self.args) self.env = make_env(self.args)
#self.args.seed = np.random.randint(0, 1000) #self.args.seed = np.random.randint(0, 1000)
self.env.seed(self.args.seed) self.env.seed(self.args.seed)
self.global_episodes = 0
# noiseless environment setup # noiseless environment setup
self.args.version = 2 # env_id changes to v2 self.args.version = 2 # env_id changes to v2
@ -160,7 +163,7 @@ class DPI:
self.obs_decoder = ObservationDecoder( self.obs_decoder = ObservationDecoder(
state_size=self.args.state_size, # 128 state_size=self.args.state_size, # 128
output_shape=(self.args.channels,self.args.image_size,self.args.image_size) # (3,84,84) output_shape=(self.args.channels*self.args.channels,self.args.image_size,self.args.image_size) # (3,84,84)
).to(device) ).to(device)
self.transition_model = TransitionModel( self.transition_model = TransitionModel(
@ -176,6 +179,8 @@ class DPI:
hidden_size=self.args.hidden_size, # 256, hidden_size=self.args.hidden_size, # 256,
action_size=self.env.action_space.shape[0], # 6 action_size=self.env.action_space.shape[0], # 6
).to(device) ).to(device)
self.actor_model.apply(self.init_weights)
# Value Models # Value Models
self.value_model = ValueModel( self.value_model = ValueModel(
@ -210,22 +215,32 @@ class DPI:
hidden_size=self.args.hidden_size, # 256 hidden_size=self.args.hidden_size, # 256
).to(device) ).to(device)
self.club_sample = CLUBSample(
x_dim=self.args.state_size, # 128
y_dim=self.args.state_size, # 128
hidden_size=self.args.hidden_size, # 256
).to(device)
# model parameters # model parameters
self.world_model_parameters = list(self.obs_encoder.parameters()) + list(self.prjoection_head.parameters()) + \ self.world_model_parameters = list(self.obs_encoder.parameters()) + list(self.prjoection_head.parameters()) + \
list(self.transition_model.parameters()) + list(self.obs_decoder.parameters()) + \ list(self.transition_model.parameters()) + list(self.obs_decoder.parameters()) + \
list(self.reward_model.parameters()) list(self.reward_model.parameters()) + list(self.club_sample.parameters())
self.past_transition_parameters = self.transition_model.parameters() self.past_transition_parameters = self.transition_model.parameters()
# optimizers # optimizers
self.world_model_opt = torch.optim.Adam(self.world_model_parameters, self.args.world_model_lr) self.world_model_opt = torch.optim.Adam(self.world_model_parameters, self.args.world_model_lr)
self.value_opt = torch.optim.Adam(self.value_model.parameters(), self.args.value_lr) self.value_opt = torch.optim.Adam(self.value_model.parameters(), self.args.value_lr)
self.actor_opt = torch.optim.Adam(self.actor_model.parameters(), self.args.actor_lr) self.actor_opt = torch.optim.Adam(self.actor_model.parameters(), self.args.actor_lr)
#self.reward_opt = torch.optim.Adam(self.reward_model.parameters(), 1e-5)
#self.decoder_opt = torch.optim.Adam(self.obs_decoder.parameters(), 1e-4)
# Create Modules # Create Modules
self.world_model_modules = [self.obs_encoder, self.prjoection_head, self.transition_model, self.obs_decoder, self.reward_model] self.world_model_modules = [self.obs_encoder, self.prjoection_head, self.transition_model, self.obs_decoder, self.reward_model, self.club_sample]
self.value_modules = [self.value_model] self.value_modules = [self.value_model]
self.actor_modules = [self.actor_model] self.actor_modules = [self.actor_model]
#self.reward_modules = [self.reward_model]
#self.decoder_modules = [self.obs_decoder]
if use_saved: if use_saved:
self._use_saved_models(saved_model_dir) self._use_saved_models(saved_model_dir)
@ -240,63 +255,85 @@ class DPI:
done = False done = False
all_rews = [] all_rews = []
#video = VideoRecorder(self.video_dir if args.save_video else None, resource_files=args.resource_files)
for episode_count in tqdm.tqdm(range(episodes), desc='Collecting episodes'): for episode_count in tqdm.tqdm(range(episodes), desc='Collecting episodes'):
#if args.save_video: self.global_episodes += 1
# self.env.video.init(enabled=True)
epi_reward = 0 epi_reward = 0
for i in range(self.args.episode_length): while not done:
if random: if random:
action = self.env.action_space.sample() action = self.env.action_space.sample()
else: else:
with torch.no_grad(): with torch.no_grad():
obs_torch = torch.unsqueeze(torch.tensor(obs).float(),0).to(device) obs = torch.tensor(obs.copy(), dtype=torch.float32).unsqueeze(0)
state = self.obs_encoder(obs_torch)["distribution"].sample() obs_processed = preprocess_obs(obs).to(device)
action = self.actor_model(state).cpu().detach().numpy().squeeze() state = self.obs_encoder(obs_processed)["distribution"].sample()
action = self.actor_model(state).cpu().numpy().squeeze()
#action = self.env.action_space.sample()
next_obs, rew, done, _ = self.env.step(action) next_obs, rew, done, _ = self.env.step(action)
self.data_buffer.add(obs, action, next_obs, rew, episode_count+1, done) self.data_buffer.add(obs, action, next_obs, rew, done, self.global_episodes)
#if args.save_video:
# self.env.video.record(self.env)
if done: #or i == self.args.episode_length-1:
obs = self.env.reset()
done=False
else:
obs = next_obs obs = next_obs
epi_reward += rew epi_reward += rew
obs = self.env.reset()
done=False
all_rews.append(epi_reward) all_rews.append(epi_reward)
#if args.save_video:
# self.env.video.save('noisy/%d.mp4' % episode_count)
#print("Collected {} random episodes".format(episode_count+1))
return all_rews return all_rews
def train(self, step, total_steps): def train(self, step, total_steps):
counter = 0 counter = 0
import matplotlib.pyplot as plt
fig, ax = plt.subplots()
while step < total_steps: while step < total_steps:
# collect experience # collect experience
if step !=0: if step !=0:
encoder = self.obs_encoder encoder = self.obs_encoder
actor = self.actor_model actor = self.actor_model
all_rews = self.collect_sequences(10, random=False, actor_model=actor, encoder_model=encoder) all_rews = self.collect_sequences(self.args.episode_collection, random=False, actor_model=actor, encoder_model=encoder)
else: else:
all_rews = self.collect_sequences(self.args.batch_size, random=True) all_rews = self.collect_sequences(self.args.episodes_buffer, random=True)
# Group by steps and sample random batch # collect sequences
#random_indices = self.data_buffer.sample_random_idx(self.args.batch_size * ((step//self.args.collection_interval)+1)) # random indices for batch non_zero_indices = np.nonzero(self.data_buffer.episode_count)[0]
#random_indices = self.data_buffer.sample_random_idx(self.data_buffer.steps//self.args.episode_length) current_obs = self.data_buffer.observations[non_zero_indices]
final_idx = self.data_buffer.group_steps(self.data_buffer, "observations").shape[1] next_obs = self.data_buffer.next_observations[non_zero_indices]
random_indices = self.data_buffer.sample_random_idx(final_idx, last=True) actions_raw = self.data_buffer.actions[non_zero_indices]
rewards = self.data_buffer.rewards[non_zero_indices]
self.terms = np.where(self.data_buffer.terminals[non_zero_indices]!=0)[0]
last_observations = self.data_buffer.group_and_sample_random_batch(self.data_buffer,"observations", "cpu", random_indices=random_indices) # Group by episodes
current_observations = self.data_buffer.group_and_sample_random_batch(self.data_buffer,"next_observations", device="cpu", random_indices=random_indices) current_obs = self.grouped_arrays(current_obs)
next_observations = self.data_buffer.group_and_sample_random_batch(self.data_buffer,"next_observations", device="cpu", offset=1, random_indices=random_indices) next_obs = self.grouped_arrays(next_obs)
actions = self.data_buffer.group_and_sample_random_batch(self.data_buffer,"actions", device=device, is_obs=False, random_indices=random_indices) actions_raw = self.grouped_arrays(actions_raw)
next_actions = self.data_buffer.group_and_sample_random_batch(self.data_buffer,"actions", device=device, is_obs=False, offset=1, random_indices=random_indices) rewards_ = self.grouped_arrays(rewards)
rewards = self.data_buffer.group_and_sample_random_batch(self.data_buffer,"rewards", device=device, is_obs=False, offset=1, random_indices=random_indices)
# Train encoder
if step == 0:
step += 1
update_steps = 1 if step > 1 else 1
#for _ in range(self.args.collection_interval // self.args.episode_length+1):
for _ in range(update_steps):
counter += 1
# Select random chunks of episodes
if current_obs.shape[0] < self.args.batch_size:
random_episode_number = np.random.randint(0, current_obs.shape[0], self.args.batch_size)
else:
random_episode_number = random.sample(range(current_obs.shape[0]), self.args.batch_size)
if current_obs[0].shape[0]-self.args.episode_length < self.args.batch_size:
init_index = np.random.randint(0, current_obs[0].shape[0]-self.args.episode_length-2, self.args.batch_size)
else:
init_index = np.asarray(random.sample(range(current_obs[0].shape[0]-self.args.episode_length), self.args.batch_size))
random.shuffle(random_episode_number)
random.shuffle(init_index)
last_observations = self.select_first_k(current_obs, init_index, random_episode_number)[:-1]
current_observations = self.select_first_k(current_obs, init_index, random_episode_number)[1:]
next_observations = self.select_first_k(next_obs, init_index, random_episode_number)[:-1]
actions = self.select_first_k(actions_raw, init_index, random_episode_number)[:-1].to(device)
next_actions = self.select_first_k(actions_raw, init_index, random_episode_number)[1:].to(device)
rewards = self.select_first_k(rewards_, init_index, random_episode_number)[1:].to(device)
# Preprocessing # Preprocessing
last_observations = preprocess_obs(last_observations).to(device) last_observations = preprocess_obs(last_observations).to(device)
@ -307,12 +344,9 @@ class DPI:
self.transition_model.init_states(self.args.batch_size, device) # (N,128) self.transition_model.init_states(self.args.batch_size, device) # (N,128)
self.history = self.transition_model.prev_history # (N,128) self.history = self.transition_model.prev_history # (N,128)
# Train encoder past_world_model_loss = 0
if step == 0: past_action_loss = 0
step += 1 past_value_loss = 0
for _ in range(1):#(self.args.collection_interval // self.args.episode_length+1):
counter += 1
past_encoder_loss = 0
for i in range(self.args.episode_length-1): for i in range(self.args.episode_length-1):
if i > 0: if i > 0:
# Encode observations and next_observations # Encode observations and next_observations
@ -323,11 +357,11 @@ class DPI:
self.next_action = next_actions[i] # (N,6) self.next_action = next_actions[i] # (N,6)
history = self.transition_model.prev_history history = self.transition_model.prev_history
# Encode negative observations # Encode negative observations fro upper bound loss
idx = torch.randperm(current_observations[i].shape[0]) # random permutation on batch idx = torch.randperm(current_observations[i].shape[0]) # random permutation on batch
random_time_index = torch.randint(0, self.args.episode_length-2, (1,)).item() # random time index random_time_index = torch.randint(0, current_observations.shape[0]-2, (1,)).item() # random time index
negative_current_observations = current_observations[random_time_index][idx] negative_current_observations = current_observations[random_time_index][idx]
self.negative_current_states_dict = self.obs_encoder(negative_current_observations) self.negative_current_states_dict = self.get_features(negative_current_observations)
# Predict current state from past state with transition model # Predict current state from past state with transition model
last_states_sample = self.last_states_dict["sample"] last_states_sample = self.last_states_dict["sample"]
@ -335,25 +369,24 @@ class DPI:
self.history = predicted_current_state_dict["history"] self.history = predicted_current_state_dict["history"]
# Calculate upper bound loss # Calculate upper bound loss
likeli_loss, ub_loss = self._upper_bound_minimization(self.last_states_dict, likeli_loss, ub_loss = self._upper_bound_minimization(self.last_states_dict["sample"].detach(),
self.current_states_dict, self.current_states_dict["sample"].detach(),
self.negative_current_states_dict, self.negative_current_states_dict["sample"].detach(),
predicted_current_state_dict predicted_current_state_dict["sample"].detach(),
) )
# Calculate encoder loss # Calculate encoder loss
encoder_loss = past_encoder_loss + self._past_encoder_loss(self.current_states_dict, predicted_current_state_dict) encoder_loss = self._past_encoder_loss(self.current_states_dict, predicted_current_state_dict)
past_encoder_loss = encoder_loss.item()
# decoder loss # decoder loss
horizon = np.minimum(self.args.imagine_horizon, self.args.episode_length-1-i) horizon = np.minimum(self.args.imagine_horizon, self.args.episode_length-1-i)
nxt_obs = next_observations[i:i+horizon].view(-1,9,84,84) nxt_obs = next_observations[i:i+horizon].reshape(-1,9,84,84)
next_states_encodings = self.get_features(nxt_obs)["sample"].view(horizon,self.args.batch_size, -1) next_states_encodings = self.get_features(nxt_obs)["sample"].view(horizon,self.args.batch_size, -1)
obs_dist = self.obs_decoder(next_states_encodings) obs_dist = self.obs_decoder(next_states_encodings)
decoder_loss = -torch.mean(obs_dist.log_prob(next_observations[i:i+horizon][:,:,:3,:,:])) decoder_loss = -torch.mean(obs_dist.log_prob(next_observations[i:i+horizon]))
# contrastive projection # contrastive projection
vec_anchor = predicted_current_state_dict["sample"] vec_anchor = predicted_current_state_dict["sample"].detach()
vec_positive = self.next_states_dict["sample"].detach() vec_positive = self.next_states_dict["sample"].detach()
z_anchor = self.prjoection_head(vec_anchor, self.action) z_anchor = self.prjoection_head(vec_anchor, self.action)
z_positive = self.prjoection_head_momentum(vec_positive, next_actions[i]).detach() z_positive = self.prjoection_head_momentum(vec_positive, next_actions[i]).detach()
@ -365,22 +398,23 @@ class DPI:
# reward loss # reward loss
reward_dist = self.reward_model(self.current_states_dict["sample"]) reward_dist = self.reward_model(self.current_states_dict["sample"])
reward_loss = -torch.mean(reward_dist.log_prob(rewards[:-1])) reward_loss = -torch.mean(reward_dist.log_prob(rewards[i]))
# world model loss # world model loss
world_model_loss = encoder_loss + ub_loss + lb_loss + reward_loss + decoder_loss world_model_loss = (10*encoder_loss + 10*ub_loss + 1e-1*lb_loss + reward_loss + 1e-3*decoder_loss + past_world_model_loss) * 1e-3
past_world_model_loss = world_model_loss.item()
# actor loss # actor loss
with FreezeParameters(self.world_model_modules): with FreezeParameters(self.world_model_modules):
imagine_horizon = self.args.imagine_horizon #np.minimum(self.args.imagine_horizon, self.args.episode_length-1-i) imagine_horizon = self.args.imagine_horizon #np.minimum(self.args.imagine_horizon, self.args.episode_length-1-i)
action = self.actor_model(self.current_states_dict["sample"]) action = self.actor_model(self.current_states_dict["sample"])
imagined_rollout = self.transition_model.imagine_rollout(self.current_states_dict["sample"].detach(), imagined_rollout = self.transition_model.imagine_rollout(self.current_states_dict["sample"],
action, self.history.detach(), action, self.history,
imagine_horizon) imagine_horizon)
with FreezeParameters(self.world_model_modules + self.value_modules): with FreezeParameters(self.world_model_modules + self.value_modules):
imag_rewards = self.reward_model(imagined_rollout["sample"]).mean imag_rewards = self.reward_model(imagined_rollout["sample"]).mean
imag_values = self.value_model(imagined_rollout["sample"]).mean imag_values = self.target_value_model(imagined_rollout["sample"]).mean
discounts = self.args.discount * torch.ones_like(imag_rewards).detach() discounts = self.args.discount * torch.ones_like(imag_rewards).detach()
@ -392,7 +426,8 @@ class DPI:
discounts = torch.cat([torch.ones_like(discounts[:1]), discounts[1:-1]], 0) discounts = torch.cat([torch.ones_like(discounts[:1]), discounts[1:-1]], 0)
self.discounts = torch.cumprod(discounts, 0).detach() self.discounts = torch.cumprod(discounts, 0).detach()
actor_loss = -torch.mean(self.discounts * self.returns) actor_loss = -torch.mean(self.discounts * self.returns) + past_action_loss
past_action_loss = actor_loss.item()
# value loss # value loss
with torch.no_grad(): with torch.no_grad():
@ -400,36 +435,17 @@ class DPI:
value_targ = self.returns.detach() value_targ = self.returns.detach()
value_dist = self.value_model(value_feat) value_dist = self.value_model(value_feat)
value_loss = -torch.mean(self.discounts * value_dist.log_prob(value_targ).unsqueeze(-1))
# update models value_loss = -torch.mean(self.discounts * value_dist.log_prob(value_targ).unsqueeze(-1)) + past_value_loss
self.world_model_opt.zero_grad() past_value_loss = value_loss.item()
self.actor_opt.zero_grad()
self.value_opt.zero_grad()
world_model_loss.backward()
actor_loss.backward()
value_loss.backward()
nn.utils.clip_grad_norm_(self.world_model_parameters, self.args.grad_clip_norm)
nn.utils.clip_grad_norm_(self.actor_model.parameters(), self.args.grad_clip_norm)
nn.utils.clip_grad_norm_(self.value_model.parameters(), self.args.grad_clip_norm)
self.world_model_opt.step()
self.actor_opt.step()
self.value_opt.step()
# update momentum encoder and projection head
soft_update_params(self.obs_encoder, self.obs_encoder_momentum, self.args.encoder_tau)
soft_update_params(self.prjoection_head, self.prjoection_head_momentum, self.args.encoder_tau)
# update target value # update target value
#if step % self.args.value_target_update_freq == 0: if step % self.args.value_target_update_freq == 0:
# self.target_value_model = copy.deepcopy(self.value_model) self.target_value_model = copy.deepcopy(self.value_model)
# counter for reward # counter for reward
count = np.arange((counter-1) * (self.args.batch_size), (counter) * (self.args.batch_size)) #count = np.arange((counter-1) * (self.args.batch_size), (counter) * (self.args.batch_size))
count = (counter-1) * (self.args.batch_size)
if step % self.args.logging_freq: if step % self.args.logging_freq:
writer.add_scalar('World Loss/World Loss', world_model_loss.detach().item(), self.data_buffer.steps) writer.add_scalar('World Loss/World Loss', world_model_loss.detach().item(), self.data_buffer.steps)
writer.add_scalar('Main Models Loss/Encoder Loss', encoder_loss.detach().item(), self.data_buffer.steps) writer.add_scalar('Main Models Loss/Encoder Loss', encoder_loss.detach().item(), self.data_buffer.steps)
@ -439,21 +455,38 @@ class DPI:
writer.add_scalar('Actor Critic Loss/Reward Loss', reward_loss.detach().item(), self.data_buffer.steps) writer.add_scalar('Actor Critic Loss/Reward Loss', reward_loss.detach().item(), self.data_buffer.steps)
writer.add_scalar('Bound Loss/Upper Bound Loss', ub_loss.detach().item(), self.data_buffer.steps) writer.add_scalar('Bound Loss/Upper Bound Loss', ub_loss.detach().item(), self.data_buffer.steps)
writer.add_scalar('Bound Loss/Lower Bound Loss', lb_loss.detach().item(), self.data_buffer.steps) writer.add_scalar('Bound Loss/Lower Bound Loss', lb_loss.detach().item(), self.data_buffer.steps)
step += 1 step += 1
print(world_model_loss, actor_loss, value_loss)
# save model # update actor model
#if step % 500 == 0:#self.args.saving_interval == 0: self.actor_opt.zero_grad()
# print("Saving model") actor_loss.backward()
# path = os.path.dirname(os.path.realpath(__file__)) + "/saved_models/models.pth" nn.utils.clip_grad_norm_(self.actor_model.parameters(), self.args.grad_clip_norm)
# self.save_models(path) self.actor_opt.step()
# update world model
self.world_model_opt.zero_grad()
world_model_loss.backward()
nn.utils.clip_grad_norm_(self.world_model_parameters, self.args.grad_clip_norm)
self.world_model_opt.step()
# update value model
self.value_opt.zero_grad()
value_loss.backward()
nn.utils.clip_grad_norm_(self.value_model.parameters(), self.args.grad_clip_norm)
self.value_opt.step()
# update momentum encoder and projection head
soft_update_params(self.obs_encoder, self.obs_encoder_momentum, self.args.encoder_tau)
soft_update_params(self.prjoection_head, self.prjoection_head_momentum, self.args.encoder_tau)
rew_len = np.arange(count, count+self.args.episode_collection) if count != 0 else np.arange(0, self.args.batch_size)
for j in range(len(all_rews)): for j in range(len(all_rews)):
writer.add_scalar('Rewards/Rewards', all_rews[j], count[j]) writer.add_scalar('Rewards/Rewards', all_rews[j], rew_len[j])
#print(self.data_buffer.steps , ((self.args.episode_length-1) * self.args.batch_size * 5)) print(step)
if self.data_buffer.steps % 5100 == 0 and self.data_buffer.steps!=0: #self.args.evaluation_interval == 0: if step % 2850 == 0 and self.data_buffer.steps!=0: #self.args.evaluation_interval == 0:
print("Saving model") print("Saving model")
path = os.path.dirname(os.path.realpath(__file__)) + "/saved_models/models.pth" path = os.path.dirname(os.path.realpath(__file__)) + "/saved_models/models.pth"
self.save_models(path) self.save_models(path)
@ -476,8 +509,9 @@ class DPI:
done = False done = False
while not done: while not done:
with torch.no_grad(): with torch.no_grad():
obs_torch = torch.unsqueeze(torch.tensor(obs).float(),0).to(device) obs = torch.tensor(obs.copy(), dtype=torch.float32).unsqueeze(0)
state = self.obs_encoder(obs_torch)["distribution"].sample() obs_processed = preprocess_obs(obs).to(device)
state = self.obs_encoder(obs_processed)["distribution"].sample()
action = self.actor_model(state).cpu().detach().numpy().squeeze() action = self.actor_model(state).cpu().detach().numpy().squeeze()
next_obs, rew, done, _ = self.env.step(action) next_obs, rew, done, _ = self.env.step(action)
@ -492,15 +526,50 @@ class DPI:
print("Episodic rewards: ", episodic_rewards) print("Episodic rewards: ", episodic_rewards)
print("Average episodic reward: ", np.mean(episodic_rewards)) print("Average episodic reward: ", np.mean(episodic_rewards))
def init_weights(self, m):
if isinstance(m, nn.Linear):
torch.nn.init.xavier_uniform_(m.weight)
m.bias.data.fill_(0.01)
def grouped_arrays(self,array):
indices = [0] + self.terms.tolist()
def subarrays():
for start, end in zip(indices[:-1], indices[1:]):
yield array[start:end]
try:
subarrays = np.stack(list(subarrays()), axis=0)
except ValueError:
subarrays = np.asarray(list(subarrays()))
return subarrays
def select_first_k(self, array, init_index, episode_number):
term_index = init_index + self.args.episode_length
array = array[episode_number]
array_list = []
for i in range(array.shape[0]):
array_list.append(array[i][init_index[i]:term_index[i]])
array = np.asarray(array_list)
if array.ndim == 5:
transposed_array = np.transpose(array, (1, 0, 2, 3, 4))
elif array.ndim == 4:
transposed_array = np.transpose(array, (1, 0, 2, 3))
elif array.ndim == 3:
transposed_array = np.transpose(array, (1, 0, 2))
elif array.ndim == 2:
transposed_array = np.transpose(array, (1, 0))
else:
transposed_array = np.expand_dims(array, axis=0)
return torch.tensor(transposed_array).float()
def _upper_bound_minimization(self, last_states, current_states, negative_current_states, predicted_current_states): def _upper_bound_minimization(self, last_states, current_states, negative_current_states, predicted_current_states):
club_sample = CLUBSample(last_states, club_loss = self.club_sample(current_states, predicted_current_states, negative_current_states)
current_states, likelihood_loss = 0
negative_current_states,
predicted_current_states)
likelihood_loss = club_sample.learning_loss()
club_loss = club_sample()
return likelihood_loss, club_loss return likelihood_loss, club_loss
def _past_encoder_loss(self, curr_states_dict, predicted_curr_states_dict): def _past_encoder_loss(self, curr_states_dict, predicted_curr_states_dict):
@ -511,7 +580,7 @@ class DPI:
predicted_curr_states_dist = predicted_curr_states_dict["distribution"] predicted_curr_states_dist = predicted_curr_states_dict["distribution"]
# KL divergence loss # KL divergence loss
loss = torch.mean(torch.distributions.kl.kl_divergence(curr_states_dist, predicted_curr_states_dist)) loss = torch.mean(torch.distributions.kl.kl_divergence(curr_states_dist,predicted_curr_states_dist))
return loss return loss
@ -573,7 +642,7 @@ if __name__ == '__main__':
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
step = 0 step = 0
total_steps = 200000 total_steps = 500000
dpi = DPI(args) dpi = DPI(args)
dpi.train(step,total_steps) dpi.train(step,total_steps)
dpi.evaluate() dpi.evaluate()