def can_get_d2b(inputs):
if 'actions' in inputs and inputs.actions.shape[2] == 3 and \
torch.equal(torch.unique(inputs.actions[:, :, 2]), like(list2ten, inputs.actions)([-1, 1])):
# Looks like sawyer
return True
else:
return False
def get_sequence(self, inputs):
# Create a sequence of indices and add zero elements
seq = like(torch.eye, inputs.reference_tensor)(self._hp.max_seq_len, dtype=torch.float32)
missing_length = (2 ** self._hp.hierarchy_levels - 1) - self._hp.max_seq_len
def merge_seqs(seq1, seq2):
mode = "shuffle"
if mode == "last":
seq = torch.cat([seq1, seq2], 0)
elif mode == "second_to_last":
seq = torch.cat([seq1[:-1], seq2, seq1[-1:]], 0)
elif mode == "shuffle":
assert (seq2 == 0).all()
np.random.seed(11)
ids = np.random.choice(self._hp.max_seq_len + missing_length, self._hp.max_seq_len, replace=False)
ids.sort()
seq = seq1.new_zeros(self._hp.max_seq_len + missing_length, seq1.shape[1])
seq[ids] = seq1
return seq
seq = merge_seqs(seq, seq.new_zeros(missing_length, self._hp.max_seq_len))
seq = depthfirst2breadthfirst(seq, dim=0)
seq = seq[None].repeat_interleave(self._hp.batch_size, 0)
return seq
def get_w(self, pad_mask, inputs, model_output, log=False):
""" Match according to the fraction """
self.apply_tree(model_output.tree, inputs)
timesteps = model_output.tree.bf.timesteps
gt_timesteps = like(torch.arange, timesteps)(pad_mask.shape[1])[None,
None]
scaled_dists = (gt_timesteps - timesteps[..., None])**2 / self.temp
if self._hp.leaf_nodes_only: # only use leaf nodes -> set distance for all other to inf
scaled_dists[:, :2**(self._hp.hierarchy_levels - 1) -
1] = float("Inf")
gt_match_dists = nn.Softmax(dim=1)(-scaled_dists)
return gt_match_dists
def forward(self, seq):
sh = list(seq.shape)
for s in seq.shape[3:]:
assert s == 1
seq = seq.view(sh[:2] + [-1])
if self.add_time:
time = like(torch.arange,
seq)(seq.shape[1])[None, :, None].repeat([sh[0], 1, 1])
seq = torch.cat([seq, time], dim=2)
proc_seq = self.run_net(seq)
proc_seq = proc_seq.view(sh[:2] + [-1] + sh[3:])
return proc_seq