From ec19cef4eca8f8536324118e4b639230e7a8c240 Mon Sep 17 00:00:00 2001 From: dmarzoug Date: Wed, 2 Oct 2024 13:56:40 +0200 Subject: [PATCH 1/3] [BrittleStar] No longer use frame sensors --- .../environment/shared/observables.py | 48 ++++--------------- .../mjcf/morphology/parts/disk.py | 41 +++------------- .../brittle_star/usage_examples/env_test.py | 6 +++ 3 files changed, 22 insertions(+), 73 deletions(-) diff --git a/biorobot/brittle_star/environment/shared/observables.py b/biorobot/brittle_star/environment/shared/observables.py index d0b7d17..535a11a 100644 --- a/biorobot/brittle_star/environment/shared/observables.py +++ b/biorobot/brittle_star/environment/shared/observables.py @@ -209,76 +209,46 @@ def get_base_brittle_star_observables( retriever=get_actuator_frc_fn(mj_model=mj_model, backend=backend), ) - # disk framepos - framepos_sensor = [ - sensor - for sensor in sensors - if sensor.type[0] == mujoco.mjtSensor.mjSENS_FRAMEPOS - ][0] + # disk pos + disk_body_id = [i for i in range(mj_model.nbody) if "central_disk" in mj_model.body(i).name][0] disk_position_observable = observable_class( name="disk_position", low=-bnp.inf * bnp.ones(3), high=bnp.inf * bnp.ones(3), retriever=lambda state: bnp.array( - get_data(state).sensordata[ - framepos_sensor.adr[0] : framepos_sensor.adr[0] + framepos_sensor.dim[0] - ] + get_data(state).xpos[disk_body_id] ), ) - # disk framequat - framequat_sensor = [ - sensor - for sensor in sensors - if sensor.type[0] == mujoco.mjtSensor.mjSENS_FRAMEQUAT - ][0] + # disk rotation disk_rotation_observable = observable_class( name="disk_rotation", low=-bnp.pi * bnp.ones(3), high=bnp.pi * bnp.ones(3), retriever=lambda state: bnp.array( get_quat2eueler_fn(backend=backend)( - get_data(state).sensordata[ - framequat_sensor.adr[0] : framequat_sensor.adr[0] - + framequat_sensor.dim[0] - ] + get_data(state).xquat[disk_body_id] ) ), ) - # disk framelinvel - framelinvel_sensor = [ - sensor - for sensor in sensors - if sensor.type[0] == mujoco.mjtSensor.mjSENS_FRAMELINVEL - ][0] + # disk com linvel disk_linvel_observable = observable_class( name="disk_linear_velocity", low=-bnp.inf * bnp.ones(3), high=bnp.inf * bnp.ones(3), retriever=lambda state: bnp.array( - get_data(state).sensordata[ - framelinvel_sensor.adr[0] : framelinvel_sensor.adr[0] - + framelinvel_sensor.dim[0] - ] + get_data(state).cvel[disk_body_id, 3:] ), ) - # disk frameangvel - frameangvel_sensor = [ - sensor - for sensor in sensors - if sensor.type[0] == mujoco.mjtSensor.mjSENS_FRAMEANGVEL - ][0] + # disk com angvel disk_angvel_observable = observable_class( name="disk_angular_velocity", low=-bnp.inf * bnp.ones(3), high=bnp.inf * bnp.ones(3), retriever=lambda state: bnp.array( - get_data(state).sensordata[ - frameangvel_sensor.adr[0] : frameangvel_sensor.adr[0] - + frameangvel_sensor.dim[0] - ] + get_data(state).cvel[disk_body_id, :3] ), ) diff --git a/biorobot/brittle_star/mjcf/morphology/parts/disk.py b/biorobot/brittle_star/mjcf/morphology/parts/disk.py index 805981a..41c189b 100644 --- a/biorobot/brittle_star/mjcf/morphology/parts/disk.py +++ b/biorobot/brittle_star/mjcf/morphology/parts/disk.py @@ -11,13 +11,13 @@ class MJCFBrittleStarDisk(MJCFMorphologyPart): def __init__( - self, - parent: Union[MJCFMorphology, MJCFMorphologyPart], - name: str, - pos: np.array, - euler: np.array, - *args, - **kwargs, + self, + parent: Union[MJCFMorphology, MJCFMorphologyPart], + name: str, + pos: np.array, + euler: np.array, + *args, + **kwargs, ) -> None: super().__init__(parent, name, pos, euler, *args, **kwargs) @@ -30,7 +30,6 @@ def _build(self) -> None: self._build_pentagon() self._build_arm_connections() - self._configure_sensors() def _build_pentagon(self) -> None: # Todo: replace this with a dynamically generated mesh @@ -93,29 +92,3 @@ def _build_arm_connections(self) -> None: contype=0, conaffinity=0, ) - - def _configure_sensors(self) -> None: - self.mjcf_model.sensor.add( - "framepos", - name=f"{self.base_name}_framepos", - objtype="body", - objname=self._name, - ) - self.mjcf_model.sensor.add( - "framequat", - name=f"{self.base_name}_framequat", - objtype="body", - objname=self._name, - ) - self.mjcf_model.sensor.add( - "framelinvel", - name=f"{self.base_name}_framelinvel", - objtype="body", - objname=self._name, - ) - self.mjcf_model.sensor.add( - "frameangvel", - name=f"{self.base_name}_frameangvel", - objtype="body", - objname=self._name, - ) diff --git a/biorobot/brittle_star/usage_examples/env_test.py b/biorobot/brittle_star/usage_examples/env_test.py index b3b1b56..b6074a1 100644 --- a/biorobot/brittle_star/usage_examples/env_test.py +++ b/biorobot/brittle_star/usage_examples/env_test.py @@ -1,6 +1,8 @@ import jax.numpy as jnp import jax.random +import mujoco from brax.v1.jumpy import ones_like +from transforms3d.euler import quat2euler from biorobot.brittle_star.environment.directed_locomotion.dual import ( BrittleStarDirectedLocomotionEnvironment, @@ -76,6 +78,7 @@ def clip_and_rescale( rng = jax.random.PRNGKey(0) state = reset_fn(rng) + print(state.observations["disk_rotation"]) i = 0 while True: rng, sub_rng = jax.random.split(rng, 2) @@ -114,3 +117,6 @@ def clip_and_rescale( env.render(state=state) i += 1 + print(state.observations["disk_linear_velocity"]) + print(state.observations["disk_angular_velocity"]) + print("-" * 20) From 8157c44441bd923a684c679f37d3f191d7256919 Mon Sep 17 00:00:00 2001 From: dmarzoug Date: Wed, 2 Oct 2024 13:58:13 +0200 Subject: [PATCH 2/3] update version --- biorobot/brittle_star/usage_examples/env_test.py | 3 +-- pyproject.toml | 2 +- setup.py | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/biorobot/brittle_star/usage_examples/env_test.py b/biorobot/brittle_star/usage_examples/env_test.py index b6074a1..bb61bc1 100644 --- a/biorobot/brittle_star/usage_examples/env_test.py +++ b/biorobot/brittle_star/usage_examples/env_test.py @@ -117,6 +117,5 @@ def clip_and_rescale( env.render(state=state) i += 1 - print(state.observations["disk_linear_velocity"]) - print(state.observations["disk_angular_velocity"]) + print(state.observations["disk_position"]) print("-" * 20) diff --git a/pyproject.toml b/pyproject.toml index aa79fa1..d907fbc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "biorobot" -version = "0.2.4" +version = "0.2.5" authors = [ { name = "Dries Marzougui", email = "dries.marzougui@gmail.com" }, ] diff --git a/setup.py b/setup.py index 45a6882..496895b 100644 --- a/setup.py +++ b/setup.py @@ -11,7 +11,7 @@ setup( name='biorobot', - version='0.2.4', + version='0.2.5', description='The Bio-inspired Robotics Testbed.', long_description=readme, url='https://github.com/Co-Evolve/brb', From 7251176342f66e88226615775ffd827d27547429 Mon Sep 17 00:00:00 2001 From: dmarzoug Date: Wed, 2 Oct 2024 14:12:17 +0200 Subject: [PATCH 3/3] formatting --- .../environment/shared/observables.py | 20 +++++++------------ .../mjcf/morphology/parts/disk.py | 14 ++++++------- .../brittle_star/usage_examples/env_test.py | 1 + 3 files changed, 15 insertions(+), 20 deletions(-) diff --git a/biorobot/brittle_star/environment/shared/observables.py b/biorobot/brittle_star/environment/shared/observables.py index 535a11a..e3c8423 100644 --- a/biorobot/brittle_star/environment/shared/observables.py +++ b/biorobot/brittle_star/environment/shared/observables.py @@ -210,14 +210,14 @@ def get_base_brittle_star_observables( ) # disk pos - disk_body_id = [i for i in range(mj_model.nbody) if "central_disk" in mj_model.body(i).name][0] + disk_body_id = [ + i for i in range(mj_model.nbody) if "central_disk" in mj_model.body(i).name + ][0] disk_position_observable = observable_class( name="disk_position", low=-bnp.inf * bnp.ones(3), high=bnp.inf * bnp.ones(3), - retriever=lambda state: bnp.array( - get_data(state).xpos[disk_body_id] - ), + retriever=lambda state: bnp.array(get_data(state).xpos[disk_body_id]), ) # disk rotation @@ -226,9 +226,7 @@ def get_base_brittle_star_observables( low=-bnp.pi * bnp.ones(3), high=bnp.pi * bnp.ones(3), retriever=lambda state: bnp.array( - get_quat2eueler_fn(backend=backend)( - get_data(state).xquat[disk_body_id] - ) + get_quat2eueler_fn(backend=backend)(get_data(state).xquat[disk_body_id]) ), ) @@ -237,9 +235,7 @@ def get_base_brittle_star_observables( name="disk_linear_velocity", low=-bnp.inf * bnp.ones(3), high=bnp.inf * bnp.ones(3), - retriever=lambda state: bnp.array( - get_data(state).cvel[disk_body_id, 3:] - ), + retriever=lambda state: bnp.array(get_data(state).cvel[disk_body_id, 3:]), ) # disk com angvel @@ -247,9 +243,7 @@ def get_base_brittle_star_observables( name="disk_angular_velocity", low=-bnp.inf * bnp.ones(3), high=bnp.inf * bnp.ones(3), - retriever=lambda state: bnp.array( - get_data(state).cvel[disk_body_id, :3] - ), + retriever=lambda state: bnp.array(get_data(state).cvel[disk_body_id, :3]), ) num_contacts, get_segment_contacts_fn = get_num_contacts_and_segment_contacts_fn( diff --git a/biorobot/brittle_star/mjcf/morphology/parts/disk.py b/biorobot/brittle_star/mjcf/morphology/parts/disk.py index 41c189b..d955735 100644 --- a/biorobot/brittle_star/mjcf/morphology/parts/disk.py +++ b/biorobot/brittle_star/mjcf/morphology/parts/disk.py @@ -11,13 +11,13 @@ class MJCFBrittleStarDisk(MJCFMorphologyPart): def __init__( - self, - parent: Union[MJCFMorphology, MJCFMorphologyPart], - name: str, - pos: np.array, - euler: np.array, - *args, - **kwargs, + self, + parent: Union[MJCFMorphology, MJCFMorphologyPart], + name: str, + pos: np.array, + euler: np.array, + *args, + **kwargs, ) -> None: super().__init__(parent, name, pos, euler, *args, **kwargs) diff --git a/biorobot/brittle_star/usage_examples/env_test.py b/biorobot/brittle_star/usage_examples/env_test.py index bb61bc1..3091444 100644 --- a/biorobot/brittle_star/usage_examples/env_test.py +++ b/biorobot/brittle_star/usage_examples/env_test.py @@ -79,6 +79,7 @@ def clip_and_rescale( state = reset_fn(rng) print(state.observations["disk_rotation"]) + print(state.observations["disk_position"]) i = 0 while True: rng, sub_rng = jax.random.split(rng, 2)