Merge pull request #38 from Co-Evolve/brittle-star/tendons

[brittle-star] Added tendon transmission

biorobot/brittle_star/README.md

@@ -165,14 +165,21 @@ Two example in-silico brittle star morphologies.
 
 Every segment has two degrees of freedom (DoF), one for in-plane motion and one for out-of-plane motion.
 
+In terms of actuation, the morphology specification can be used to select either position based control, and torque
+based control.
+The morphology specification can also be used to use tendon-based transmission (but only with torque based control). In
+this case, four tendons will be added, similar to the muscle architecture of the brittle star as shown above.
+
 In terms of sensing, the following sensors are implemented. These sensors define the base set of observations that every
 brittle star environment returns as observations (further discussed below).
 
 - Proprioception
     - Joint positions (two per segment, in-plane and out-of-plane, in radians)
     - Joint velocities (two per segment, in-plane and out-of-plane, in radians / second)
     - Joint actuator force (i.e. the total actuator force acting on a joint, in Newton meters) (two per segment)
-    - Actuator force (the scalar actuator force, in Newtons) (two per segment)
+    - Actuator force (the scalar actuator force, in Newtons) (four per segment in case of tendon transmission, otherwise two)
+    - Tendon position (in case tendon transmission is used, four per segment, in meters)
+    - Tendon velocity (in case tendon transmission is used, four per segment, in meters / second)
     - Central disk's position (w.r.t. world frame)
     - Central disk's rotation (w.r.t. world frame, in radians)
     - Central disk's velocity (w.r.t. world frame, in m/s)

biorobot/brittle_star/environment/shared/observables.py

@@ -247,6 +247,46 @@ def get_base_brittle_star_observables(
         ],
     )
 
+    # tendons
+    tendon_pos_sensors = [
+        sensor
+        for sensor in sensors
+        if sensor.type[0] == mujoco.mjtSensor.mjSENS_TENDONPOS
+    ]
+    tendon_pos_observable = observable_class(
+        name="tendon_position",
+        low=-bnp.inf * bnp.ones(len(tendon_pos_sensors)),
+        high=bnp.inf * bnp.ones(len(tendon_pos_sensors)),
+        retriever=lambda state: bnp.array(
+            [
+                get_data(state).sensordata[
+                    sensor.adr[0] : sensor.adr[0] + sensor.dim[0]
+                ]
+                for sensor in tendon_pos_sensors
+            ]
+        ).flatten(),
+    )
+
+    tendon_vel_sensors = [
+        sensor
+        for sensor in sensors
+        if sensor.type[0] == mujoco.mjtSensor.mjSENS_TENDONVEL
+    ]
+    tendon_vel_observable = observable_class(
+        name="tendon_velocity",
+        low=-bnp.inf * bnp.ones(len(tendon_vel_sensors)),
+        high=bnp.inf * bnp.ones(len(tendon_vel_sensors)),
+        retriever=lambda state: bnp.array(
+            [
+                get_data(state).sensordata[
+                    sensor.adr[0] : sensor.adr[0] + sensor.dim[0]
+                ]
+                for sensor in tendon_vel_sensors
+            ]
+        ).flatten(),
+    )
+
+    # contacts
     num_contacts, get_segment_contacts_fn = get_num_contacts_and_segment_contacts_fn(
         mj_model=mj_model, backend=backend
     )
@@ -266,5 +306,7 @@ def get_base_brittle_star_observables(
         disk_rotation_observable,
         disk_linvel_observable,
         disk_angvel_observable,
+        tendon_pos_observable,
+        tendon_vel_observable,
         segment_contact_observable,
     ]

biorobot/brittle_star/mjcf/morphology/morphology.py

@@ -73,6 +73,11 @@ def _configure_camera(self) -> None:
 
 if __name__ == "__main__":
     spec = default_brittle_star_morphology_specification(
-        num_arms=5, num_segments_per_arm=5, use_p_control=False, use_torque_control=True
+        num_arms=5,
+        num_segments_per_arm=5,
+        use_p_control=False,
+        use_torque_control=True,
+        use_tendons=True,
+        radius_to_strength_factor=200,
     )
     MJCFBrittleStarMorphology(spec).export_to_xml_with_assets("./mjcf")

