def forward(self, inputs, phase='train'): """ forward pass at training time """ if not 'enc_traj_seq' in inputs: enc_traj_seq, _ = self.encoder(inputs.traj_seq[:, 0]) if self._hp.train_first_action_only \ else batch_apply(self.encoder, inputs.traj_seq) if self._hp.train_first_action_only: enc_traj_seq = enc_traj_seq[:, None] enc_traj_seq = enc_traj_seq.detach( ) if self.detach_enc else enc_traj_seq enc_goal, _ = self.encoder(inputs.I_g) n_dim = len(enc_goal.shape) fused_enc = torch.cat((enc_traj_seq, enc_goal[:, None].repeat( 1, enc_traj_seq.shape[1], *([1] * (n_dim - 1)))), dim=2) #fused_enc = torch.cat((enc_traj_seq, enc_goal[:, None].repeat(1, enc_traj_seq.shape[1], 1, 1, 1)), dim=2) if self._hp.reactive: actions_pred = batch_apply(self.policy, fused_enc) else: policy_output = self.policy(fused_enc) actions_pred = policy_output # remove last time step to match ground truth if training on full sequence actions_pred = actions_pred[:, : -1] if not self._hp.train_first_action_only else actions_pred output = AttrDict() output.actions = remove_spatial(actions_pred) if len( actions_pred.shape) > 3 else actions_pred return output
def full_seq_forward(self, inputs): if 'model_enc_seq' in inputs: enc_seq_1 = inputs.model_enc_seq[:, 1:] if self._hp.train_im0_enc and 'enc_traj_seq' in inputs: enc_seq_0 = inputs.enc_traj_seq.reshape( inputs.enc_traj_seq.shape[:2] + (self._hp.nz_enc, ))[:, :-1] enc_seq_0 = enc_seq_0[:, :enc_seq_1.shape[1]] else: enc_seq_0 = inputs.model_enc_seq[:, :-1] else: enc_seq = batch_apply(self.encoder, inputs.traj_seq) enc_seq_0, enc_seq_1 = enc_seq[:, :-1], enc_seq[:, 1:] if self.detach_enc: enc_seq_0 = enc_seq_0.detach() enc_seq_1 = enc_seq_1.detach() # TODO quite sure the concatenation is automatic actions_pred = batch_apply(self.action_pred, torch.cat([enc_seq_0, enc_seq_1], dim=2)) output = AttrDict() output.actions = actions_pred #remove_spatial(actions_pred) if 'actions' in inputs: output.action_targets = inputs.actions output.pad_mask = inputs.pad_mask return output
def apply_fn(self, inputs, fn, left_parents, right_parents): """ Recursively applies fn to the tree. :param inputs: :param fn: a function that takes in (inputs, subgoal, left_parent, right_parent) and outputs a dict :param left_parents: :param right_parents: :return: """ if self.depth == 0: return assert self.subgoals is not None # need subgoal info to match to ground truth sequence self.subgoals.update(batch_apply(fn, inputs, self.subgoals, left_parents, right_parents, unshape_inputs=True)) self.child_layer.apply_fn(rec_interleave([inputs, inputs]), fn, rec_interleave([left_parents, self.subgoals]), rec_interleave([self.subgoals, right_parents]))
def produce_tree(self, inputs, layerwise_inputs, start_inds, end_inds, left_parents, right_parents, producer, depth): """no done mask checks, assumes start_ind never None. all input tensors are of shape [batch, num_parent_nodes, ...] """ self.depth = depth if depth == 0: return # slice out inputs for this layer layer_inputs = rmap(lambda x: depthfirst2layers(reduce_dim(x, dim=1))[-depth].contiguous(), layerwise_inputs) out = batch_apply(lambda x: producer.produce_subgoal(inputs, *x, depth=depth), [layer_inputs, start_inds.float(), end_inds.float(), left_parents, right_parents]) self.subgoals, left_parents, right_parents = out self.child_layer = SubgoalTreeLayer(self) self.child_layer.produce_tree(inputs, layerwise_inputs, rec_interleave([start_inds.float(), self.subgoals.ind.clone()]), rec_interleave([self.subgoals.ind.clone(), end_inds.float()]), rec_interleave([left_parents, self.subgoals]), rec_interleave([self.subgoals, right_parents]), producer, depth - 1)