fix(traj-pred-vae): decode_only passes unsupported z (+3 more)

[andrewwhitecdw]

Small fixes in traj_pred/models/traj_pred_vae.py:

fix: decode_only passes unsupported z argument

Fix: Replace:

        self.data_decoder(data, mode, sample_num=sample_num, z=z)

with:

        self.data_decoder(data, mode, sample_num=sample_num)

fix: get_prior calls a DataDecoder method that does not exist

Fix: Apply patch:

--- a/traj_pred/models/traj_pred_vae.py
+++ b/traj_pred/models/traj_pred_vae.py
@@ -204,6 +204,15 @@
         return trans, orient_q
 
 
+    def get_prior(self, data):
+        context = data['context']
+        if self.pooling == 'mean':
+            h = context.mean(dim=0)
+        else:
+            h = context.max(dim=0)[0]
+        if self.prior_mlp is not None:
+            h = self.prior_mlp(h)
+        return self.p_z_net(h)
+
 """ 
 Main Model
 """

fix: joint_vel not computed when only context encoder needs it

Fix: Replace:

            if self.data_encoder.use_jvel or self.data_decoder.use_jvel:
                data['joint_vel_tp'] = (data['joint_pos_tp'][1:] - data['joint_pos_tp'][:-1]) * 30  # 30 fps
                data['joint_vel_tp'] = torch.cat([data['joint_vel_tp'][[0]], data['joint_vel_tp']], dim=0)

with:

            if self.data_encoder.use_jvel or self.data_decoder.use_jvel or self.context_encoder.use_jvel:
                data['joint_vel_tp'] = (data['joint_pos_tp'][1:] - data['joint_pos_tp'][:-1]) * 30  # 30 fps
                data['joint_vel_tp'] = torch.cat([data['joint_vel_tp'][[0]], data['joint_vel_tp']], dim=0)

fix: multi_step inference assumes 6d local orient

Fix: Replace:

            cur_data[f'{mode}_orig_out_local_traj_tp'][0, ..., 9:]  = heading_to_vec(get_heading(rot6d_to_quat(data[f'{mode}_out_local_traj_tp'][-1, ..., 3:-2])))

with:

            local_orient = data[f'{mode}_out_local_traj_tp'][-1, ..., 3:-2]
            if self.local_orient_type == '6d':
                last_q = rot6d_to_quat(local_orient)
            else:
                last_q = torch.cat([local_orient, torch.zeros_like(local_orient[..., [0]])], dim=-1)
                last_q = normalize(last_q)
            cur_data[f'{mode}_orig_out_local_traj_tp'][0, ..., 9:] = heading_to_vec(get_heading(last_q))

Files changed

• traj_pred/models/traj_pred_vae.py