Compare commits
No commits in common. "47a0772c9de300582b82c451b2ad64a97ba91c1a" and "11f00ad695cab2fd3208c4de8ea0f5c5588530ec" have entirely different histories.
47a0772c9d
...
11f00ad695
@ -8,7 +8,6 @@ def make(
|
|||||||
resource_files,
|
resource_files,
|
||||||
img_source,
|
img_source,
|
||||||
total_frames,
|
total_frames,
|
||||||
version,
|
|
||||||
seed=1,
|
seed=1,
|
||||||
visualize_reward=True,
|
visualize_reward=True,
|
||||||
from_pixels=False,
|
from_pixels=False,
|
||||||
@ -21,7 +20,7 @@ def make(
|
|||||||
video_recording=False,
|
video_recording=False,
|
||||||
video_recording_dir=None,
|
video_recording_dir=None,
|
||||||
):
|
):
|
||||||
env_id = 'dmc_%s_%s_%s-v1' % (domain_name, task_name, version)
|
env_id = 'dmc_%s_%s_%s-v1' % (domain_name, task_name, seed)
|
||||||
|
|
||||||
if from_pixels:
|
if from_pixels:
|
||||||
assert not visualize_reward, 'cannot use visualize reward when learning from pixels'
|
assert not visualize_reward, 'cannot use visualize reward when learning from pixels'
|
||||||
|
50
DPI/train.py
50
DPI/train.py
@ -26,7 +26,6 @@ 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)
|
||||||
@ -114,17 +113,9 @@ 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,
|
||||||
@ -133,12 +124,6 @@ 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)
|
||||||
@ -160,11 +145,6 @@ 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
|
||||||
@ -173,8 +153,7 @@ class DPI:
|
|||||||
)
|
)
|
||||||
|
|
||||||
# model parameters
|
# model parameters
|
||||||
self.model_parameters = list(self.obs_encoder.parameters()) + list(self.obs_encoder_momentum.parameters()) + \
|
self.model_parameters = list(self.obs_encoder.parameters()) + list(self.obs_decoder.parameters()) + list(self.transition_model.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)
|
||||||
@ -187,45 +166,33 @@ 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_sequences(self, episodes):
|
def collect_random_episodes(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'):
|
||||||
if args.save_video:
|
#self.env.video.init(enabled=True)
|
||||||
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_clean)
|
# self.env.video.record(self.env)
|
||||||
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
|
||||||
obs_clean = next_obs_clean
|
#self.env.video.save('%d.mp4' % episode_count)
|
||||||
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_sequences(self.args.batch_size)
|
self.collect_random_episodes(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()
|
||||||
@ -256,9 +223,6 @@ 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
|
||||||
|
@ -193,7 +193,6 @@ def make_env(args):
|
|||||||
frame_skip=args.action_repeat,
|
frame_skip=args.action_repeat,
|
||||||
video_recording=args.save_video,
|
video_recording=args.save_video,
|
||||||
video_recording_dir=args.work_dir,
|
video_recording_dir=args.work_dir,
|
||||||
version=args.version,
|
|
||||||
)
|
)
|
||||||
return env
|
return env
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user