v0.2.7

biorobot/brittle_star/environment/directed_locomotion/mjx_env.py

@@ -171,9 +171,13 @@ def reset(
         qvel = jnp.zeros(mjx_model.nv)
 
         # Set morphology position
-        morphology_qpos_adr = mj_model.joint("BrittleStarMorphology/freejoint/").qposadr[0]
+        morphology_qpos_adr = mj_model.joint(
+            "BrittleStarMorphology/freejoint/"
+        ).qposadr[0]
         morphology_pos = jnp.array([0.0, 0.0, 0.11])
-        qpos = qpos.at[morphology_qpos_adr:morphology_qpos_adr + 3].set(morphology_pos)
+        qpos = qpos.at[morphology_qpos_adr : morphology_qpos_adr + 3].set(
+            morphology_pos
+        )
 
         # Add noise to initial qpos and qvel of segment joints
         joint_qpos_adrs = self._get_segment_joints_qpos_adrs(mj_model=mj_model)

biorobot/brittle_star/environment/light_escape/mjx_env.py

@@ -30,10 +30,10 @@ class BrittleStarLightEscapeMJXEnvironment(
     metadata = {"render_modes": ["human", "rgb_array"]}
 
     def __init__(
-            self,
-            mjcf_str: str,
-            mjcf_assets: Dict[str, Any],
-            configuration: BrittleStarLightEscapeEnvironmentConfiguration,
+        self,
+        mjcf_str: str,
+        mjcf_assets: Dict[str, Any],
+        configuration: BrittleStarLightEscapeEnvironmentConfiguration,
     ) -> None:
         BrittleStarLightEscapeEnvironmentBase.__init__(self)
         MJXEnv.__init__(
@@ -46,16 +46,16 @@ def __init__(
 
     @property
     def environment_configuration(
-            self,
+        self,
     ) -> BrittleStarLightEscapeEnvironmentConfiguration:
         return super(MJXEnv, self).environment_configuration
 
     @classmethod
     def from_morphology_and_arena(
-            cls,
-            morphology: MJCFBrittleStarMorphology,
-            arena: MJCFAquariumArena,
-            configuration: BrittleStarLightEscapeEnvironmentConfiguration,
+        cls,
+        morphology: MJCFBrittleStarMorphology,
+        arena: MJCFAquariumArena,
+        configuration: BrittleStarLightEscapeEnvironmentConfiguration,
     ) -> BrittleStarLightEscapeMJXEnvironment:
         assert arena.arena_configuration.sand_ground_color, (
             "This environment requires the 'sand_ground_color' "
@@ -83,13 +83,13 @@ def _get_light_of_segment(geom_id: int) -> float:
             segment_xy_position = state.mjx_data.geom_xpos[geom_id, :2]
             shifted_segment_xy_position = segment_xy_position + arena_size
             normalized_segment_xy_position = shifted_segment_xy_position / (
-                    2 * arena_size
+                2 * arena_size
             )
             normalized_segment_xy_position = normalized_segment_xy_position.at[1].set(
                 1 - normalized_segment_xy_position[1]
             )
             segment_yx_light_map_position = (
-                    normalized_segment_xy_position[::-1] * light_map_size
+                normalized_segment_xy_position[::-1] * light_map_size
             )
 
             light_map_coords = jnp.round(segment_yx_light_map_position).astype(int)
@@ -125,7 +125,7 @@ def _get_time(state: MJXEnvState) -> float:
         return state.mjx_data.time
 
     def _get_mj_models_and_datas_to_render(
-            self, state: MJXEnvState
+        self, state: MJXEnvState
     ) -> Tuple[List[mujoco.MjModel], List[mujoco.MjData]]:
         mj_models, mj_datas = super()._get_mj_models_and_datas_to_render(state=state)
         self._update_mj_models_tex_data(mj_models=mj_models, state=state)
@@ -136,11 +136,11 @@ def _get_mj_models_and_datas_to_render(
         return mj_models, mj_datas
 
     def get_renderer(
-            self,
-            identifier: int,
-            model: mujoco.MjModel,
-            data: mujoco.MjData,
-            state: MJXEnvState,
+        self,
+        identifier: int,
+        model: mujoco.MjModel,
+        data: mujoco.MjData,
+        state: MJXEnvState,
     ) -> Union[MujocoRenderer, mujoco.Renderer]:
         renderer = super().get_renderer(
             identifier=identifier, model=model, data=data, state=state
@@ -149,14 +149,14 @@ def get_renderer(
         return renderer
 
     def get_renderer_context(
-            self, renderer: Union[MujocoRenderer, mujoco.Renderer]
+        self, renderer: Union[MujocoRenderer, mujoco.Renderer]
     ) -> mujoco.MjrContext:
         return super(MJXEnv, self).get_renderer_context(renderer)
 
     def render(self, state: MJXEnvState) -> list[RenderFrame] | None:
         render_output = super().render(state=state)
         state.info["_light_map_has_changed"] = (
-                state.info["_light_map_has_changed"] * False
+            state.info["_light_map_has_changed"] * False
         )
         return render_output
 
@@ -191,7 +191,7 @@ def _update_light_map(self, state: MJXEnvState) -> MJXEnvState:
             light_map = light_map * light_noise
 
             light_map = (light_map - jnp.min(light_map)) / (
-                    jnp.max(light_map) - jnp.min(light_map)
+                jnp.max(light_map) - jnp.min(light_map)
             )
 
         info = state.info
@@ -200,11 +200,11 @@ def _update_light_map(self, state: MJXEnvState) -> MJXEnvState:
         return state.replace(rng=rng, info=info)
 
     def _finish_reset(
-            self,
-            models_and_datas: Tuple[
-                Tuple[mujoco.MjModel, mujoco.MjData], Tuple[mjx.Model, mjx.Data]
-            ],
-            rng: np.random.RandomState,
+        self,
+        models_and_datas: Tuple[
+            Tuple[mujoco.MjModel, mujoco.MjData], Tuple[mjx.Model, mjx.Data]
+        ],
+        rng: np.random.RandomState,
     ) -> MJXEnvState:
         (mj_model, mj_data), (mjx_model, mjx_data) = models_and_datas
         mjx_data = mjx.forward(m=mjx_model, d=mjx_data)
@@ -235,7 +235,9 @@ def reset(self, rng: chex.PRNGKey, *args, **kwargs) -> MJXEnvState:
         qvel = jnp.zeros(mjx_model.nv)
 
         # Set morphology position
-        morphology_qpos_adr = mj_model.joint("BrittleStarMorphology/freejoint/").qposadr[0]
+        morphology_qpos_adr = mj_model.joint(
+            "BrittleStarMorphology/freejoint/"
+        ).qposadr[0]
         morphology_pos = jnp.array(
             [
                 BrittleStarLightEscapeEnvironmentBase._get_x_start_position(
@@ -245,7 +247,9 @@ def reset(self, rng: chex.PRNGKey, *args, **kwargs) -> MJXEnvState:
                 0.11,
             ]
         )
-        qpos = qpos.at[morphology_qpos_adr:morphology_qpos_adr + 3].set(morphology_pos)
+        qpos = qpos.at[morphology_qpos_adr : morphology_qpos_adr + 3].set(
+            morphology_pos
+        )
 
         # Add noise to initial qpos and qvel of segment joints
         joint_qpos_adrs = self._get_segment_joints_qpos_adrs(mj_model=mj_model)

biorobot/brittle_star/environment/shared/observables.py

@@ -77,54 +77,6 @@ def get_segment_contacts(state: MJCEnvState) -> np.ndarray:
     return num_contacts, get_segment_contacts
 
 
-def get_actuator_frc_fn(
-    mj_model: mujoco.MjModel, backend: str
-) -> Callable[[BaseEnvState], chex.Array]:
-    # Temporary fix until https://github.com/google-deepmind/mujoco/issues/2068 is resolved
-    if backend == "mjx":
-        low_actuator_force_limit = jnp.array(
-            [limits[0] for limits in mj_model.actuator_forcerange]
-        )
-        high_actuator_force_limit = jnp.array(
-            [limits[1] for limits in mj_model.actuator_forcerange]
-        )
-
-        def calculate_actuator_force(state: MJXEnvState) -> jnp.ndarray:
-            l = state.mjx_data.actuator_length
-            v = state.mjx_data.actuator_velocity
-            gain = state.mjx_model.actuator_gainprm[:, 0]
-            b0 = state.mjx_model.actuator_biasprm[:, 0]
-            b1 = state.mjx_model.actuator_biasprm[:, 1]
-            b2 = state.mjx_model.actuator_biasprm[:, 2]
-            ctrl = state.mjx_data.ctrl
-
-            return jnp.clip(
-                gain * ctrl + b0 + b1 * l + b2 * v,
-                low_actuator_force_limit,
-                high_actuator_force_limit,
-            )
-
-        return calculate_actuator_force
-    else:
-        actuator_frc_sensors = [
-            mj_model.sensor(i)
-            for i in range(mj_model.nsensor)
-            if mj_model.sensor(i).type[0] == mujoco.mjtSensor.mjSENS_ACTUATORFRC
-        ]
-
-        def get_actuator_force(state: MJCEnvState) -> np.ndarray:
-            return np.array(
-                [
-                    state.mj_data.sensordata[
-                        sensor.adr[0] : sensor.adr[0] + sensor.dim[0]
-                    ]
-                    for sensor in actuator_frc_sensors
-                ]
-            ).flatten()
-
-        return get_actuator_force
-
-
 def get_base_brittle_star_observables(
     mj_model: mujoco.MjModel, backend: str
 ) -> List[BaseObservable]:
@@ -202,11 +154,23 @@ def get_base_brittle_star_observables(
     )
 
     # All actuator forces
+    actuator_frc_sensors = [
+        mj_model.sensor(i)
+        for i in range(mj_model.nsensor)
+        if mj_model.sensor(i).type[0] == mujoco.mjtSensor.mjSENS_ACTUATORFRC
+    ]
     actuator_force_observable = observable_class(
         name="actuator_force",
         low=bnp.array([limits[0] for limits in mj_model.actuator_forcerange]),
         high=bnp.array([limits[1] for limits in mj_model.actuator_forcerange]),
-        retriever=get_actuator_frc_fn(mj_model=mj_model, backend=backend),
+        retriever=lambda state: bnp.array(
+            [
+                get_data(state).sensordata[
+                    sensor.adr[0] : sensor.adr[0] + sensor.dim[0]
+                ]
+                for sensor in actuator_frc_sensors
+            ]
+        ).flatten(),
     )
 
     # disk pos

biorobot/brittle_star/environment/undirected_locomotion/mjx_env.py

@@ -93,9 +93,13 @@ def reset(self, rng: chex.PRNGKey, *args, **kwargs) -> MJXEnvState:
         qvel = jnp.zeros(mjx_model.nv)
 
         # Set morphology position
-        morphology_qpos_adr = mj_model.joint("BrittleStarMorphology/freejoint/").qposadr[0]
+        morphology_qpos_adr = mj_model.joint(
+            "BrittleStarMorphology/freejoint/"
+        ).qposadr[0]
         morphology_pos = jnp.array([0.0, 0.0, 0.11])
-        qpos = qpos.at[morphology_qpos_adr:morphology_qpos_adr + 3].set(morphology_pos)
+        qpos = qpos.at[morphology_qpos_adr : morphology_qpos_adr + 3].set(
+            morphology_pos
+        )
 
         # Add noise to initial qpos and qvel of segment joints
         joint_qpos_adrs = self._get_segment_joints_qpos_adrs(mj_model=mj_model)

biorobot/brittle_star/mjcf/arena/aquarium.py

@@ -121,7 +121,7 @@ def _build_ground(self) -> None:
             material=ground_material,
             rgba=rgba,
             size=list(self.arena_configuration.size) + [0.25],
-            condim=6,
+            condim=4,
             contype=0,
             conaffinity=1,
         )