[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

Diff for: docker/manyskill-lerobot-gpu/10_nvidia.json

@@ -3,4 +3,4 @@
     "ICD" : {
         "library_path" : "libEGL_nvidia.so.0"
     }
-}
+}

Diff for: docker/manyskill-lerobot-gpu/nvidia_icd.json

@@ -4,4 +4,4 @@
         "library_path": "libGLX_nvidia.so.0",
         "api_version" : "1.2.155"
     }
-}
+}

Diff for: docker/manyskill-lerobot-gpu/nvidia_layers.json

@@ -18,4 +18,4 @@
             "DISABLE_LAYER_NV_OPTIMUS_1": ""
         }
     }
-}
+}

Diff for: examples/3_train_policy.py

@@ -8,7 +8,10 @@
 
 import torch
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.common.datasets.lerobot_dataset import (
+    LeRobotDataset,
+    LeRobotDatasetMetadata,
+)
 from lerobot.common.datasets.utils import dataset_to_policy_features
 from lerobot.common.policies.diffusion.configuration_diffusion import DiffusionConfig
 from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy
@@ -34,12 +37,18 @@ def main():
     #   - dataset stats: for normalization and denormalization of input/outputs
     dataset_metadata = LeRobotDatasetMetadata("lerobot/pusht")
     features = dataset_to_policy_features(dataset_metadata.features)
-    output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
-    input_features = {key: ft for key, ft in features.items() if key not in output_features}
+    output_features = {
+        key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION
+    }
+    input_features = {
+        key: ft for key, ft in features.items() if key not in output_features
+    }
 
     # Policies are initialized with a configuration class, in this case `DiffusionConfig`. For this example,
     # we'll just use the defaults and so no arguments other than input/output features need to be passed.
-    cfg = DiffusionConfig(input_features=input_features, output_features=output_features)
+    cfg = DiffusionConfig(
+        input_features=input_features, output_features=output_features
+    )
 
     # We can now instantiate our policy with this config and the dataset stats.
     policy = DiffusionPolicy(cfg, dataset_stats=dataset_metadata.stats)
@@ -49,8 +58,12 @@ def main():
     # Another policy-dataset interaction is with the delta_timestamps. Each policy expects a given number frames
     # which can differ for inputs, outputs and rewards (if there are some).
     delta_timestamps = {
-        "observation.image": [i / dataset_metadata.fps for i in cfg.observation_delta_indices],
-        "observation.state": [i / dataset_metadata.fps for i in cfg.observation_delta_indices],
+        "observation.image": [
+            i / dataset_metadata.fps for i in cfg.observation_delta_indices
+        ],
+        "observation.state": [
+            i / dataset_metadata.fps for i in cfg.observation_delta_indices
+        ],
         "action": [i / dataset_metadata.fps for i in cfg.action_delta_indices],
     }
 
@@ -63,7 +76,24 @@ def main():
         # Load the previous action (-0.1), the next action to be executed (0.0),
         # and 14 future actions with a 0.1 seconds spacing. All these actions will be
         # used to supervise the policy.
-        "action": [-0.1, 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4],
+        "action": [
+            -0.1,
+            0.0,
+            0.1,
+            0.2,
+            0.3,
+            0.4,
+            0.5,
+            0.6,
+            0.7,
+            0.8,
+            0.9,
+            1.0,
+            1.1,
+            1.2,
+            1.3,
+            1.4,
+        ],
     }
 
     # We can then instantiate the dataset with these delta_timestamps configuration.

Diff for: examples/advanced/2_calculate_validation_loss.py

@@ -40,7 +40,24 @@ def main():
         # Load the previous action (-0.1), the next action to be executed (0.0),
         # and 14 future actions with a 0.1 seconds spacing. All these actions will be
         # used to calculate the loss.
-        "action": [-0.1, 0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4],
+        "action": [
+            -0.1,
+            0.0,
+            0.1,
+            0.2,
+            0.3,
+            0.4,
+            0.5,
+            0.6,
+            0.7,
+            0.8,
+            0.9,
+            1.0,
+            1.1,
+            1.2,
+            1.3,
+            1.4,
+        ],
     }
 
     # Load the last 10% of episodes of the dataset as a validation set.
@@ -59,7 +76,9 @@ def main():
     train_dataset = LeRobotDataset(
         "lerobot/pusht", episodes=train_episodes, delta_timestamps=delta_timestamps
     )
