import numpy as np import torch import argparse import os import gym import time import json import dmc2gym import tqdm import wandb import utils 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 #from agents.navigation.carla_env import CarlaEnv def parse_args(): parser = argparse.ArgumentParser() # environment parser.add_argument('--domain_name', default='cheetah') parser.add_argument('--version', default=1, type=int) parser.add_argument('--task_name', default='run') parser.add_argument('--image_size', default=84, type=int) parser.add_argument('--channels', default=3, type=int) parser.add_argument('--action_repeat', default=1, type=int) parser.add_argument('--frame_stack', default=3, type=int) parser.add_argument('--resource_files', type=str) parser.add_argument('--eval_resource_files', type=str) parser.add_argument('--img_source', default=None, type=str, choices=['color', 'noise', 'images', 'video', 'none']) parser.add_argument('--total_frames', default=1000, type=int) # 10000 parser.add_argument('--high_noise', action='store_true') # replay buffer parser.add_argument('--replay_buffer_capacity', default=50000, type=int) #50000 parser.add_argument('--episode_length', default=51, type=int) # train 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('--num_train_steps', default=10000, type=int) parser.add_argument('--batch_size', default=20, type=int) #512 parser.add_argument('--state_size', default=256, type=int) parser.add_argument('--hidden_size', default=128, type=int) parser.add_argument('--history_size', default=128, type=int) 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) # 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) parser.add_argument('--td_lambda', default=0.95, type=int) # reward parser.add_argument('--reward_lr', default=1e-4, type=float) # actor 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) # 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('--past_transition_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) parser.add_argument('--decoder_type', default='pixel', type=str, choices=['pixel', 'identity', 'contrastive', 'reward', 'inverse', 'reconstruction']) parser.add_argument('--decoder_lr', default=1e-3, type=float) parser.add_argument('--decoder_update_freq', default=1, type=int) 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) parser.add_argument('--alpha_lr', default=1e-3, type=float) parser.add_argument('--alpha_beta', default=0.9, type=float) # misc parser.add_argument('--seed', default=1, type=int) parser.add_argument('--logging_freq', default=100, type=int) parser.add_argument('--work_dir', default='.', type=str) parser.add_argument('--save_tb', default=False, action='store_true') parser.add_argument('--save_model', default=False, action='store_true') parser.add_argument('--save_buffer', default=False, action='store_true') parser.add_argument('--save_video', default=False, action='store_true') parser.add_argument('--transition_model_type', default='', type=str, choices=['', 'deterministic', 'probabilistic', 'ensemble']) parser.add_argument('--render', default=False, action='store_true') parser.add_argument('--port', default=2000, type=int) parser.add_argument('--num_likelihood_updates', default=5, type=int) args = parser.parse_args() return args class DPI: def __init__(self, args): # wandb config #run = wandb.init(project="dpi") self.args = args # set environment noise set_global_var(self.args.high_noise) # environment setup self.env = make_env(self.args) #self.args.seed = np.random.randint(0, 1000) self.env.seed(self.args.seed) # noiseless environment setup self.args.version = 2 # env_id changes to v2 self.args.img_source = None # no image noise self.args.resource_files = None self.env_clean = make_env(self.args) self.env_clean.seed(self.args.seed) # stack several consecutive frames together if self.args.encoder_type.startswith('pixel'): self.env = utils.FrameStack(self.env, k=self.args.frame_stack) self.env_clean = utils.FrameStack(self.env_clean, k=self.args.frame_stack) # create replay buffer self.data_buffer = ReplayBuffer(size=self.args.replay_buffer_capacity, obs_shape=(self.args.frame_stack*self.args.channels,self.args.image_size,self.args.image_size), action_size=self.env.action_space.shape[0], seq_len=self.args.episode_length, batch_size=args.batch_size, args=self.args) # create work directory utils.make_dir(self.args.work_dir) self.video_dir = utils.make_dir(os.path.join(self.args.work_dir, 'video')) self.model_dir = utils.make_dir(os.path.join(self.args.work_dir, 'model')) self.buffer_dir = utils.make_dir(os.path.join(self.args.work_dir, 'buffer')) # create models 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 ) self.obs_encoder_momentum = 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 ) self.obs_decoder = ObservationDecoder( state_size=self.args.state_size, # 128 output_shape=(self.args.frame_stack*self.args.channels,self.args.image_size,self.args.image_size) # (12,84,84) ) self.transition_model = TransitionModel( state_size=self.args.state_size, # 128 hidden_size=self.args.hidden_size, # 256 action_size=self.env.action_space.shape[0], # 6 history_size=self.args.history_size, # 128 ) # 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 ) self.target_value_model = ValueModel( state_size=self.args.state_size, # 128 hidden_size=self.args.hidden_size, # 256 ) self.reward_model = RewardModel( 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.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()) self.past_transition_parameters = self.transition_model.parameters() # 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) self.past_transition_opt = torch.optim.Adam(self.past_transition_parameters, self.args.past_transition_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) def _use_saved_models(self, saved_model_dir): self.obs_encoder.load_state_dict(torch.load(os.path.join(saved_model_dir, 'obs_encoder.pt'))) self.obs_decoder.load_state_dict(torch.load(os.path.join(saved_model_dir, 'obs_decoder.pt'))) self.transition_model.load_state_dict(torch.load(os.path.join(saved_model_dir, 'transition_model.pt'))) 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 #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'): if args.save_video: self.env.video.init(enabled=True) #self.env_clean.video.init(enabled=True) for i in range(self.args.episode_length): action = self.env.action_space.sample() next_obs, rew, done, _ = self.env.step(action) #next_obs_clean, _, done, _ = self.env_clean.step(action) self.data_buffer.add(obs, action, next_obs, rew, episode_count+1, done) #self.data_buffer_clean.add(obs_clean, action, next_obs_clean, episode_count+1, done) if args.save_video: self.env.video.record(self.env) #self.env_clean.video.record(self.env_clean) if done or i == self.args.episode_length-1: obs = self.env.reset() #obs_clean = self.env_clean.reset() done=False else: obs = next_obs #obs_clean = next_obs_clean if args.save_video: self.env.video.save('noisy/%d.mp4' % episode_count) #self.env_clean.video.save('clean/%d.mp4' % episode_count) print("Collected {} random episodes".format(episode_count+1)) def train(self): # collect experience self.collect_sequences(self.args.batch_size) # Group observations and next_observations by steps from past to present last_observations = torch.tensor(self.data_buffer.group_steps(self.data_buffer,"observations")).float()[:self.args.episode_length-1] current_observations = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"next_observations")).float()[:self.args.episode_length-1] next_observations = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"next_observations")).float()[1:] actions = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"actions",obs=False)).float()[:self.args.episode_length-1] next_actions = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"actions",obs=False)).float()[1:] rewards = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"rewards",obs=False)).float()[1:] # Initialize transition model states self.transition_model.init_states(self.args.batch_size, device="cpu") # (N,128) self.history = self.transition_model.prev_history # (N,128) # Train encoder step = 0 total_steps = 10000 metrics = {} 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) self.next_action = next_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 likeli_loss, ub_loss = self._upper_bound_minimization(self.last_states_dict, self.current_states_dict, self.negative_current_states_dict, predicted_current_state_dict ) #likeli_loss = torch.tensor(likeli_loss.numpy(),dtype=torch.float32, requires_grad=True) #ikeli_loss = likeli_loss.mean() # Calculate encoder loss encoder_loss = self._past_encoder_loss(self.current_states_dict, predicted_current_state_dict) #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() # 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 """ print(likeli_loss) for i in range(self.args.num_likelihood_updates): self.past_transition_opt.zero_grad() print(likeli_loss) likeli_loss.backward() nn.utils.clip_grad_norm_(self.past_transition_parameters, self.args.grad_clip_norm) self.past_transition_opt.step() print(encoder_loss, ub_loss, lb_loss, step) """ world_model_loss = encoder_loss + ub_loss + lb_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() """ if step % self.args.logging_freq: metrics['Upper Bound Loss'] = ub_loss.item() metrics['Encoder Loss'] = encoder_loss.item() metrics["Lower Bound Loss"] = lb_loss.item() metrics["World Model Loss"] = world_model_loss.item() wandb.log(metrics) """ # behaviour learning with FreezeParameters(self.world_model_modules): imagine_horizon = self.args.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.next_action, self.history.detach(), imagine_horizon) #print(imagined_rollout["sample"].shape, imagined_rollout["distribution"][0].sample().shape) # actor loss with FreezeParameters(self.world_model_modules + self.value_modules): imag_rew_dist = self.reward_model(imagined_rollout["sample"]) target_imag_val_dist = self.target_value_model(imagined_rollout["sample"]) imag_rews = imag_rew_dist.mean target_imag_vals = target_imag_val_dist.mean discounts = self.args.discount * torch.ones_like(imag_rews).detach() self.target_returns = self._compute_lambda_return(imag_rews[:-1], target_imag_vals[:-1], discounts[:-1] , self.args.td_lambda, target_imag_vals[-1]) discounts = torch.cat([torch.ones_like(discounts[:1]), discounts[1:-1]], 0) self.discounts = torch.cumprod(discounts, 0).detach() actor_loss = -torch.mean(self.discounts * self.target_returns) self.actor_opt.zero_grad() actor_loss.backward() nn.utils.clip_grad_norm_(self.actor_model.parameters(), self.args.grad_clip_norm) self.actor_opt.step() # value loss with torch.no_grad(): value_feat = imagined_rollout["sample"][:-1].detach() value_targ = self.target_returns.detach() value_dist = self.value_model(value_feat) value_loss = -torch.mean(self.discounts * value_dist.log_prob(value_targ).unsqueeze(-1)) 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() step += 1 if step>total_steps: print("Training finished") break #print(total_ub_loss, total_encoder_loss) def _upper_bound_minimization(self, last_states, current_states, negative_current_states, predicted_current_states): club_sample = CLUBSample(last_states, current_states, negative_current_states, predicted_current_states) likelihood_loss = club_sample.learning_loss() club_loss = club_sample() return likelihood_loss, club_loss def _past_encoder_loss(self, curr_states_dict, predicted_curr_states_dict): # current state distribution curr_states_dist = curr_states_dict["distribution"] # 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 get_features(self, x, momentum=False): 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(): if momentum: x = self.obs_encoder_momentum(x) else: x = self.obs_encoder(x) return x def _compute_lambda_return(self, rewards, values, discounts, td_lam, last_value): next_values = torch.cat([values[1:], last_value.unsqueeze(0)],0) targets = rewards + discounts * next_values * (1-td_lam) rets =[] last_rew = last_value for t in range(rewards.shape[0]-1, -1, -1): last_rew = targets[t] + discounts[t] * td_lam *(last_rew) rets.append(last_rew) returns = torch.flip(torch.stack(rets), [0]) return returns if __name__ == '__main__': args = parse_args() dpi = DPI(args) dpi.train()