nit in control_robot.py

lerobot/common/datasets/image_writer.py

@@ -28,7 +28,7 @@ def wrapper(*args, **kwargs):
         try:
             return func(*args, **kwargs)
         except Exception as e:
-            dataset = kwargs.get("dataset", None)
+            dataset = kwargs.get("dataset")
             image_writer = getattr(dataset, "image_writer", None) if dataset else None
             if image_writer is not None:
                 print("Waiting for image writer to terminate...")

lerobot/common/logger.py

@@ -25,13 +25,13 @@
 from pathlib import Path
 
 import torch
+import wandb
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from omegaconf import DictConfig, OmegaConf
 from termcolor import colored
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 
-import wandb
 from lerobot.common.policies.policy_protocol import Policy
 from lerobot.common.utils.utils import get_global_random_state, set_global_random_state
 

lerobot/common/policies/hilserl/configuration_hilserl.py

@@ -15,7 +15,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 
 
 @dataclass

lerobot/common/policies/hilserl/modeling_hilserl.py

@@ -15,11 +15,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import torch
 import torch.nn as nn
-import torch.nn.functional as F  # noqa: N812
 from huggingface_hub import PyTorchModelHubMixin
 
+
 class HILSerlPolicy(
     nn.Module,
     PyTorchModelHubMixin,

lerobot/scripts/control_robot.py

@@ -224,12 +224,7 @@ def record(
     else:
         raise NotImplementedError("Only single-task recording is supported for now")
 
-    if single_task:
-        task = single_task
-    else:
-        raise NotImplementedError("Only single-task recording is supported for now")
-
-    # Load pretrained policy
+        # Load pretrained policy
     if pretrained_policy_name_or_path is not None:
         policy, policy_fps, device, use_amp = init_policy(pretrained_policy_name_or_path, policy_overrides)
 
@@ -263,6 +258,7 @@ def record(
             use_videos=video,
             image_writer_processes=num_image_writer_processes,
             image_writer_threads=num_image_writer_threads_per_camera * len(robot.cameras),
+            features=extra_features,
         )
 
     if not robot.is_connected:

lerobot/scripts/train_hilserl_classifier.py

@@ -22,6 +22,7 @@
 import hydra
 import torch
 import torch.nn as nn
+import wandb
 from deepdiff import DeepDiff
 from omegaconf import DictConfig, OmegaConf
 from termcolor import colored
@@ -30,7 +31,6 @@
 from torch.utils.data import DataLoader, WeightedRandomSampler, random_split
 from tqdm import tqdm
 
-import wandb
 from lerobot.common.datasets.factory import resolve_delta_timestamps
 from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.common.logger import Logger