-    val_dataset = LeRobotDataset("lerobot/pusht", episodes=val_episodes, delta_timestamps=delta_timestamps)
+    val_dataset = LeRobotDataset(
+        "lerobot/pusht", episodes=val_episodes, delta_timestamps=delta_timestamps
+    )
     print(f"Number of frames in training dataset (90% subset): {len(train_dataset)}")
     print(f"Number of frames in validation dataset (10% subset): {len(val_dataset)}")
 

Diff for: lerobot/common/datasets/compute_stats.py

@@ -45,20 +45,20 @@ def get_stats_einops_patterns(dataset, num_workers=0):
         if key in dataset.meta.camera_keys:
             # sanity check that images are channel first
             _, c, h, w = batch[key].shape
-            assert (
-                c < h and c < w
-            ), f"expect channel first images, but instead {batch[key].shape}"
+            assert c < h and c < w, (
+                f"expect channel first images, but instead {batch[key].shape}"
+            )
 
             # sanity check that images are float32 in range [0,1]
-            assert (
-                batch[key].dtype == torch.float32
-            ), f"expect torch.float32, but instead {batch[key].dtype=}"
-            assert (
-                batch[key].max() <= 1
-            ), f"expect pixels lower than 1, but instead {batch[key].max()=}"
-            assert (
-                batch[key].min() >= 0
-            ), f"expect pixels greater than 1, but instead {batch[key].min()=}"
+            assert batch[key].dtype == torch.float32, (
+                f"expect torch.float32, but instead {batch[key].dtype=}"
+            )
+            assert batch[key].max() <= 1, (
+                f"expect pixels lower than 1, but instead {batch[key].max()=}"
+            )
+            assert batch[key].min() >= 0, (
+                f"expect pixels greater than 1, but instead {batch[key].min()=}"
+            )
 
             stats_patterns[key] = "b c h w -> c 1 1"
         elif batch[key].ndim == 2:

Diff for: lerobot/common/datasets/factory.py

