def compute()

in robogym/mujoco/forward_kinematics.py [0:0]

• 35 lines of code
• 7 McCabe index (conditional complexity)

    def compute(self, qpos, return_joint_pos=False):
        """
            given joint positions, calculate the endpoint positions.

            currently only return positions, but orientation also computed
            but just not returned.
            later could be useful for other end effectors.

            currently only support hinge joints.
        """
        num_sites = len(self.site_computations)
        site_positions = []
        joint_positions = [None] * len(self.joint_info)

        def cached_joint_calculator(computations, cidx):
            joint_idx = computations[cidx]
            assert isinstance(
                joint_idx, int
            ), "computation should be body interweaving with joints."

            if joint_positions[joint_idx] is not None:
                return joint_positions[joint_idx]
            else:
                (joint_pos, joint_axis) = self.joint_info[joint_idx]
                joint_matrix = get_joint_matrix(joint_pos, qpos[joint_idx], joint_axis)

                if cidx == len(computations) - 2:
                    joint_pos = computations[-1] @ joint_matrix
                else:
                    joint_pos = (
                        cached_joint_calculator(computations, cidx + 2)
                        @ computations[cidx + 1]
                        @ joint_matrix
                    )

                joint_positions[joint_idx] = joint_pos
                return joint_pos

        for i in range(num_sites):
            computations = self.site_computations[i]
            m = computations[0]
            if len(computations) > 1:
                m = cached_joint_calculator(computations, 1) @ m

            # only return position for now
            site_positions.append(m[:3, 3])

        if return_joint_pos:
            joint_xpos = list(map(lambda v: v[:3, 3], joint_positions))
            return np.array(site_positions + joint_xpos)
        else:
            return np.array(site_positions)