Add bounding box to figure8

Author: amacati

crazyflow/envs/figure_8_env.py

@@ -112,6 +112,13 @@ def _reward(terminated: Array, pos: Array, goal: Array) -> Array:
         reward = jnp.where(terminated, -1.0, reward)
         return reward
 
+    @staticmethod
+    @jax.jit
+    def _terminated(pos: Array) -> Array:
+        hit_floor = pos[:, 0, 2] < 0.0  # Terminate if the drone has crashed into the ground
+        bounding_box = jnp.any(jnp.abs(pos[:, 0, :3]) > jnp.array([2.0, 2.0, 2.0]), axis=-1)
+        return hit_floor | bounding_box
+
     @staticmethod
     def _reset_randomization(data: SimData, mask: Array) -> SimData:
         """Randomize the initial position and velocity of the drones.

examples/figure8.py

@@ -28,7 +28,7 @@ def main():
 
     # dummy action for going up (in attitude control)
     action = np.zeros((20, 4), dtype=np.float32)
-    action[..., 0] = -0.2
+    action[..., 3] = 0.3
 
     obs, info = envs.reset()
     # Step through the environment