Update gradient example. Reduce sampling example duration

Author: amacati

examples/control/sampling.py

@@ -32,7 +32,7 @@
 
 # Simulation configuration
 DRONE_MODEL = "cf21B_500"
-DURATION = 18.0
+DURATION = 10.0
 FPS = 60
 RENDER = True
 

examples/jax/gradient.py

@@ -7,14 +7,18 @@
 from crazyflow.control import Control
 from crazyflow.sim import Physics, Sim
 from crazyflow.sim.data import SimData
-from crazyflow.sim.pipeline import remove_fn
 
 
 def main():
     sim = Sim(control=Control.attitude, physics=Physics.first_principles, attitude_freq=50)
     # Remove clipping floor function which kills gradients
-    remove_fn(sim.step_pipeline, "clip_floor_pos")
     sim_step = sim.build_step_fn()
+    # If the drone starts on the floor, the gradient gets killed by the floor clipping function. We
+    # thus start in the air to avoid zero gradients. Alternatively, we could also remove the floor
+    # clipping function
+    sim.data = sim.data.replace(
+        states=sim.data.states.replace(pos=sim.data.states.pos.at[..., 2].set(0.5))
+    )
 
     def step(cmd: NDArray, data: SimData) -> jax.Array:
         data = data.replace(