def forward(self, root, inputs):
outputs = AttrDict()
# TODO implement soft interpolation
sg_times, sg_encs = [], []
for segment in root:
sg_times.append(segment.subgoal.ind)
sg_encs.append(segment.subgoal.e_g_prime)
sg_times = torch.stack(sg_times, dim=1)
sg_encs = torch.stack(sg_encs, dim=1)
# compute time difference weights
seq_length = self._hp.max_seq_len
target_ind = torch.arange(end=seq_length, dtype=sg_times.dtype)
time_diffs = torch.abs(target_ind[None, None, :] -
sg_times[:, :, None])
weights = nn.functional.softmax(-time_diffs, dim=-1)
# compute weighted sum outputs
weighted_sg = weights[:, :, :, None, None,
None] * sg_encs.unsqueeze(2).repeat(
1, 1, seq_length, 1, 1, 1)
outputs.encodings = torch.sum(weighted_sg, dim=1)
outputs.update(
self._dense_decode(inputs, outputs.encodings, seq_length))
return outputs
def forward(self, root, inputs):
# TODO implement stopping probability prediction
# TODO make the low-level network not predict subgoals
batch_size, time = self._hp.batch_size, self._hp.max_seq_len
outputs = AttrDict()
lstm_inputs = self._get_lstm_inputs(root, inputs)
lstm_outputs = self.lstm(lstm_inputs, time)
outputs.encodings = torch.stack(lstm_outputs, dim=1)
outputs.update(self._dense_decode(inputs, outputs.encodings, time))
return outputs
