Rework camera logic

flygym/__init__.py

@@ -1,6 +1,6 @@
 from .simulation import Simulation, SingleFlySimulation
 from .fly import Fly
-from .camera import Camera, NeckCamera
+from .camera import Camera, YawOnlyCamera, ZStabilizedCamera, GravityAlignedCamera
 from .util import get_data_path, load_config
 from dm_control.rl.control import PhysicsError
 
@@ -9,3 +9,17 @@
 is_rendering_skipped = (
     "SKIP_RENDERING" in environ and environ["SKIP_RENDERING"] == "true"
 )
+
+__all__ = [
+    "Simulation",
+    "SingleFlySimulation",
+    "Fly",
+    "Camera",
+    "YawOnlyCamera",
+    "ZStabilizedCamera",
+    "GravityAlignedCamera",
+    "PhysicsError",
+    "get_data_path",
+    "load_config",
+    "is_rendering_skipped",
+]

flygym/arena/base.py

@@ -188,6 +188,19 @@ def step(self, dt: float, physics: mjcf.Physics, *args, **kwargs) -> None:
         """
         return
 
+    @abstractmethod
+    def _get_max_floor_height(self) -> float:
+        """Get the height of the floor of the arena. This is useful for
+        camera rendering. The camera should be placed at a height that is
+        slightly above the floor to avoid rendering artifacts.
+
+        Returns
+        -------
+        float
+            The height of the floor of the arena in mm.
+        """
+        pass
+
 
 class FlatTerrain(BaseArena):
     """Flat terrain with no obstacles.
@@ -247,6 +260,7 @@ def __init__(
             material=grid,
             size=ground_size,
             friction=friction,
+            conaffinity=0,
         )
         self.friction = friction
         if scale_bar_pos:
@@ -263,3 +277,12 @@ def get_spawn_position(
         self, rel_pos: np.ndarray, rel_angle: np.ndarray
     ) -> tuple[np.ndarray, np.ndarray]:
         return rel_pos, rel_angle
+
+    def _get_max_floor_height(self) -> float:
+        geom = self.root_element.find("geom", "ground")
+        try:
+            plane_height = geom.pos[2]
+        except TypeError:
+            plane_height = 0.0
+
+        return plane_height

flygym/arena/complex_terrain.py

@@ -114,6 +114,9 @@ def get_spawn_position(
         adj_pos = rel_pos + np.array([0, 0, self.gap_depth / 2])
         return adj_pos, rel_angle
 
+    def _get_max_floor_height(self):
+        return self.gap_depth / 2
+
 
 class BlocksTerrain(BaseArena):
     """Terrain formed by blocks at random heights.
@@ -185,6 +188,8 @@ def __init__(
         self.height_range = height_range
         rand_state = np.random.RandomState(rand_seed)
 
+        self.max_height = -np.inf
+
         x_centers = np.arange(x_range[0] + block_size / 2, x_range[1], block_size)
         y_centers = np.arange(y_range[0] + block_size / 2, y_range[1], block_size)
         for i, x_pos in enumerate(x_centers):
@@ -197,6 +202,8 @@ def __init__(
                 else:
                     height = 0.1 + rand_state.uniform(*height_range)
 
+                self.max_height = max(self.max_height, height)
+
                 self.root_element.worldbody.add(
                     "geom",
                     type="box",
@@ -231,6 +238,9 @@ def get_spawn_position(
         adj_pos = rel_pos + np.array([0, 0, 0.1])
         return adj_pos, rel_angle
 
+    def _get_max_floor_height(self):
+        return self.max_height
+
 
 class MixedTerrain(BaseArena):
     """A mixture of flat, blocks, and gaps terrains.
@@ -289,6 +299,8 @@ def __init__(
 
         self._height_expected_value = np.mean([*height_range])
 
+        self._max_block_height = -np.inf
+
         # 3 repetitions, each consisting of a block part, 2 gaps, and a flat part
         for x_range in [(-4, 5), (5, 14), (14, 23)]:
             # block part
@@ -316,6 +328,10 @@ def __init__(
                         y_pos,
                         height / 2 - block_size / 2 - self._height_expected_value - 0.1,
                     )
+                    self._max_block_height = max(
+                        self._max_block_height,
+                        height - self._height_expected_value - 0.1,
+                    )
                     self.root_element.worldbody.add(
                         "geom",
                         type="box",
@@ -391,3 +407,7 @@ def get_spawn_position(
     ) -> tuple[np.ndarray, np.ndarray]:
         adj_pos = rel_pos + np.array([0, 0, -1 * self._height_expected_value])
         return adj_pos, rel_angle
+
+    def _get_max_floor_height(self):
+        # The floor and gap tops are at z=0
+        return max(0, self._max_block_height)

flygym/arena/sensory_environment.py

@@ -34,11 +34,6 @@ class OdorArena(BaseArena):
         The function that, given a distance from the odor source, returns
         the relative intensity of the odor. By default, this is an inverse
         square relationship.
-    birdeye_cam : dm_control.mujoco.Camera
-        MuJoCo camera that gives a birdeye view of the arena.
-    birdeye_cam_zoom : dm_control.mujoco.Camera
-         MuJoCo camera that gives a birdeye view of the arena, zoomed in
-         toward the fly.
 
     Parameters
     ----------
@@ -119,24 +114,6 @@ def __init__(
             )
         self.diffuse_func = diffuse_func
 
-        # Add birdeye camera
-        self.birdeye_cam = self.root_element.worldbody.add(
-            "camera",
-            name="birdeye_cam",
-            mode="fixed",
-            pos=(self.odor_source[:, 0].max() / 2, 0, 35),
-            euler=(0, 0, 0),
-            fovy=45,
-        )
-        self.birdeye_cam_zoom = self.root_element.worldbody.add(
-            "camera",
-            name="birdeye_cam_zoom",
-            mode="fixed",
-            pos=(11, 0, 29),
-            euler=(0, 0, 0),
-            fovy=45,
-        )
-
         # Add markers at the odor sources
         if marker_colors is None:
             color_cycle_rgb = load_config()["color_cycle_rgb"]
@@ -209,3 +186,12 @@ def get_olfaction(self, antennae_pos: np.ndarray) -> np.ndarray:
     @property
     def odor_dimensions(self) -> int:
         return self.peak_odor_intensity.shape[1]
+
+    def _get_max_floor_height(self) -> float:
+        geom = self.root_element.find("geom", "ground")
+        try:
+            plane_height = geom.pos[2]
+        except TypeError:
+            plane_height = 0.0
+
+        return plane_height

flygym/arena/tethered.py

@@ -47,6 +47,9 @@ def spawn_entity(
             "joint", name="prismatic_support_1", limited=True, range=(0, 1e-10)
         )
 
+    def _get_max_floor_height(self) -> float:
+        raise NotImplementedError
+
 
 class Ball(Tethered):
     """Fly tethered on a spherical treadmill.
@@ -126,3 +129,6 @@ def __init__(
             "joint", name="treadmill_joint", type="ball", limited="false"
         )
         treadmill_body.add("inertial", pos=[0, 0, 0], mass=mass)
+
+    def _get_max_floor_height(self) -> float:
+        raise NotImplementedError