From 30cbea056e22683378fe981c4ed7a68c787da9df Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Sun, 5 Jul 2026 17:32:54 +0200 Subject: [PATCH] style: apply ruff format/lint to onnx examples --- examples/onnx/eval_act_onnx.py | 9 ++++----- examples/onnx/export_act.py | 10 +++++----- 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/examples/onnx/eval_act_onnx.py b/examples/onnx/eval_act_onnx.py index d98938804..2c11f8b1f 100644 --- a/examples/onnx/eval_act_onnx.py +++ b/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] - else: - state = batch[OBS_ENV_STATE] + state = batch[OBS_STATE] if self.has_state else batch[OBS_ENV_STATE] ref = state ort_inputs = {"state": state.detach().cpu().numpy().astype(np.float32)} images = batch[OBS_IMAGES] diff --git a/examples/onnx/export_act.py b/examples/onnx/export_act.py index 9acc7d10e..37c30f220 100644 --- a/examples/onnx/export_act.py +++ b/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} " - f"chunk_size={cfg.chunk_size} action_dim={cfg.action_feature.shape[0]}") + print( + 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 = [ - torch.rand(1, *cfg.image_features[k].shape, device=args.device) for k in image_keys - ] + image_examples = [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))]