style: apply ruff format/lint to onnx examples

examples/onnx/eval_act_onnx.py

@@ -38,7 +38,9 @@ from lerobot.utils.random_utils import set_seed
 class ONNXACTModel(nn.Module):
     """Drop-in replacement for ``ACTPolicy.model`` backed by onnxruntime."""
 
-    def __init__(self, onnx_path: str, image_keys: list[str], has_state: bool, has_env_state: bool, device: str):
+    def __init__(
+        self, onnx_path: str, image_keys: list[str], has_state: bool, has_env_state: bool, device: str
+    ):
         super().__init__()
         import onnxruntime as ort
 
@@ -56,10 +58,7 @@ class ONNXACTModel(nn.Module):
         print(f"[onnx] providers in use: {self.sess.get_providers()}")
 
     def forward(self, batch: dict):
-        if self.has_state:
+        state = batch[OBS_STATE] if self.has_state else batch[OBS_ENV_STATE]
-            state = batch[OBS_STATE]
-        else:
-            state = batch[OBS_ENV_STATE]
         ref = state
         ort_inputs = {"state": state.detach().cpu().numpy().astype(np.float32)}
         images = batch[OBS_IMAGES]

examples/onnx/export_act.py

@@ -68,16 +68,16 @@ def main():
     has_env_state = cfg.env_state_feature is not None
     state_dim = (cfg.robot_state_feature or cfg.env_state_feature).shape[0]
 
-    print(f"      image_keys={image_keys} state_dim={state_dim} "
+    print(
-          f"chunk_size={cfg.chunk_size} action_dim={cfg.action_feature.shape[0]}")
+        f"      image_keys={image_keys} state_dim={state_dim} "
+        f"chunk_size={cfg.chunk_size} action_dim={cfg.action_feature.shape[0]}"
+    )
 
     wrapper = ACTExportWrapper(policy.model, image_keys, has_state, has_env_state).eval().to(args.device)
 
     # Build example inputs (batch size 1) from the config feature shapes.
     state_example = torch.randn(1, state_dim, device=args.device)
-    image_examples = [
+    image_examples = [torch.rand(1, *cfg.image_features[k].shape, device=args.device) for k in image_keys]
-        torch.rand(1, *cfg.image_features[k].shape, device=args.device) for k in image_keys
-    ]
     example_inputs = (state_example, *image_examples)
 
     input_names = ["state"] + [f"image_{i}" for i in range(len(image_keys))]