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

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

Author: pre-commit-ci[bot]

src/lerobot/policies/octo/base.py

@@ -15,7 +15,9 @@
 from dataclasses import dataclass
 from enum import Enum
 from fnmatch import fnmatch
-from typing import Any, Dict, List, Mapping
+from typing import Any, Dict, List
+from collections.abc import Mapping
+
 import torch
 
 
@@ -38,7 +40,7 @@ class AttentionRule(Enum):
 }
 
 
-def find_match(pattern_dict: Dict[str, Any], name: str, default: Any) -> Any:
+def find_match(pattern_dict: dict[str, Any], name: str, default: Any) -> Any:
     """Find the first matching pattern in the dictionary, or return the default value."""
     for pattern, value in pattern_dict.items():
         if fnmatch(name, pattern):
@@ -61,7 +63,7 @@ def __post_init__(self):
             )
 
     @classmethod
-    def concatenate(cls, group_list: List["TokenGroup"], axis: int = -2) -> "TokenGroup":
+    def concatenate(cls, group_list: list["TokenGroup"], axis: int = -2) -> "TokenGroup":
         """Concatenates a list of TokenGroups along a specified axis."""
         if not group_list:
             raise ValueError("Cannot concatenate an empty list of TokenGroups")

src/lerobot/policies/octo/diffusion.py

@@ -82,7 +82,7 @@ class MLPResNetBlock(nn.Module):
     """Implementation of MLPResNetBlock."""
 
     def __init__(
-        self, features: int, activation, dropout_rate: Optional[float] = None, use_layer_norm: bool = False
+        self, features: int, activation, dropout_rate: float | None = None, use_layer_norm: bool = False
     ):
         super().__init__()
         self.features = features
