def __init__()

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

• 55 lines of code
• 16 McCabe index (conditional complexity)

    def __init__(self, prefix, sim):
        self.prefix = prefix
        self.sim = sim

        self.drivers = [
            x
            for x in self.sim.model.joint_names
            if x.startswith("{}cubelet:driver:".format(self.prefix))
        ]

        self.rotators = [
            x
            for x in self.sim.model.joint_names
            if x.startswith("{}cubelet:rot".format(self.prefix))
        ]

        self.joints = self.drivers + self.rotators

        self.joints_qpos_map = {}
        self.joints_qpos_idx = []

        for j in self.joints:
            current_index = self.sim.model.get_joint_qpos_addr(j)

            self.joints_qpos_map[j] = current_index
            self.joints_qpos_idx.append(current_index)

        # Info about the cubelets - indexed by cubelet index
        self.cubelet_meta_info = []

        # Populate cubelet meta information
        for i in range(3):
            for j in range(3):
                for k in range(3):
                    indicators = collections.OrderedDict(
                        [("x", i - 1), ("y", j - 1), ("z", k - 1)]
                    )

                    # Build cubelet nam
                    name_pieces = []
                    keys = []

                    for key, value in indicators.items():
                        if value == -1:
                            keys.append(key)
                            name_pieces.append("neg_{}".format(key))
                        elif value == 1:
                            keys.append(key)
                            name_pieces.append("pos_{}".format(key))

                    name = "_".join(name_pieces)

                    data = {"name": name, "coords": np.array([i - 1, j - 1, k - 1])}

                    if len(name_pieces) > 1:
                        # Not a driver just a normal cubelet
                        idxs = [
                            self.joints.index(
                                "{}cubelet:rot{}:{}".format(prefix, key, name)
                            )
                            for key in indicators
                        ]

                        data["type"] = "cubelet"
                        data["euler_qpos"] = [self.joints_qpos_idx[i] for i in idxs]
                    elif len(name_pieces) == 1:
                        data["driver"] = self.drivers.index(
                            "{}cubelet:driver:{}".format(prefix, name)
                        )
                        data["type"] = "driver"
                    else:
                        data["type"] = "null"

                    self.cubelet_meta_info.append(data)