diff --git a/DPI/train.py b/DPI/train.py index 0fe08a2..a23b5cb 100644 --- a/DPI/train.py +++ b/DPI/train.py @@ -10,14 +10,17 @@ import dmc2gym import tqdm import wandb import utils -from utils import ReplayBuffer, make_env, save_image -from models import ObservationEncoder, ObservationDecoder, TransitionModel, CLUBSample, Actor, ValueModel, RewardModel +from utils import ReplayBuffer, FreezeParameters, make_env, soft_update_params, save_image +from models import ObservationEncoder, ObservationDecoder, TransitionModel, Actor, ValueModel, RewardModel, ProjectionHead, ContrastiveHead, CLUBSample from logger import Logger from video import VideoRecorder from dmc2gym.wrappers import set_global_var +import torch.nn as nn +import torch.nn.functional as F import torchvision.transforms as T + #from agent.baseline_agent import BaselineAgent #from agent.bisim_agent import BisimAgent #from agent.deepmdp_agent import DeepMDPAgent @@ -53,23 +56,27 @@ def parse_args(): parser.add_argument('--num-units', type=int, default=200, help='num hidden units for reward/value/discount models') parser.add_argument('--load_encoder', default=None, type=str) parser.add_argument('--imagine_horizon', default=15, type=str) + parser.add_argument('--grad_clip_norm', type=float, default=100.0, help='Gradient clipping norm') # eval parser.add_argument('--eval_freq', default=10, type=int) # TODO: master had 10000 parser.add_argument('--num_eval_episodes', default=20, 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.005, type=float) - parser.add_argument('--critic_target_update_freq', default=2, type=int) + # value + parser.add_argument('--value_lr', default=1e-4, 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_target_update_freq', default=2, type=int) + # reward + parser.add_argument('--reward_lr', default=1e-4, type=float) # actor - parser.add_argument('--actor_lr', default=1e-3, type=float) + parser.add_argument('--actor_lr', default=1e-4, 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/decoder + # world/encoder/decoder parser.add_argument('--encoder_type', default='pixel', type=str, choices=['pixel', 'pixelCarla096', 'pixelCarla098', 'identity']) parser.add_argument('--encoder_feature_dim', default=50, type=int) + parser.add_argument('--world_model_lr', default=1e-3, type=float) parser.add_argument('--encoder_lr', default=1e-3, type=float) parser.add_argument('--encoder_tau', default=0.005, type=float) parser.add_argument('--encoder_stride', default=1, type=int) @@ -79,6 +86,7 @@ def parse_args(): parser.add_argument('--decoder_weight_lambda', default=0.0, type=float) parser.add_argument('--num_layers', default=4, type=int) parser.add_argument('--num_filters', default=32, type=int) + parser.add_argument('--aug', action='store_true') # sac parser.add_argument('--discount', default=0.99, type=float) parser.add_argument('--init_temperature', default=0.01, type=float) @@ -154,6 +162,7 @@ class DPI: self.build_models(use_saved=False, saved_model_dir=self.model_dir) def build_models(self, use_saved, saved_model_dir=None): + # World Models self.obs_encoder = ObservationEncoder( obs_shape=(self.args.frame_stack*self.args.channels,self.args.image_size,self.args.image_size), # (12,84,84) state_size=self.args.state_size # 128 @@ -176,12 +185,14 @@ class DPI: history_size=self.args.history_size, # 128 ) - self.action_model = Actor( + # Actor Model + self.actor_model = Actor( state_size=self.args.state_size, # 128 hidden_size=self.args.hidden_size, # 256, action_size=self.env.action_space.shape[0], # 6 ) + # Value Models self.value_model = ValueModel( state_size=self.args.state_size, # 128 hidden_size=self.args.hidden_size, # 256 @@ -196,13 +207,39 @@ class DPI: state_size=self.args.state_size, # 128 hidden_size=self.args.hidden_size, # 256 ) + + # Contrastive Models + self.prjoection_head = ProjectionHead( + state_size=self.args.state_size, # 128 + action_size=self.env.action_space.shape[0], # 6 + hidden_size=self.args.hidden_size, # 256 + ) + + self.prjoection_head_momentum = ProjectionHead( + state_size=self.args.state_size, # 128 + action_size=self.env.action_space.shape[0], # 6 + hidden_size=self.args.hidden_size, # 256 + ) + + self.contrastive_head = ContrastiveHead( + hidden_size=self.args.hidden_size, # 256 + ) + # model parameters - self.model_parameters = list(self.obs_encoder.parameters()) + list(self.obs_encoder_momentum.parameters()) + \ - list(self.obs_decoder.parameters()) + list(self.transition_model.parameters()) + self.world_model_parameters = list(self.obs_encoder.parameters()) + list(self.obs_decoder.parameters()) + \ + list(self.value_model.parameters()) + list(self.transition_model.parameters()) + \ + list(self.prjoection_head.parameters()) - # optimizer - self.optimizer = torch.optim.Adam(self.model_parameters, lr=self.args.encoder_lr) + # optimizers + 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.actor_opt = torch.optim.Adam(self.actor_model.parameters(), self.args.actor_lr) + + # Create Modules + self.world_model_modules = [self.obs_encoder, self.obs_decoder, self.value_model, self.transition_model, self.prjoection_head] + self.value_modules = [self.value_model] + self.actor_modules = [self.actor_model] if use_saved: self._use_saved_models(saved_model_dir) @@ -214,6 +251,8 @@ class DPI: def collect_sequences(self, episodes): obs = self.env.reset() + self.ob_mean = np.mean(obs, 0).astype(np.float32) + self.ob_std = np.std(obs, 0).mean().astype(np.float32) #obs_clean = self.env_clean.reset() done = False @@ -265,48 +304,84 @@ class DPI: self.history = self.transition_model.prev_history # (N,128) # Train encoder - total_ub_loss = 0 - total_encoder_loss = 0 - for i in range(self.args.episode_length-1): - if i > 0: - # Encode observations and next_observations - self.last_states_dict = self.obs_encoder(last_observations[i]) - self.current_states_dict = self.obs_encoder(current_observations[i]) - self.next_states_dict = self.obs_encoder_momentum(next_observations[i]) - self.action = actions[i] # (N,6) - history = self.transition_model.prev_history - - # Encode negative observations - 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 - negative_current_observations = current_observations[random_time_index][idx] - self.negative_current_states_dict = self.obs_encoder(negative_current_observations) + step = 0 + total_steps = 10000 + while step < total_steps: + for i in range(self.args.episode_length-1): + if i > 0: + # Encode observations and next_observations + self.last_states_dict = self.get_features(last_observations[i]) + self.current_states_dict = self.get_features(current_observations[i]) + self.next_states_dict = self.get_features(next_observations[i], momentum=True) + self.action = actions[i] # (N,6) + history = self.transition_model.prev_history + + # Encode negative observations + 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 + negative_current_observations = current_observations[random_time_index][idx] + self.negative_current_states_dict = self.obs_encoder(negative_current_observations) - # Predict current state from past state with transition model - last_states_sample = self.last_states_dict["sample"] - predicted_current_state_dict = self.transition_model.imagine_step(last_states_sample, self.action, self.history) - self.history = predicted_current_state_dict["history"] - - # Calculate upper bound loss - ub_loss = self._upper_bound_minimization(self.last_states_dict, - self.current_states_dict, - self.negative_current_states_dict, - predicted_current_state_dict - ) - - # Calculate encoder loss - encoder_loss = self._past_encoder_loss(self.current_states_dict, - predicted_current_state_dict) + # Predict current state from past state with transition model + last_states_sample = self.last_states_dict["sample"] + predicted_current_state_dict = self.transition_model.imagine_step(last_states_sample, self.action, self.history) + self.history = predicted_current_state_dict["history"] - total_ub_loss += ub_loss - total_encoder_loss += encoder_loss - imagine_horizon = np.minimum(self.args.imagine_horizon, self.args.episode_length-1-i) - imagined_rollout = self.transition_model.imagine_rollout(self.current_states_dict["sample"], self.action, self.history, imagine_horizon) + + # Calculate upper bound loss + ub_loss = self._upper_bound_minimization(self.last_states_dict, + self.current_states_dict, + self.negative_current_states_dict, + predicted_current_state_dict + ) + + # Calculate encoder loss + encoder_loss = self._past_encoder_loss(self.current_states_dict, + predicted_current_state_dict) - #exit() + #total_ub_loss += ub_loss + #total_encoder_loss += encoder_loss + + # contrastive projection + vec_anchor = predicted_current_state_dict["sample"] + vec_positive = self.next_states_dict["sample"].detach() + z_anchor = self.prjoection_head(vec_anchor, self.action) + z_positive = self.prjoection_head_momentum(vec_positive, next_actions[i]).detach() - #print(total_ub_loss, total_encoder_loss) + # contrastive loss + logits = self.contrastive_head(z_anchor, z_positive) + labels = labels = torch.arange(logits.shape[0]).long() + lb_loss = F.cross_entropy(logits, labels) + + + # update models + world_model_loss = encoder_loss + 1e-1 * ub_loss + lb_loss #1e-1 * ub_loss + 1e-5 * encoder_loss + 1e-1 * lb_loss + print("ub_loss: {:.4f}, encoder_loss: {:.4f}, lb_loss: {:.4f}".format(ub_loss, encoder_loss, lb_loss)) + print("world_model_loss: {:.4f}".format(world_model_loss)) + 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() + + + # behaviour learning + with FreezeParameters(self.world_model_modules): + imagine_horizon = np.minimum(self.args.imagine_horizon, self.args.episode_length-1-i) + imagined_rollout = self.transition_model.imagine_rollout(self.current_states_dict["sample"].detach(), + self.action, self.history.detach(), + imagine_horizon) + print(imagined_rollout["sample"].shape, imagined_rollout["distribution"][0].sample().shape) + #exit() + + step += 1 + + if step>total_steps: + print("Training finished") + break + #exit() + + #print(total_ub_loss, total_encoder_loss) @@ -315,7 +390,7 @@ class DPI: current_states, negative_current_states, predicted_current_states) - club_loss = club_sample() + club_loss = club_sample.loglikeli() return club_loss def _past_encoder_loss(self, curr_states_dict, predicted_curr_states_dict): @@ -325,42 +400,27 @@ class DPI: # predicted current state distribution predicted_curr_states_dist = predicted_curr_states_dict["distribution"] + + # KL divergence loss loss = torch.distributions.kl.kl_divergence(curr_states_dist, predicted_curr_states_dist).mean() return loss - - """ - def _past_encoder_loss(self, states, next_states, states_dist, next_states_dist, actions, history, step): - # Imagine next state - if step == 0: - actions = torch.zeros(self.args.batch_size, self.env.action_space.shape[0]).float() # Zero action for first step - imagined_next_states = self.transition_model.imagine_step(states, actions, history) - self.history = imagined_next_states["history"] - else: - imagined_next_states = self.transition_model.imagine_step(states, actions, self.history) # (N,128) - - # State Distribution - imagined_next_states_dist = imagined_next_states["distribution"] - - # KL divergence loss - loss = torch.distributions.kl.kl_divergence(imagined_next_states_dist, next_states_dist["distribution"]).mean() - - return loss - """ def get_features(self, x, momentum=False): - if self.aug: + import torchvision.transforms.functional as fn + x = x/255.0 - 0.5 # Preprocessing + + if self.args.aug: x = T.RandomCrop((80, 80))(x) # (None,80,80,4) x = T.functional.pad(x, (4, 4, 4, 4), "symmetric") # (None,88,88,4) x = T.RandomCrop((84, 84))(x) # (None,84,84,4) with torch.no_grad(): - x = (x.float() - self.ob_mean) / self.ob_std if momentum: - x = self.obs_encoder(x).detach() - else: x = self.obs_encoder_momentum(x) + else: + x = self.obs_encoder(x) return x if __name__ == '__main__':