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

Outputting only 3 channels from the decoder

Browse Source
This commit is contained in:
Vedant Dave 2023-04-12 17:30:20 +02:00
parent a83149f61e
commit c8fdd11d8c
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
DPI/train.py
View File

@ -173,7 +173,7 @@ class DPI:
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) # (9,84,84)
output_shape=(self.args.channels,self.args.image_size,self.args.image_size) # (3,84,84)
)
self.transition_model = TransitionModel(
@ -362,7 +362,7 @@ class DPI:
# decoder loss
horizon = np.minimum(50-i, imagine_horizon)
obs_dist = self.obs_decoder(imagined_rollout["sample"][:horizon])
decoder_loss = -torch.mean(obs_dist.log_prob(next_observations[i:i+horizon]))
decoder_loss = -torch.mean(obs_dist.log_prob(next_observations[i:i+horizon][:,:,:3,:,:]))
# reward loss
reward_dist = self.reward_model(self.current_states_dict["sample"])