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

Collecting dataset from noiseless environment

Browse Source
This commit is contained in:
Vedant Dave 2023-03-28 20:21:26 +02:00
parent 11f00ad695
commit 41dcf22262
1 changed files with 43 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
DPI/train.py
View File

@ -26,6 +26,7 @@ def parse_args():
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
# environment # environment
parser.add_argument('--domain_name', default='cheetah') 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('--task_name', default='run')
parser.add_argument('--image_size', default=84, type=int) parser.add_argument('--image_size', default=84, type=int)
parser.add_argument('--channels', default=3, type=int) parser.add_argument('--channels', default=3, type=int)
@ -113,9 +114,17 @@ class DPI:
self.env = make_env(self.args) self.env = make_env(self.args)
self.env.seed(self.args.seed) 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 # stack several consecutive frames together
if self.args.encoder_type.startswith('pixel'): if self.args.encoder_type.startswith('pixel'):
self.env = utils.FrameStack(self.env, k=self.args.frame_stack) 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 # create replay buffer
self.data_buffer = ReplayBuffer(size=self.args.replay_buffer_capacity, self.data_buffer = ReplayBuffer(size=self.args.replay_buffer_capacity,
@ -124,6 +133,12 @@ class DPI:
seq_len=self.args.episode_length, seq_len=self.args.episode_length,
batch_size=args.batch_size, batch_size=args.batch_size,
args=self.args) args=self.args)
self.data_buffer_clean = 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 # create work directory
utils.make_dir(self.args.work_dir) utils.make_dir(self.args.work_dir)
@ -145,6 +160,11 @@ class DPI:
output_shape=(self.args.frame_stack*self.args.channels,self.args.image_size,self.args.image_size) # (12,84,84) output_shape=(self.args.frame_stack*self.args.channels,self.args.image_size,self.args.image_size) # (12,84,84)
) )
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.transition_model = TransitionModel( self.transition_model = TransitionModel(
state_size=self.args.state_size, # 128 state_size=self.args.state_size, # 128
hidden_size=self.args.hidden_size, # 256 hidden_size=self.args.hidden_size, # 256
@ -153,7 +173,8 @@ class DPI:
) )
# model parameters # model parameters
self.model_parameters = list(self.obs_encoder.parameters()) + list(self.obs_decoder.parameters()) + list(self.transition_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())
# optimizer # optimizer
self.optimizer = torch.optim.Adam(self.model_parameters, lr=self.args.encoder_lr) self.optimizer = torch.optim.Adam(self.model_parameters, lr=self.args.encoder_lr)
@ -166,33 +187,45 @@ class DPI:
self.obs_decoder.load_state_dict(torch.load(os.path.join(saved_model_dir, 'obs_decoder.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'))) self.transition_model.load_state_dict(torch.load(os.path.join(saved_model_dir, 'transition_model.pt')))
def collect_random_episodes(self, episodes): def collect_sequences(self, episodes):
obs = self.env.reset() obs = self.env.reset()
obs_clean = self.env_clean.reset()
done = False done = False
#video = VideoRecorder(self.video_dir if args.save_video else None, resource_files=args.resource_files) #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'):
#self.env.video.init(enabled=True) 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): for i in range(self.args.episode_length):
action = self.env.action_space.sample() action = self.env.action_space.sample()
next_obs, _, done, _ = self.env.step(action) next_obs, _, done, _ = self.env.step(action)
next_obs_clean, _, done, _ = self.env_clean.step(action)
self.data_buffer.add(obs, action, next_obs, episode_count+1, done) self.data_buffer.add(obs, action, next_obs, episode_count+1, done)
self.data_buffer_clean.add(obs_clean, action, next_obs_clean, episode_count+1, done)
#if args.save_video: if args.save_video:
# self.env.video.record(self.env) self.env.video.record(self.env_clean)
self.env_clean.video.record(self.env_clean)
if done: if done:
obs = self.env.reset() obs = self.env.reset()
obs_clean = self.env_clean.reset()
done=False done=False
else: else:
obs = next_obs obs = next_obs
#self.env.video.save('%d.mp4' % episode_count) 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)) print("Collected {} random episodes".format(episode_count+1))
def train(self): def train(self):
# collect experience # collect experience
self.collect_random_episodes(self.args.batch_size) self.collect_sequences(self.args.batch_size)
# Group observations and next_observations by steps # Group observations and next_observations by steps
observations = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"observations")).float() observations = torch.Tensor(self.data_buffer.group_steps(self.data_buffer,"observations")).float()
@ -223,6 +256,9 @@ class DPI:
past_encoder_loss = previous_encoder_loss + self._past_encoder_loss(self.states, self.next_states, past_encoder_loss = previous_encoder_loss + self._past_encoder_loss(self.states, self.next_states,
self.states_dist, self.next_states_dist, self.states_dist, self.next_states_dist,
self.actions, self.history, i) self.actions, self.history, i)
previous_information_loss = past_latent_loss previous_information_loss = past_latent_loss
previous_encoder_loss = past_encoder_loss previous_encoder_loss = past_encoder_loss