def __init__()

in differentiable_robot_model/robot_model.py [0:0]

• 23 lines of code
• 6 McCabe index (conditional complexity)

    def __init__(self, urdf_path: str, name="", device=None):

        super().__init__()

        self.name = name

        self._device = (
            torch.device(device) if device is not None else torch.device("cpu")
        )

        self._urdf_model = URDFRobotModel(urdf_path=urdf_path, device=self._device)
        self._bodies = torch.nn.ModuleList()
        self._n_dofs = 0
        self._controlled_joints = []

        # here we're making the joint a part of the rigid body
        # while urdfs model joints and rigid bodies separately
        # joint is at the beginning of a link
        self._name_to_idx_map = dict()

        for (i, link) in enumerate(self._urdf_model.robot.links):

            rigid_body_params = self._urdf_model.get_body_parameters_from_urdf(i, link)

            body = DifferentiableRigidBody(
                rigid_body_params=rigid_body_params, device=self._device
            )

            body.joint_idx = None
            if rigid_body_params["joint_type"] != "fixed":
                body.joint_idx = self._n_dofs
                self._n_dofs += 1
                self._controlled_joints.append(i)

            self._bodies.append(body)
            self._name_to_idx_map[body.name] = i