def update_beacon_coordinates()

in envs/thor_beacons.py [0:0]

• 40 lines of code
• 17 McCabe index (conditional complexity)

    def update_beacon_coordinates(self, history):

        rotation_matrices = {} 
        def rot_matrix_cache(rot, inv):
            key = tuple(rot) + (inv,)
            if key in rotation_matrices:
                return rotation_matrices[key]
            return map_util.get_rotation_matrix_3D(rot, inv)

        beacon_R1_0 = torch.stack([map_util.get_rotation_matrix_3D(beacon['rot'], True) for beacon in self.beacons], 0)
        beacon_dp = torch.stack([torch.Tensor(beacon['pt']) - torch.Tensor(beacon['p']) for beacon in self.beacons], 0)

        beacons = [{'pt':beacon['pt'], 'target':{'objectId':beacon['target']['objectId']}, 'action':beacon['action'], 'success':beacon['success'], 't':beacon['t']} for beacon in self.beacons]
        # ^ throw away a lot of stuff we no longer need

        for t, hist in enumerate(history):

            beacon_t = copy.deepcopy(beacons)

            obj_poses = hist['obj_poses']
            update_idx = []
            R0_2 = []
            p2 = []
            for idx, beacon in enumerate(beacon_t):

                # fixed point on wall/floor
                if beacon['target']['objectId']=='unknown':
                    continue

                # object has changed state/disappeared
                if beacon['target']['objectId'] not in obj_poses or beacon['target']['objectId']==hist['inv_obj'] or beacon['target']['objectId'] not in hist['visible_objs']:
                    beacon['pt'] = None
                    continue

                pose = obj_poses[beacon['target']['objectId']] # (x, y, z) -- pos + (x, y, z) -- rot
                R0_2.append(rot_matrix_cache(pose[3:], False))
                p2.append(pose[:3])
                update_idx.append(idx)


            if len(update_idx)==0:
                hist['beacons'] = [beacon for beacon in beacon_t if beacon['pt'] is not None]
                continue

            R1_0 = beacon_R1_0[update_idx]
            dp = beacon_dp[update_idx].unsqueeze(-1)


            R0_2 = torch.stack(R0_2, 0)
            R1_2 = torch.bmm(R1_0, R0_2)
            p2 = torch.Tensor(p2)

            dp = torch.bmm(R1_2, dp)[:, :, 0] # local (x, y, z)
            new_p = dp + p2
            
            for enum_idx, t in enumerate(update_idx):
                beacon_t[t]['pt'] = tuple(new_p[enum_idx].tolist())

            hist['beacons'] =[beacon for beacon in beacon_t if beacon['pt'] is not None]

        return history