Fix ruff. Fix testing. Update benchmark

Author: amacati

.github/workflows/testing.yml

@@ -1,22 +1,27 @@
-name: Testing  # Skips RL tests because stable-baselines3 comes with a lot of heavy-weight dependencies
-
-on: [push]
+name: Testing
+on: [push, pull_request]
 
 jobs:
   test:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
-      - uses: mamba-org/setup-micromamba@v1
+
+      - name: Setup Pixi (installs pixi + caches envs) # https://github.com/marketplace/actions/setup-pixi
+        uses: prefix-dev/setup-pixi@v0.9.0            # pin the action version
         with:
-          micromamba-version: '2.0.2-1' # any version from https://github.com/mamba-org/micromamba-releases
-          environment-name: test-env
-          init-shell: bash
-          create-args: python=3.11
-          cache-environment: true
-      - name: Install dependencies and package
-        run: pip install .[test]
-        shell: micromamba-shell {0}
-      - name: Test with pytest
-        run: pytest tests --cov=crazyflow
-        shell: micromamba-shell {0}
+          pixi-version: v0.49.0                       # pin the pixi binary version (optional)
+          cache: true                                 # enable caching of installed envs
+          # only write new caches on main pushes (TODO: Enable)
+          # cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
+          # ensure the 'test' environment(s) are installed
+          environments: test
+          # don't activate env (we'll call pixi run -e test explicitly)
+          activate-environment: false
+          # prefer using existing lockfile if present (faster, deterministic)
+          locked: true
+
+      - name: Verify pixi and run tests
+        run: |
+          pixi --version
+          pixi run -e test pytest

crazyflow/sim/structs.py

@@ -1,15 +1,10 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 import jax
 import jax.numpy as jnp
 from flax.struct import dataclass, field
 from jax import Array, Device
 
-if TYPE_CHECKING:
-    from mujoco.mjx import Data, Model
-
 
 @dataclass
 class SimState:

docs/img/performance.png

@@ -81,7 +81,7 @@ target-version = "py312"
 
 [tool.ruff.lint]
 select = ["E4", "E7", "E9", "F", "I", "D", "TCH", "ANN"]
-ignore = ["ANN401", "ANN101"]
+ignore = ["ANN401"]
 fixable = ["ALL"]
 unfixable = []
 
@@ -122,3 +122,11 @@ gpu = { features = ["gpu"], solve-group = "default" }
 test = { features = ["test", "gpu"], solve-group = "default" }
 
 [tool.pixi.tasks]
+
+[tool.pixi.dependencies]  
+# TODO: Remove compilers once scipy 1.17 is released and drone-models have been updated.
+gcc = ">=15.1.0,<15.2"
+cxx-compiler = ">=1.0.0,<2"
+gfortran = ">=15.1.0,<15.2"
+openblas = ">=0.3.30,<0.4"
+ruff = "*"