Fix example scripts

Author: amacati

examples/cache.py

@@ -21,6 +21,7 @@
 or when the cache directory is deleted.
 """
 
+import shutil
 import time
 from pathlib import Path
 
@@ -32,7 +33,7 @@ def main():
     cache_dir = Path("/tmp/jax_cache_test")
     if use_cache := cache_dir.exists():
         print("Cache directory exists. This run will be fast.")
-        print(f"\nTo run without cache, delete the directory {cache_dir}.")
+        print("\nTo run without cache, run this script again.")
     else:
         print("Cache directory does not exist. This run will be slow.")
         print("\nTo run with cache, run this script again.")
@@ -44,6 +45,9 @@ def main():
     t2 = time.perf_counter()
     prefix = "Using cache: " if use_cache else "Not using cache: "
     print(f"{prefix}\n Init: {t1 - t0:.3f}s\n Step: {t2 - t1:.3f}s")
+    if use_cache:
+        shutil.rmtree(cache_dir)  # Clean up cache so that the next run is slow again
+    sim.close()
 
 
 if __name__ == "__main__":

examples/contacts.py

@@ -9,7 +9,8 @@ def main():
     sim = Sim(n_worlds=n_worlds, n_drones=n_drones, physics=Physics.so_rpy, device="cpu")
     fps = 60
 
-    cmd = np.array([[[0.3, 0, 0, 0] for _ in range(n_drones)] for _ in range(n_worlds)])
+    cmd = np.zeros((sim.n_worlds, sim.n_drones, 4))
+    cmd[..., 3] = sim.data.params.mass[0, 0, 0] * 9.81 * 1.04
     for i in range(int(2 * sim.control_freq)):
         sim.attitude_control(cmd)
         sim.step(sim.freq // sim.control_freq)

examples/crash.py

@@ -35,6 +35,7 @@ def main():
         if ((i * fps) % sim.control_freq) < fps:
             sim.render()
             print(f"Crash detected: {sim.contacts().any()}")
+    sim.close()
 
 
 if __name__ == "__main__":

examples/gradient.py

@@ -11,7 +11,9 @@
 
 def main():
     sim = Sim(control=Control.attitude, physics=Physics.first_principles, attitude_freq=50)
-    sim_step = sim._step
+    # Remove clipping floor function which kills gradients
+    sim.step_pipeline = sim.step_pipeline[:-1]
+    sim_step = sim.build_step_fn()
 
     def step(cmd: NDArray, data: SimData) -> jax.Array:
         data = data.replace(
@@ -23,7 +25,7 @@ def step(cmd: NDArray, data: SimData) -> jax.Array:
     step_grad = jax.jit(jax.grad(step))
 
     cmd = jnp.zeros((1, 1, 4), dtype=jnp.float32)
-    cmd = cmd.at[..., 3].set(0.3)
+    cmd = cmd.at[..., 3].set(sim.data.params.mass[0, 0, 0] * 9.81 * 1.05)
 
     # Trigger jax's jit to compile the gradient function. This is not necessary, but it ensures that
     # the timings are not affected by the compilation time.

examples/randomize.py

@@ -17,9 +17,9 @@ def main():
     # Randomize the inertia and mass of the drones
     mask = np.array([True, False, False])  # Only randomize the first world
     mass = sim.data.params.mass
-    mass_rng = mass + jax.random.normal(jax.random.key(0), (sim.n_worlds, sim.n_drones, 1)) * 1e-4
+    mass_rng = mass + jax.random.normal(jax.random.key(0), (sim.n_worlds, sim.n_drones, 1)) * 5e-3
     J = sim.data.params.J
-    J_rng = J + jax.random.normal(jax.random.key(0), (sim.n_worlds, sim.n_drones, 3, 3)) * 1e-5
+    J_rng = J + jax.random.normal(jax.random.key(0), (sim.n_worlds, sim.n_drones, 3, 3)) * 1e-6
 
     randomize_mass(sim, mass_rng, mask)
     # Note: The mask is optional. We can also randomize all worlds at once

examples/render.py

@@ -12,7 +12,7 @@ def main():
     sim = Sim(n_worlds=n_worlds, n_drones=n_drones, physics=Physics.so_rpy, device="cpu")
     fps = 60
     cmd = np.zeros((sim.n_worlds, sim.n_drones, 4))
-    cmd[..., 3] = sim.data.params.mass[0, 0, 0] * 9.81
+    cmd[..., 3] = sim.data.params.mass[0, 0, 0] * 9.81 * 1.05
     rgbas = np.random.default_rng(0).uniform(0, 1, (n_drones, 4))
     rgbas[..., 3] = 1.0