diff --git a/biorobot/brittle_star/README.md b/biorobot/brittle_star/README.md index 93a554a..7b2dd6d 100644 --- a/biorobot/brittle_star/README.md +++ b/biorobot/brittle_star/README.md @@ -165,6 +165,11 @@ 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). @@ -172,7 +177,9 @@ brittle star environment returns as observations (further discussed below). - 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) diff --git a/biorobot/brittle_star/environment/shared/observables.py b/biorobot/brittle_star/environment/shared/observables.py index 2d3a991..94bcf05 100644 --- a/biorobot/brittle_star/environment/shared/observables.py +++ b/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, ] diff --git a/biorobot/brittle_star/mjcf/morphology/morphology.py b/biorobot/brittle_star/mjcf/morphology/morphology.py index 92be91f..3566fac 100644 --- a/biorobot/brittle_star/mjcf/morphology/morphology.py +++ b/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") diff --git a/biorobot/brittle_star/mjcf/morphology/parts/arm_segment.py b/biorobot/brittle_star/mjcf/morphology/parts/arm_segment.py index 4a87484..751f0f5 100644 --- a/biorobot/brittle_star/mjcf/morphology/parts/arm_segment.py +++ b/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() diff --git a/biorobot/brittle_star/mjcf/morphology/parts/disk.py b/biorobot/brittle_star/mjcf/morphology/parts/disk.py index ab03761..5ffc870 100644 --- a/biorobot/brittle_star/mjcf/morphology/parts/disk.py +++ b/biorobot/brittle_star/mjcf/morphology/parts/disk.py @@ -2,11 +2,13 @@ import numpy as np from moojoco.mjcf.morphology import MJCFMorphology, MJCFMorphologyPart +from scipy.spatial.transform import Rotation as R from biorobot.brittle_star.mjcf.morphology.specification.specification import ( BrittleStarMorphologySpecification, ) from biorobot.utils import colors +from biorobot.utils.colors import rgba_red class MJCFBrittleStarDisk(MJCFMorphologyPart): @@ -30,6 +32,7 @@ def _build(self) -> None: self._build_pentagon() self._build_arm_connections() + self._configure_tendon_attachment_points() self._configure_sensors() def _build_pentagon(self) -> None: @@ -77,7 +80,6 @@ def _build_arm_connections(self) -> None: arm_angles = np.linspace(0, 2 * np.pi, 6)[:-1] - # connector_length = radius / 2 connector_length = radius * 0.2 for i, angle in enumerate(arm_angles): pos = (radius - connector_length) * np.array( @@ -89,12 +91,61 @@ def _build_arm_connections(self) -> None: name=f"{self.base_name}_arm_connector_{i}", pos=pos, euler=[0.0, 0.0, angle], - size=[connector_length, 1.1 * height, height], + size=[connector_length, height, height], rgba=colors.rgba_green, contype=0, conaffinity=0, ) + def _configure_tendon_attachment_points(self) -> None: + if self.morphology_specification.actuation_specification.use_tendons.value: + self.distal_taps = [] + + disk_radius = self.morphology_specification.disk_specification.radius.value + center_pos = np.array([disk_radius, 0, 0]) + + arm_angles = np.linspace(0, 2 * np.pi, 6)[:-1] + tap_angles = np.linspace(np.pi / 4, 7 * np.pi / 4, 4) + + for arm_index, arm_angle in enumerate(arm_angles): + arm_specification = self.morphology_specification.arm_specifications[ + arm_index + ] + if arm_specification.number_of_segments == 0: + continue + + base_segment_radius = arm_specification.segment_specifications[ + 0 + ].radius.value + + arm_taps = [] + positions = [] + for angle in tap_angles: + pos = center_pos + 0.8 * base_segment_radius * np.array( + [0, np.cos(angle), np.sin(angle)] + ) + positions.append(pos) + + for tap_index, position in enumerate(positions): + # rotate position around arm_angle degress + # Define the rotation + rotation = R.from_euler("z", arm_angle, degrees=False) + + # Rotate point A around point B + rotated_point = rotation.apply(position) + + arm_taps.append( + self.mjcf_body.add( + "site", + name=f"{self.base_name}_arm_{arm_index}_tap_{tap_index}", + type="sphere", + rgba=rgba_red, + pos=rotated_point, + size=[0.001], + ) + ) + self.distal_taps.append(arm_taps) + def _configure_sensors(self) -> None: self.mjcf_model.sensor.add( "framepos", diff --git a/biorobot/brittle_star/mjcf/morphology/specification/default.py b/biorobot/brittle_star/mjcf/morphology/specification/default.py index 5493d9f..fa5b158 100644 --- a/biorobot/brittle_star/mjcf/morphology/specification/default.py +++ b/biorobot/brittle_star/mjcf/morphology/specification/default.py @@ -74,6 +74,7 @@ def default_arm_specification(num_segments_per_arm: int) -> BrittleStarArmSpecif def default_brittle_star_morphology_specification( num_arms: int = 5, num_segments_per_arm: Union[int, List[int]] = 5, + use_tendons: bool = False, use_p_control: bool = False, use_torque_control: bool = False, radius_to_strength_factor: float = 200, @@ -96,6 +97,7 @@ def default_brittle_star_morphology_specification( arm_specifications.append(arm_specification) actuation_specification = BrittleStarActuationSpecification( + use_tendons=use_tendons, use_p_control=use_p_control, use_torque_control=use_torque_control, radius_to_strength_factor=radius_to_strength_factor, diff --git a/biorobot/brittle_star/mjcf/morphology/specification/specification.py b/biorobot/brittle_star/mjcf/morphology/specification/specification.py index 00fa6f9..7f5f64b 100644 --- a/biorobot/brittle_star/mjcf/morphology/specification/specification.py +++ b/biorobot/brittle_star/mjcf/morphology/specification/specification.py @@ -61,6 +61,7 @@ def __init__( class BrittleStarActuationSpecification(Specification): def __init__( self, + use_tendons: bool, use_p_control: bool, use_torque_control: bool, radius_to_strength_factor: float, @@ -69,7 +70,10 @@ def __init__( assert ( use_p_control + use_torque_control == 1 ), "Only one actuation method can be used." - + assert ( + not use_tendons or use_torque_control + ), "Only torque control is supported with tendons." + self.use_tendons = FixedParameter(use_tendons) self.use_p_control = FixedParameter(use_p_control) self.use_torque_control = FixedParameter(use_torque_control) self.radius_to_strength_factor = FixedParameter(radius_to_strength_factor) diff --git a/biorobot/brittle_star/usage_examples/directed_locomotion_single.py b/biorobot/brittle_star/usage_examples/directed_locomotion_single.py index c14ddf7..d0b8c8d 100644 --- a/biorobot/brittle_star/usage_examples/directed_locomotion_single.py +++ b/biorobot/brittle_star/usage_examples/directed_locomotion_single.py @@ -49,7 +49,11 @@ def create_env( backend: str, render_mode: str ) -> BrittleStarDirectedLocomotionEnvironment: morphology_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=False, ) morphology = MJCFBrittleStarMorphology(morphology_spec) arena_config = AquariumArenaConfiguration(attach_target=True) @@ -88,6 +92,7 @@ def action_sample_fn(_: None) -> Tuple[np.ndarray, None]: step_fn = jax.jit(env.step) reset_fn = jax.jit(env.reset) + reset_fn = env.reset def action_sample_fn(rng: chex.PRNGKey) -> Tuple[jnp.ndarray, chex.PRNGKey]: rng, sub_rng = jax.random.split(rng, 2) @@ -99,7 +104,7 @@ def action_sample_fn(rng: chex.PRNGKey) -> Tuple[jnp.ndarray, chex.PRNGKey]: action, action_rng = action_sample_fn(action_rng) state = step_fn(state=state, action=action) post_render(env.render(state=state), env.environment_configuration) - print(state.observations["actuator_force"]) + print(state.observations["tendon_position"]) if state.terminated | state.truncated: state = reset_fn(env_rng) env.close()