def to_pycuber()

in robogym/envs/dactyl/common/cube_manipulator.py [0:0]

• 31 lines of code
• 8 McCabe index (conditional complexity)

    def to_pycuber(self) -> pycuber.Cube:
        """ Return current cubelet state as a pycuber state """
        self.soft_align_faces()
        qpos_copy = self.sim.data.qpos.copy()

        cubies = []

        for i in range(27):
            cubelet_meta = self.cubelet_meta_info[i]

            if cubelet_meta["type"] == "cubelet":
                mtx = self._cubelet_rotation_matrix(cubelet_meta, qpos_copy)

                original_coords = cubelet_meta["coords"]
                # current_coords = (mtx @ cubelet_meta['coords'].astype(float)).round().astype(int)

                cubie_desc = {}

                for prev_axis, sign in enumerate(original_coords):
                    if sign != 0:
                        vec = mtx[:, prev_axis] * sign
                        new_axis = np.abs(vec).argmax()
                        new_sign = vec[new_axis]

                        color = PYCUBER_REVERSE_COLORS[prev_axis, sign]
                        loc = PYCUBER_REVERSE_LOCATIONS[new_axis, new_sign]

                        cubie_desc[loc] = pycuber.Square(color)

                if len(cubie_desc) == 3:
                    cubies.append(pycuber.Corner(**cubie_desc))
                elif len(cubie_desc) == 2:
                    cubies.append(pycuber.Edge(**cubie_desc))
            if cubelet_meta["type"] == "driver":
                original_coords = cubelet_meta["coords"]
                axis = np.abs(original_coords).argmax()
                sign = original_coords[axis]

                color = PYCUBER_REVERSE_COLORS[axis, sign]
                loc = PYCUBER_REVERSE_LOCATIONS[axis, sign]

                cubie_desc = {loc: pycuber.Square(color)}
                cubies.append(pycuber.Centre(**cubie_desc))

        return pycuber.Cube(cubies=cubies)