biorobot/brittle_star/mjcf/morphology/parts/arm_segment.py

@@ -1,14 +1,18 @@
-from typing import Union
+from __future__ import annotations
+
+from typing import Union, List
 
 import numpy as np
 from dm_control.mjcf.element import _ElementImpl
 from moojoco.mjcf.morphology import MJCFMorphology, MJCFMorphologyPart
 
+from biorobot.brittle_star.mjcf.morphology.parts.disk import MJCFBrittleStarDisk
 from biorobot.brittle_star.mjcf.morphology.specification.specification import (
     BrittleStarJointSpecification,
     BrittleStarMorphologySpecification,
 )
 from biorobot.utils import colors
+from biorobot.utils.colors import rgba_red, rgba_tendon_contracted, rgba_tendon_relaxed
 
 
 class MJCFBrittleStarArmSegment(MJCFMorphologyPart):
@@ -23,6 +27,10 @@ def __init__(
     ) -> None:
         super().__init__(parent, name, pos, euler, *args, **kwargs)
 
+    @property
+    def parent(self) -> Union[MJCFBrittleStarDisk, MJCFBrittleStarArmSegment]:
+        return super().parent
+
     @property
     def morphology_specification(self) -> BrittleStarMorphologySpecification:
         return super().morphology_specification
@@ -41,6 +49,7 @@ def _build(self, arm_index: int, segment_index: int) -> None:
         self._build_capsule()
         self._build_connector()
         self._configure_joints()
+        self._configure_tendons()
         self._configure_actuators()
         self._configure_sensors()
 
