fix(traj-pred-vae): global_traj axis_angle uses empty key (+3 more)

[andrewwhitecdw]

Small fixes in traj_pred/models/traj_pred_vae.py:

fix: global_traj axis_angle uses empty key causing KeyError

Fix: Replace:

            orient_key = {'axis_angle': '', '6d': '_6d', 'quat': '_q_'}[self.orient_type]

with:

            orient_key = {'axis_angle': 'orient_tp', '6d': 'orient_6d_tp', 'quat': 'orient_q_tp'}[self.orient_type]

fix: DataEncoder max pooling returns tuple instead of tensor

Fix: Replace:

        else:
            x = x.max(dim=0)

with:

        else:
            x = x.max(dim=0)[0]

fix: DataDecoder max pooling returns tuple instead of tensor

Fix: Replace:

            else:
                h = context.max(dim=0)

with:

            else:
                h = context.max(dim=0)[0]

fix: DataDecoder calls prior_mlp unconditionally when None

Fix: Replace:

            h = self.prior_mlp(h)

with:

            if self.prior_mlp is not None:
                h = self.prior_mlp(h)

Files changed

• traj_pred/models/traj_pred_vae.py