Merge pull request #27 from Co-Evolve/develop

Develop
Co-Evolve · Oct 2, 2024 · 30a8294 · 30a8294
2 parents feef0f5 + 7251176
commit 30a8294
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 75 deletions.
diff --git a/biorobot/brittle_star/environment/shared/observables.py b/biorobot/brittle_star/environment/shared/observables.py
@@ -209,77 +209,41 @@ 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
+    # 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]
-            ]
-        ),
+        retriever=lambda state: bnp.array(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_quat2eueler_fn(backend=backend)(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]
-            ]
-        ),
+        retriever=lambda state: bnp.array(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]
-            ]
-        ),
+        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
@@ -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
@@ -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,8 @@ def clip_and_rescale(
     rng = jax.random.PRNGKey(0)
     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)
@@ -114,3 +118,5 @@ def clip_and_rescale(
         env.render(state=state)
 
         i += 1
+        print(state.observations["disk_position"])
+        print("-" * 20)
diff --git 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
@@ -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',