@@ -98,16 +107,83 @@ def _configure_joint(
         return joint
 
     def _configure_joints(self) -> None:
-        self._in_plane_joint = self._configure_joint(
-            name=f"{self.base_name}_in_plane_joint",
-            axis=[0, 0, 1],
-            joint_specification=self._segment_specification.in_plane_joint_specification,
-        )
-        self._out_of_plane_joint = self._configure_joint(
-            name=f"{self.base_name}_out_of_plane_joint",
-            axis=[0, -1, 0],
-            joint_specification=self._segment_specification.out_of_plane_joint_specification,
-        )
+        self._joints = [
+            self._configure_joint(
+                name=f"{self.base_name}_in_plane_joint",
+                axis=[0, 0, 1],
+                joint_specification=self._segment_specification.in_plane_joint_specification,
+            ),
+            self._configure_joint(
+                name=f"{self.base_name}_out_of_plane_joint",
+                axis=[0, -1, 0],
+                joint_specification=self._segment_specification.out_of_plane_joint_specification,
+            ),
+        ]
+
+    def _configure_tendon_attachment_points(self) -> None:
+        angles = np.linspace(np.pi / 4, 7 * np.pi / 4, 4)
+        self._proximal_taps = []
+        self.distal_taps = []
+        for i, angle in enumerate(angles):
+            # proximal
+            pos = (
+                0.8
+                * self._segment_specification.radius.value
+                * np.array([0, np.cos(angle), np.sin(angle)])
+            )
+            pos[0] = self._segment_specification.radius.value
+            self._proximal_taps.append(
+                self.mjcf_body.add(
+                    "site",
+                    name=f"{self.base_name}_proximal_tap_{i}",
+                    type="sphere",
+                    rgba=rgba_red,
+                    pos=pos,
+                    size=[0.001],
+                )
+            )
+
+            # distal
+            pos[0] = (
+                self._segment_specification.radius.value
+                + self._segment_specification.length.value
+            )
+            self.distal_taps.append(
+                self.mjcf_body.add(
+                    "site",
+                    name=f"{self.base_name}_distal_tap_{i}",
+                    type="sphere",
+                    rgba=rgba_red,
+                    pos=pos,
+                    size=[0.001],
+                )
+            )
+
+    def _build_tendons(self) -> None:
+        if self._segment_index == 0:
+            parent: MJCFBrittleStarDisk = self.parent.parent
+            distal_taps = parent.distal_taps[self._arm_index]
+        else:
+            distal_taps = self.parent.distal_taps
+
+        self._tendons = []
+        for tendon_index, (parent_tap, segment_tap) in enumerate(
+            zip(distal_taps, self._proximal_taps)
+        ):
+            tendon = self.mjcf_model.tendon.add(
+                "spatial",
+                name=f"{self.base_name}_tendon_{tendon_index}",
+                rgba=rgba_tendon_relaxed,
+                width=self._segment_specification.radius.value * 0.1,
+            )
+            tendon.add("site", site=parent_tap)
+            tendon.add("site", site=segment_tap)
+            self._tendons.append(tendon)
+
+    def _configure_tendons(self) -> None:
+        if self.morphology_specification.actuation_specification.use_tendons.value:
+            self._configure_tendon_attachment_points()
+            self._build_tendons()
 
     def _is_first_segment(self) -> bool:
         return self._segment_index == 0
@@ -116,55 +192,79 @@ def _is_last_segment(self) -> bool:
         number_of_segments = len(self._arm_specification.segment_specifications)
         return self._segment_index == number_of_segments - 1
 
-    def _get_strength(self, joint: _ElementImpl) -> float:
+    @property
+    def _actuator_strength(self) -> float:
         strength = (
             self._segment_specification.radius.value
             * self.morphology_specification.actuation_specification.radius_to_strength_factor.value
         )
         return strength
 
-    def _configure_p_control_actuator(self, joint: _ElementImpl) -> _ElementImpl:
-        return self.mjcf_model.actuator.add(
-            "position",
-            name=f"{joint.name}_p_control",
-            joint=joint,
-            kp=50,
-            ctrllimited=True,
-            ctrlrange=joint.range,
-            forcelimited=True,
-            forcerange=[-self._get_strength(joint), self._get_strength(joint)],
-        )
+    @property
+    def _transmissions(self) -> List[_ElementImpl]:
+        if self.morphology_specification.actuation_specification.use_tendons.value:
+            return self._tendons
+        else:
+            return self._joints
+
+    def _configure_p_control_actuator(self, transmission: _ElementImpl) -> _ElementImpl:
+        actuator_attributes = {
+            "name": f"{transmission.name}_p_control",
+            "kp": 50,
+            "ctrllimited": True,
+            "ctrlrange": transmission.range,
+            "forcelimited": True,
+            "forcerange": [-self._actuator_strength, self._actuator_strength],
+            "joint": transmission,
+        }
+
+        return self.mjcf_model.actuator.add("position", **actuator_attributes)
 
     def _configure_p_control_actuators(self) -> None:
         if self.morphology_specification.actuation_specification.use_p_control.value:
-            self._in_plane_actuator = self._configure_p_control_actuator(
-                self._in_plane_joint
-            )
-            self._out_of_plane_actuator = self._configure_p_control_actuator(
-                self._out_of_plane_joint
-            )
+            self._actuators = [
+                self._configure_p_control_actuator(transmission)
+                for transmission in self._transmissions
+            ]
 
-    def _configure_torque_control_actuator(self, joint: _ElementImpl) -> _ElementImpl:
-        return self.mjcf_model.actuator.add(
-            "motor",
-            name=f"{joint.name}_torque_control",
-            joint=joint,
-            ctrllimited=True,
-            ctrlrange=[-self._get_strength(joint), self._get_strength(joint)],
-            forcelimited=True,
-            forcerange=[-self._get_strength(joint), self._get_strength(joint)],
-        )
+    def _configure_torque_control_actuator(
+        self, transmission: _ElementImpl
+    ) -> _ElementImpl:
+        actuator_attributes = {
+            "name": f"{transmission.name}_torque_control",
+            "ctrllimited": True,
+            "forcelimited": True,
+            "ctrlrange": [-self._actuator_strength, self._actuator_strength],
+            "forcerange": [-self._actuator_strength, self._actuator_strength],
+        }
+
+        if self.morphology_specification.actuation_specification.use_tendons.value:
+            actuator_attributes["tendon"] = transmission
+            actuator_attributes["ctrlrange"] = [-self._actuator_strength, 0]
+            gear = 15
+            actuator_attributes["gear"] = [gear]
+            actuator_attributes["forcerange"] = [-self._actuator_strength * gear, 0]
+        else:
+            actuator_attributes["joint"] = transmission
+            actuator_attributes["ctrlrange"] = [
+                -self._actuator_strength,
+                self._actuator_strength,
+            ]
+            actuator_attributes["forcerange"] = [
+                -self._actuator_strength,
+                self._actuator_strength,
+            ]
+
+        return self.mjcf_model.actuator.add("motor", **actuator_attributes)
 
     def _configure_torque_control_actuators(self) -> None:
         if (
             self.morphology_specification.actuation_specification.use_torque_control.value
         ):
-            self._in_plane_actuator = self._configure_torque_control_actuator(
-                self._in_plane_joint
-            )
-            self._out_of_plane_actuator = self._configure_torque_control_actuator(
-                self._out_of_plane_joint
-            )
+            self._actuators = [
+                self._configure_torque_control_actuator(transmission)
+                for transmission in self._transmissions
+            ]
 
     def _configure_actuators(self) -> None:
         self._configure_p_control_actuators()
@@ -178,34 +278,38 @@ def _configure_position_sensor(self) -> None:
             objname=self._capsule.name,
         )
 
