Fix vmap arg range

ratheron · ratheron · commit 34db28abbe5d · 2025-10-06T20:37:57.000+02:00
diff --git a/crazyflow/sim/sim.py b/crazyflow/sim/sim.py
@@ -237,7 +237,7 @@ def build_mjx_model(self, spec: mujoco.MjSpec) -> tuple[Any, Any, Model, Data]:
         mj_data = mujoco.MjData(mj_model)
         mjx_model = mjx.put_model(mj_model, device=self.device)
         mjx_data = mjx.put_data(mj_model, mj_data, device=self.device)
-        mjx_data = jax.vmap(lambda _: mjx_data)(range(self.n_worlds))
+        mjx_data = jax.vmap(lambda _: mjx_data)(jnp.arange(self.n_worlds))
         return mj_model, mj_data, mjx_model, mjx_data
 
     def build_step_fn(self):
diff --git a/pixi.lock b/pixi.lock
diff --git a/submodules/drone-controllers b/submodules/drone-controllers
@@ -1 +1 @@
-Subproject commit a480afec7520f04795f9a20f8603cca974e13ce7
+Subproject commit 7b3df61a9314721178d7cd83857a49166c59331a
diff --git a/submodules/drone-models b/submodules/drone-models
@@ -1 +1 @@
-Subproject commit 21d93d4a9deda185b088fc870ee59454e6f54b6c
+Subproject commit f83cabcfe78013d4fee1a27159766cdd58505acb