@@ -58,7 +58,9 @@ def resolve_delta_timestamps(
         if key == "action" and cfg.action_delta_indices is not None:
             delta_timestamps[key] = [i / ds_meta.fps for i in cfg.action_delta_indices]
         if key.startswith("observation.") and cfg.observation_delta_indices is not None:
-            delta_timestamps[key] = [i / ds_meta.fps for i in cfg.observation_delta_indices]
+            delta_timestamps[key] = [
+                i / ds_meta.fps for i in cfg.observation_delta_indices
+            ]
 
     if len(delta_timestamps) == 0:
         delta_timestamps = None
@@ -79,11 +81,15 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
         LeRobotDataset | MultiLeRobotDataset
     """
     image_transforms = (
-        ImageTransforms(cfg.dataset.image_transforms) if cfg.dataset.image_transforms.enable else None
+        ImageTransforms(cfg.dataset.image_transforms)
+        if cfg.dataset.image_transforms.enable
+        else None
     )
 
     if isinstance(cfg.dataset.repo_id, str):
-        ds_meta = LeRobotDatasetMetadata(cfg.dataset.repo_id, local_files_only=cfg.dataset.local_files_only)
+        ds_meta = LeRobotDatasetMetadata(
+            cfg.dataset.repo_id, local_files_only=cfg.dataset.local_files_only
+        )
         delta_timestamps = resolve_delta_timestamps(cfg.policy, ds_meta)
         dataset = LeRobotDataset(
             cfg.dataset.repo_id,
@@ -110,6 +116,8 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
     if cfg.dataset.use_imagenet_stats:
         for key in dataset.meta.camera_keys:
             for stats_type, stats in IMAGENET_STATS.items():
-                dataset.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32)
+                dataset.meta.stats[key][stats_type] = torch.tensor(
+                    stats, dtype=torch.float32
+                )
 
     return dataset

Diff for: lerobot/common/datasets/lerobot_dataset.py

@@ -326,7 +326,9 @@ def create(
             # as this would break the dict flattening in the stats computation, which uses '/' as separator
             for key in features:
                 if "/" in key:
-                    raise ValueError(f"Feature names should not contain '/'. Found '/' in feature '{key}'.")
+                    raise ValueError(
+                        f"Feature names should not contain '/'. Found '/' in feature '{key}'."
+                    )
 
             features = {**features, **DEFAULT_FEATURES}
 

Diff for: lerobot/common/datasets/push_dataset_to_hub/aloha_hdf5_format.py

@@ -75,9 +75,9 @@ def check_format(raw_dir) -> bool:
                 else:
                     assert data[f"/observations/images/{camera}"].ndim == 4
                     b, h, w, c = data[f"/observations/images/{camera}"].shape
-                    assert (
-                        c < h and c < w
-                    ), f"Expect (h,w,c) image format but ({h=},{w=},{c=}) provided."
+                    assert c < h and c < w, (
+                        f"Expect (h,w,c) image format but ({h=},{w=},{c=}) provided."
+                    )
 
 
 def load_from_raw(

Diff for: lerobot/common/datasets/transforms.py

@@ -141,12 +141,16 @@ def _check_input(self, sharpness):
         return float(sharpness[0]), float(sharpness[1])
 
     def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
-        sharpness_factor = torch.empty(1).uniform_(self.sharpness[0], self.sharpness[1]).item()
+        sharpness_factor = (
+            torch.empty(1).uniform_(self.sharpness[0], self.sharpness[1]).item()
+        )
         return {"sharpness_factor": sharpness_factor}
 
     def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
         sharpness_factor = params["sharpness_factor"]
-        return self._call_kernel(F.adjust_sharpness, inpt, sharpness_factor=sharpness_factor)
+        return self._call_kernel(
+            F.adjust_sharpness, inpt, sharpness_factor=sharpness_factor
+        )
 
 
 @dataclass

Diff for: lerobot/common/envs/configs.py

@@ -47,10 +47,16 @@ class AlohaEnv(EnvConfig):
 
     def __post_init__(self):
         if self.obs_type == "pixels":
-            self.features["top"] = PolicyFeature(type=FeatureType.VISUAL, shape=(480, 640, 3))
+            self.features["top"] = PolicyFeature(
+                type=FeatureType.VISUAL, shape=(480, 640, 3)
+            )
         elif self.obs_type == "pixels_agent_pos":
-            self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(14,))
-            self.features["pixels/top"] = PolicyFeature(type=FeatureType.VISUAL, shape=(480, 640, 3))
+            self.features["agent_pos"] = PolicyFeature(
+                type=FeatureType.STATE, shape=(14,)
+            )
+            self.features["pixels/top"] = PolicyFeature(
+                type=FeatureType.VISUAL, shape=(480, 640, 3)
+            )
 
     @property
     def gym_kwargs(self) -> dict:
@@ -88,9 +94,13 @@ class PushtEnv(EnvConfig):
 
     def __post_init__(self):
         if self.obs_type == "pixels_agent_pos":
-            self.features["pixels"] = PolicyFeature(type=FeatureType.VISUAL, shape=(384, 384, 3))
+            self.features["pixels"] = PolicyFeature(
+                type=FeatureType.VISUAL, shape=(384, 384, 3)
+            )
         elif self.obs_type == "environment_state_agent_pos":
-            self.features["environment_state"] = PolicyFeature(type=FeatureType.ENV, shape=(16,))
+            self.features["environment_state"] = PolicyFeature(
+                type=FeatureType.ENV, shape=(16,)
+            )
 
     @property
     def gym_kwargs(self) -> dict:
@@ -129,7 +139,9 @@ class XarmEnv(EnvConfig):
 
     def __post_init__(self):
         if self.obs_type == "pixels_agent_pos":
-            self.features["agent_pos"] = PolicyFeature(type=FeatureType.STATE, shape=(4,))
+            self.features["agent_pos"] = PolicyFeature(
+                type=FeatureType.STATE, shape=(4,)
+            )
 
     @property
     def gym_kwargs(self) -> dict:

Diff for: lerobot/common/envs/factory.py

@@ -14,7 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import importlib
-from collections import deque
 
 import gymnasium as gym
 
@@ -32,7 +31,9 @@ def make_env_config(env_type: str, **kwargs) -> EnvConfig:
         raise ValueError(f"Policy type '{env_type}' is not available.")
 
 
-def make_env(cfg: EnvConfig, n_envs: int = 1, use_async_envs: bool = False) -> gym.vector.VectorEnv | None:
+def make_env(
+    cfg: EnvConfig, n_envs: int = 1, use_async_envs: bool = False
+) -> gym.vector.VectorEnv | None:
     """Makes a gym vector environment according to the config.
 
     Args:
@@ -56,16 +57,20 @@ def make_env(cfg: EnvConfig, n_envs: int = 1, use_async_envs: bool = False) -> g
     try:
         importlib.import_module(package_name)
     except ModuleNotFoundError as e:
-        print(f"{package_name} is not installed. Please install it with `pip install 'lerobot[{cfg.type}]'`")
+        print(
+            f"{package_name} is not installed. Please install it with `pip install 'lerobot[{cfg.type}]'`"
+        )
         raise e
 
     gym_handle = f"{package_name}/{cfg.task}"
 
     # batched version of the env that returns an observation of shape (b, c)
     env_cls = gym.vector.AsyncVectorEnv if use_async_envs else gym.vector.SyncVectorEnv
     env = env_cls(
-        [lambda: gym.make(gym_handle, disable_env_checker=True, **cfg.gym_kwargs) for _ in range(n_envs)]
+        [
+            lambda: gym.make(gym_handle, disable_env_checker=True, **cfg.gym_kwargs)
+            for _ in range(n_envs)
+        ]
     )
 
     return env
-

Diff for: lerobot/common/envs/utils.py

@@ -46,7 +46,9 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten
 
             # sanity check that images are channel last
             _, h, w, c = img.shape
-            assert c < h and c < w, f"expect channel last images, but instead got {img.shape=}"
+            assert c < h and c < w, (
+                f"expect channel last images, but instead got {img.shape=}"
+            )
 
         # sanity check that images are uint8
         assert img.dtype == torch.uint8, f"expect torch.uint8, but instead {img.dtype=}"
@@ -79,7 +81,9 @@ def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]:
     for key, ft in env_cfg.features.items():
         if ft.type is FeatureType.VISUAL:
             if len(ft.shape) != 3:
-                raise ValueError(f"Number of dimensions of {key} != 3 (shape={ft.shape})")
+                raise ValueError(
+                    f"Number of dimensions of {key} != 3 (shape={ft.shape})"
+                )
 
             shape = get_channel_first_image_shape(ft.shape)
             feature = PolicyFeature(type=ft.type, shape=shape)

Diff for: lerobot/common/logger.py

@@ -250,9 +250,9 @@ def load_last_training_state(
         )
         # For the case where the optimizer is a dictionary of optimizers (e.g., sac)
         if type(training_state["optimizer"]) is dict:
-            assert set(training_state["optimizer"].keys()) == set(
-                optimizer.keys()
-            ), "Optimizer dictionaries do not have the same keys during resume!"
+            assert set(training_state["optimizer"].keys()) == set(optimizer.keys()), (
+                "Optimizer dictionaries do not have the same keys during resume!"
+            )
             for k, v in training_state["optimizer"].items():
                 optimizer[k].load_state_dict(v)
         else:

Diff for: lerobot/common/optim/factory.py

@@ -34,7 +34,13 @@ def make_optimizer_and_scheduler(
     Returns:
         tuple[Optimizer, LRScheduler | None]: The couple (Optimizer, Scheduler). Scheduler can be `None`.
     """
-    params = policy.get_optim_params() if cfg.use_policy_training_preset else policy.parameters()
+    params = (
+        policy.get_optim_params()
+        if cfg.use_policy_training_preset
+        else policy.parameters()
+    )
     optimizer = cfg.optimizer.build(params)
-    lr_scheduler = cfg.scheduler.build(optimizer, cfg.steps) if cfg.scheduler is not None else None
+    lr_scheduler = (
+        cfg.scheduler.build(optimizer, cfg.steps) if cfg.scheduler is not None else None
+    )
     return optimizer, lr_scheduler

Diff for: lerobot/common/optim/optimizers.py

@@ -102,7 +102,9 @@ def save_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> No
     write_json(param_groups, save_dir / OPTIMIZER_PARAM_GROUPS)
 
 
-def load_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> torch.optim.Optimizer:
+def load_optimizer_state(
+    optimizer: torch.optim.Optimizer, save_dir: Path
+) -> torch.optim.Optimizer:
     current_state_dict = optimizer.state_dict()
     flat_state = load_file(save_dir / OPTIMIZER_STATE)
     state = unflatten_dict(flat_state)