-    def _configure_joint_sensors(self, joint: _ElementImpl) -> None:
-        self.mjcf_model.sensor.add(
-            "jointpos", joint=joint, name=f"{joint.name}_jointpos_sensor"
-        )
-        self.mjcf_model.sensor.add(
-            "jointvel", joint=joint, name=f"{joint.name}_jointvel_sensor"
-        )
-        self.mjcf_model.sensor.add(
-            "jointactuatorfrc", joint=joint, name=f"{joint.name}_actuatorfrc_sensor"
-        )
-
     def _configure_joints_sensors(self) -> None:
-        self._configure_joint_sensors(joint=self._in_plane_joint)
-        self._configure_joint_sensors(joint=self._out_of_plane_joint)
+        for joint in self._joints:
+            self.mjcf_model.sensor.add(
+                "jointpos", joint=joint, name=f"{joint.name}_jointpos_sensor"
+            )
+            self.mjcf_model.sensor.add(
+                "jointvel", joint=joint, name=f"{joint.name}_jointvel_sensor"
+            )
+            self.mjcf_model.sensor.add(
+                "jointactuatorfrc", joint=joint, name=f"{joint.name}_actuatorfrc_sensor"
+            )
 
     def _configure_actuator_sensors(self) -> None:
-        self.mjcf_model.sensor.add(
-            "actuatorfrc",
-            actuator=self._in_plane_actuator,
-            name=f"{self._in_plane_actuator.name}_actuatorfrc_sensor",
-        )
-        self.mjcf_model.sensor.add(
-            "actuatorfrc",
-            actuator=self._out_of_plane_actuator,
-            name=f"{self._out_of_plane_actuator.name}_actuatorfrc_sensor",
-        )
+        for actuator in self._actuators:
+            self.mjcf_model.sensor.add(
+                "actuatorfrc",
+                actuator=actuator,
+                name=f"{actuator.name}_actuatorfrc_sensor",
+            )
+
+    def _configure_tendon_sensors(self) -> None:
+        if self.morphology_specification.actuation_specification.use_tendons.value:
+            for tendon in self._tendons:
+                self.mjcf_model.sensor.add(
+                    "tendonpos", name=f"{tendon.name}_tendonpos_sensor", tendon=tendon
+                )
+                self.mjcf_model.sensor.add(
+                    "tendonvel", name=f"{tendon.name}_tendonvel_sensor", tendon=tendon
+                )
 
     def _configure_sensors(self) -> None:
         self._configure_position_sensor()
         self._configure_joints_sensors()
         self._configure_actuator_sensors()
+        self._configure_tendon_sensors()