@@ -173,7 +173,7 @@ def __init__(
         num_blocks: int,
         out_dim: int,
         in_dim: int,
-        dropout_rate: Optional[float] = None,
+        dropout_rate: float | None = None,
         use_layer_norm: bool = False,
         hidden_dim: int = 256,
         activation=nn.SiLU,
@@ -280,9 +280,9 @@ def _cosine_beta_schedule(self, timesteps: int, s: float = 0.008) -> torch.Tenso
 
     def forward(
         self,
-        transformer_outputs: Dict[str, TokenGroup],
-        time: Optional[torch.Tensor] = None,
-        noisy_actions: Optional[torch.Tensor] = None,
+        transformer_outputs: dict[str, TokenGroup],
+        time: torch.Tensor | None = None,
+        noisy_actions: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Performs a single forward pass through the diffusion model."""
 
@@ -336,8 +336,8 @@ def loss(self, transformer_outputs, actions, timestep_pad_mask, action_pad_mask)
 
     def predict_action(
         self,
-        transformer_outputs: Dict[str, TokenGroup],
-        embodiment_action_dim: Optional[int] = None,
+        transformer_outputs: dict[str, TokenGroup],
+        embodiment_action_dim: int | None = None,
         sample_shape: tuple = (),
     ) -> torch.Tensor:
         """Convenience method for predicting actions for the final timestep."""

src/lerobot/policies/octo/modeling_octo.py

@@ -39,7 +39,8 @@
 """
 
 from collections import deque
-from typing import Dict, Optional, Sequence
+from typing import Dict, Optional
+from collections.abc import Sequence
 
 import torch
 import torch.nn as nn
@@ -216,7 +217,7 @@ def from_pretrained(cls, *args, **kwargs):
         return super().from_pretrained(*args, **kwargs)
 
     def _prepare_batch(
-        self, batch: dict[str, Tensor], raw_tasks: Optional[Sequence[str]] = None
+        self, batch: dict[str, Tensor], raw_tasks: Sequence[str] | None = None
     ) -> dict[str, Tensor]:
         """
         Prepare batch for model input.
@@ -330,9 +331,9 @@ def _prepare_batch(
 
     def create_tasks(
         self,
-        goals: Optional[Dict[str, torch.Tensor]] = None,
-        texts: Optional[Sequence[str]] = None,
-        device: Optional[torch.device] = None,
+        goals: dict[str, torch.Tensor] | None = None,
+        texts: Sequence[str] | None = None,
+        device: torch.device | None = None,
     ):
         """Creates tasks dict from goals and texts."""
         assert goals is not None or texts is not None
@@ -401,7 +402,7 @@ def _get_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
         return actions
 
     @torch.no_grad()
-    def predict_action_chunk(self, batch: dict[str, Tensor], tasks: Optional[Sequence[str]] = None) -> Tensor:
+    def predict_action_chunk(self, batch: dict[str, Tensor], tasks: Sequence[str] | None = None) -> Tensor:
         """Predict a chunk of actions given environment observations."""
         self.eval()
 
@@ -499,10 +500,10 @@ def __init__(self, config: OctoConfig):
 
     def forward(
         self,
-        observations: Dict[str, torch.Tensor],
-        tasks: Dict[str, torch.Tensor],
+        observations: dict[str, torch.Tensor],
+        tasks: dict[str, torch.Tensor],
         timestep_pad_mask: torch.Tensor,
-        embodiment_action_dim: Optional[int] = None,
+        embodiment_action_dim: int | None = None,
     ) -> torch.Tensor:
         transformer_outputs = self.octo_transformer(observations, tasks, timestep_pad_mask)
         actions = self.head.predict_action(

src/lerobot/policies/octo/tokenizers.py

@@ -15,7 +15,8 @@
 # limitations under the License.
 
 import re
-from typing import Dict, List, Optional, Sequence, Tuple
+from typing import Dict, List, Optional, Tuple
+from collections.abc import Sequence
 
 import numpy as np
 import torch
@@ -112,10 +113,10 @@ def __init__(
         self,
         use_film: bool = False,
         patch_size: int = 32,
-        kernel_sizes: Tuple[int, ...] = (3, 3, 3, 3),
-        strides: Tuple[int, ...] = (2, 2, 2, 2),
-        features: Tuple[int, ...] = (32, 96, 192, 384),
-        padding: Tuple[int, ...] = (1, 1, 1, 1),
+        kernel_sizes: tuple[int, ...] = (3, 3, 3, 3),
+        strides: tuple[int, ...] = (2, 2, 2, 2),
+        features: tuple[int, ...] = (32, 96, 192, 384),
+        padding: tuple[int, ...] = (1, 1, 1, 1),
         num_features: int = 512,
         img_norm_type: str = "default",
     ):
@@ -167,7 +168,7 @@ def __init__(
         self.film = FilmConditioning() if use_film else None
 
     def forward(
-        self, observations: torch.Tensor, train: bool = True, cond_var: Optional[torch.Tensor] = None
+        self, observations: torch.Tensor, train: bool = True, cond_var: torch.Tensor | None = None
     ):
         """
         Args:
@@ -212,10 +213,10 @@ class SmallStem16(SmallStem):
     def __init__(
         self,
         use_film: bool = False,
-        kernel_sizes: Tuple[int, ...] = (3, 3, 3, 3),
-        strides: Tuple[int, ...] = (2, 2, 2, 2),
-        features: Tuple[int, ...] = (32, 96, 192, 384),
-        padding: Tuple[int, ...] = (1, 1, 1, 1),
+        kernel_sizes: tuple[int, ...] = (3, 3, 3, 3),
+        strides: tuple[int, ...] = (2, 2, 2, 2),
+        features: tuple[int, ...] = (32, 96, 192, 384),
+        padding: tuple[int, ...] = (1, 1, 1, 1),
         num_features: int = 512,
         img_norm_type: str = "default",
     ):
@@ -243,7 +244,7 @@ def regex_filter(regex_keys, xs):
 
 def generate_proper_pad_mask(
     tokens: torch.Tensor,
-    pad_mask_dict: Optional[Dict[str, torch.Tensor]],
+    pad_mask_dict: dict[str, torch.Tensor] | None,
     keys: Sequence[str],
 ) -> torch.Tensor:
     """Generate proper padding mask for tokens."""
@@ -286,8 +287,8 @@ def __init__(
 
     def forward(
         self,
-        observations: Dict[str, torch.Tensor],
-        tasks: Optional[Dict[str, torch.Tensor]] = None,
+        observations: dict[str, torch.Tensor],
+        tasks: dict[str, torch.Tensor] | None = None,
     ):
         """Forward pass through image tokenizer."""
 
@@ -382,7 +383,7 @@ def __init__(self, finetune_encoder: bool = False, proper_pad_mask: bool = True)
             for param in self.t5_encoder.parameters():
                 param.requires_grad = False
 
-    def forward(self, language_input: Dict[str, torch.Tensor], tasks=None) -> TokenGroup:
+    def forward(self, language_input: dict[str, torch.Tensor], tasks=None) -> TokenGroup:
         outputs = self.t5_encoder(
             input_ids=language_input["input_ids"], attention_mask=language_input["attention_mask"]
         )
@@ -411,7 +412,7 @@ def forward(self, language_input: Dict[str, torch.Tensor], tasks=None) -> TokenG
 class TextProcessor:
     """HF Tokenizer wrapper."""
 
-    def __init__(self, tokenizer_name: str = "t5-base", tokenizer_kwargs: Optional[Dict] = None):
+    def __init__(self, tokenizer_name: str = "t5-base", tokenizer_kwargs: dict | None = None):
         if tokenizer_kwargs is None:
             tokenizer_kwargs = {
                 "max_length": 16,
@@ -423,6 +424,6 @@ def __init__(self, tokenizer_name: str = "t5-base", tokenizer_kwargs: Optional[D
         self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
         self.tokenizer_kwargs = tokenizer_kwargs
 
-    def encode(self, strings: List[str]) -> Dict[str, torch.Tensor]:
+    def encode(self, strings: list[str]) -> dict[str, torch.Tensor]:
         """Encode strings to token IDs and attention masks."""
         return self.tokenizer(strings, **self.tokenizer_kwargs)