refactor(inference): improve timeout handling and report timeout percentage

- Commented out the timeout handling logic to prevent appending timeout values to the results. - Added a print statement to display the percentage of timeouts during inference.
refactor(model): change tensor data type from bfloat16 to float32
2026-05-15 00:29:52 +00:00 · 2025-09-24 14:50:58 +02:00 · 2025-09-24 14:33:11 +02:00 · 2025-09-24 14:32:47 +02:00 · 2025-09-24 01:06:13 +02:00 · 2025-09-23 22:34:52 +02:00
187 changed files with 1708 additions and 1317 deletions
@@ -31,11 +31,11 @@ env:
    Feel free to reopen if is still relevant, or to ping a collaborator if you have any questions.
  WARN_ISSUE_MESSAGE: >
    This issue has been automatically marked as stale because it has not had
-    recent activity (6 months). It will be closed if no further activity occurs.
+    recent activity (1 year). It will be closed if no further activity occurs.
    Thank you for your contributions.
  WARN_PR_MESSAGE: >
    This PR has been automatically marked as stale because it has not had
-    recent activity (6 months). It will be closed if no further activity occurs.
+    recent activity (1 year). It will be closed if no further activity occurs.
    Thank you for your contributions.

 jobs:
@@ -0,0 +1,400 @@
+"""
+Benchmark memory footprint and inference latency of a policy on arbitrary devices.
+
+This script loads a pretrained policy directly (similar to the async inference server)
+and generates dummy input data based on the policy's input_features to perform
+accurate benchmarking without requiring datasets.
+"""
+
+import argparse
+import os
+import signal
+import statistics
+from contextlib import contextmanager
+from copy import deepcopy
+from datetime import datetime
+from pathlib import Path
+
+import psutil
+import torch
+from tqdm import tqdm
+
+from lerobot.configs.types import FeatureType
+from lerobot.policies.factory import get_policy_class, make_pre_post_processors
+from lerobot.policies.pretrained import PreTrainedPolicy
+
+
+class TimeoutExceptionError(Exception):
+    pass
+
+
+@contextmanager
+def timeout(seconds):
+    def signal_handler(signum, frame):
+        raise TimeoutExceptionError(f"Timed out after {seconds} seconds")
+
+    # On Windows, signal is not available, so we can't use this timeout mechanism
+    if not hasattr(signal, "SIGALRM"):
+        yield
+        return
+
+    old_handler = signal.signal(signal.SIGALRM, signal_handler)
+    try:
+        # signal.alarm expects integer seconds
+        # for float seconds, we can use setitimer
+        signal.setitimer(signal.ITIMER_REAL, seconds)
+        yield
+    finally:
+        signal.setitimer(signal.ITIMER_REAL, 0)
+        signal.signal(signal.SIGALRM, old_handler)
+
+
+def bytes_to_human(n: int) -> str:
+    for unit in ["B", "KB", "MB", "GB", "TB"]:
+        if n < 1024:
+            return f"{n:.2f} {unit}"
+        n /= 1024
+    return f"{n:.2f} PB"
+
+
+def percentile(values: list[float], p: float) -> float:
+    if not values:
+        return float("nan")
+    k = (len(values) - 1) * (p / 100.0)
+    f = int(k)
+    c = min(f + 1, len(values) - 1)
+    if f == c:
+        return values[f]
+    return values[f] + (values[c] - values[f]) * (k - f)
+
+
+def generate_dummy_observation(input_features: dict, device: str = "cpu") -> dict:
+    """Generate dummy observation data based on policy input features."""
+    dummy_obs = {}
+
+    for key, feature in input_features.items():
+        shape = feature.shape
+
+        if feature.type == FeatureType.VISUAL:
+            # Images: random values in [0, 1] range (already normalized)
+            dummy_obs[key] = torch.rand(shape, dtype=torch.float32, device=device)
+        elif feature.type in [FeatureType.STATE, FeatureType.ACTION, FeatureType.ENV]:
+            # State/action/env: random normal distribution
+            dummy_obs[key] = torch.randn(shape, dtype=torch.float32, device=device)
+        else:
+            # Default: random normal for unknown types
+            dummy_obs[key] = torch.randn(shape, dtype=torch.float32, device=device)
+
+    # # Add batch dimension
+    # for key in dummy_obs:
+    #     dummy_obs[key] = dummy_obs[key].unsqueeze(0)
+
+    # Add task string for language-conditioned policies
+    dummy_obs["task"] = " this is a dummy task"
+
+    return dummy_obs
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Policy inference benchmark")
+    parser.add_argument(
+        "--policy-id", type=str, required=True, help="Model ID or local path to pretrained policy"
+    )
+    parser.add_argument(
+        "--policy-type", type=str, required=True, help="Type of policy (smolvla, act, diffusion, etc.)"
+    )
+    parser.add_argument(
+        "--device", type=str, default="mps", choices=["cuda", "cpu", "mps"], help="Device to run on"
+    )
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    parser.add_argument(
+        "--num-samples", type=int, default=100, help="Number of inference samples to benchmark"
+    )
+    parser.add_argument("--warmup", type=int, default=10, help="Number of warmup samples (not timed)")
+    parser.add_argument(
+        "--output-dir", type=str, default="outputs/benchmarks", help="Directory to save benchmark results"
+    )
+    parser.add_argument(
+        "--timeout",
+        type=float,
+        default=0.3,
+        help="Timeout for each inference pass in seconds (default: 0.3s = 300ms)",
+    )
+    args = parser.parse_args()
+
+    # Seed & deterministic-ish setup
+    torch.manual_seed(args.seed)
+    if args.device == "cuda":
+        torch.cuda.manual_seed_all(args.seed)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = False  # leave False to avoid perf cliffs
+
+    # Resolve device availability
+    device = args.device.lower()
+    if device == "cuda" and not torch.cuda.is_available():
+        print("[!] CUDA requested but unavailable. Falling back to CPU.")
+        device = "cpu"
+    elif device == "mps" and not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()):
+        print("[!] MPS requested but unavailable. Falling back to CPU.")
+        device = "cpu"
+
+    use_cuda = device == "cuda"
+
+    # Create output directory and log file
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    policy_name = args.policy_id.replace("/", "_").replace("\\", "_")
+    log_file = output_dir / f"benchmark_{args.policy_type}_{policy_name}_{device}_{timestamp}.txt"
+
+    # Load policy directly from pretrained (similar to async inference server)
+    print(f"Loading policy {args.policy_type} from {args.policy_id}...")
+    policy_class = get_policy_class(args.policy_type)
+    policy: PreTrainedPolicy = policy_class.from_pretrained(args.policy_id)
+    policy.eval()
+    policy.to(device, torch.float32)
+    policy.config.device = device
+    preprocessor, postprocessor = make_pre_post_processors(policy.config)
+
+    print(f"Policy loaded on {device}")
+    print(f"Input features: {list(policy.config.input_features.keys())}")
+    print(f"Output features: {list(policy.config.output_features.keys())}")
+
+    # Generate dummy observation based on policy input features
+    dummy_observation = generate_dummy_observation(policy.config.input_features, device)
+    dummy_observation["task"] = "this is a dummy task"
+
+    # Helper to sync for fair timings
+    def _sync(dev_=device):
+        if dev_ == "cuda" and torch.cuda.is_available():
+            torch.cuda.synchronize()
+        elif dev_ == "mps" and hasattr(torch, "mps"):
+            try:
+                torch.mps.synchronize()
+            except AttributeError:
+                pass  # MPS sync not available in this PyTorch version
+
+    # Warmup (to stabilize kernels/caches)
+    print("Warming up...")
+    with torch.no_grad():
+        policy.reset()
+        preprocessor.reset()
+        postprocessor.reset()
+        orginal_dummy_observation = deepcopy(dummy_observation)
+        for _ in range(args.warmup):
+            dummy_observation_model = deepcopy(orginal_dummy_observation)
+            dummy_observation_model = preprocessor(dummy_observation_model)
+            action_model = policy.select_action(dummy_observation_model)
+            _ = postprocessor(action_model)
+            policy.reset()
+        _sync()
+
+    # Memory footprint before timing
+    process = psutil.Process(os.getpid())
+    rss_before = process.memory_info().rss
+    if use_cuda:
+        torch.cuda.reset_peak_memory_stats()
+
+    # PyTorch timing with Event objects for more accurate GPU timing
+    print(f"Running benchmark: {args.num_samples} samples...")
+
+    if use_cuda:
+        # Use CUDA Events for precise GPU timing
+        start_events = []
+        end_events = []
+        timeout_count = 0
+        orginal_dummy_observation = deepcopy(dummy_observation)
+
+        with torch.no_grad():
+            for forward in tqdm(range(args.num_samples), desc="Trials"):
+                start_event = torch.cuda.Event(enable_timing=True)
+                end_event = torch.cuda.Event(enable_timing=True)
+                try:
+                    dummy_observation_model = deepcopy(orginal_dummy_observation)
+                    dummy_observation_model = preprocessor(dummy_observation)
+                    with timeout(args.timeout):
+                        start_event.record()
+                        action_model = policy.select_action(dummy_observation_model)
+                        end_event.record()
+                    _ = postprocessor(action_model)
+                    policy.reset()
+
+                    start_events.append(start_event)
+                    end_events.append(end_event)
+                except TimeoutExceptionError:
+                    timeout_count += 1
+                    # Add placeholder for timeout
+                    start_events.append(None)
+                    end_events.append(None)
+                    print(f"\n[!] Timeout on forward {forward + 1}")
+                    continue
+
+        # Synchronize and collect timing results
+        torch.cuda.synchronize()
+        per_forward_ms = []
+        for start_event, end_event in zip(start_events, end_events, strict=True):
+            if start_event is None:
+                # per_forward_ms.append(args.timeout * 1000)
+                continue
+            else:
+                per_forward_ms.append(start_event.elapsed_time(end_event))
+
+        if timeout_count > 0:
+            print(f"[!] {timeout_count} inference passes timed out (>{args.timeout * 1000:.1f}ms)")
+
+    else:
+        # Use simple time.perf_counter for CPU/MPS timing with timeout
+        import time
+
+        per_forward_ms = []
+        timeout_count = 0
+
+        with torch.no_grad():
+            for sample in tqdm(range(args.num_samples), desc="Samples"):
+                try:
+                    dummy_observation_model = deepcopy(orginal_dummy_observation)
+                    dummy_observation_model = preprocessor(dummy_observation_model)
+                    with timeout(args.timeout):
+                        start_time = time.perf_counter()
+                        action_model = policy.select_action(dummy_observation_model)
+                        end_time = time.perf_counter()
+                    policy.reset()
+
+                    per_forward_ms.append((end_time - start_time) * 1000)  # Convert to ms
+                except TimeoutExceptionError:
+                    timeout_count += 1
+                    # per_forward_ms.append(args.timeout * 1000)
+
+                    print(f"\n[!] Timeout on sample {sample + 1}")
+                    continue
+
+        if timeout_count > 0:
+            print(f"[!] {timeout_count} inference passes timed out (>{args.timeout * 1000:.1f}ms)")
+            print(f"Timeout percentage: {timeout_count / args.num_samples * 100:.1f}%")
+
+    # Memory footprint after timing
+    rss_after = process.memory_info().rss
+    rss_delta = rss_after - rss_before
+    cuda_peak = torch.cuda.max_memory_allocated() if use_cuda else 0
+
+    # Sort timing results for percentile calculations
+    per_forward_ms_sorted = sorted(per_forward_ms)
+
+    mean_ms = statistics.fmean(per_forward_ms) if per_forward_ms else float("nan")
+    std_ms = statistics.pstdev(per_forward_ms) if len(per_forward_ms) > 1 else 0.0
+    min_ms = per_forward_ms_sorted[0] if per_forward_ms_sorted else float("nan")
+    max_ms = per_forward_ms_sorted[-1] if per_forward_ms_sorted else float("nan")
+    p50_ms = percentile(per_forward_ms_sorted, 50)
+    p95_ms = percentile(per_forward_ms_sorted, 95)
+
+    # Model size
+    num_params = sum(p.numel() for p in policy.parameters())
+
+    # Prepare results for logging
+    results = {
+        "timestamp": datetime.now().isoformat(),
+        "policy_type": args.policy_type,
+        "policy_id": args.policy_id,
+        "device": device,
+        "num_trials": args.num_samples,
+        "forwards_per_trial": 1,
+        "warmup": args.warmup,
+        "timeout_ms": args.timeout * 1000,
+        "seed": args.seed,
+        "num_params": num_params,
+        "timeout_count": timeout_count,
+        "latency_mean_ms": mean_ms,
+        "latency_std_ms": std_ms,
+        "latency_min_ms": min_ms,
+        "latency_max_ms": max_ms,
+        "latency_p50_ms": p50_ms,
+        "latency_p95_ms": p95_ms,
+        "cpu_rss_before": rss_before,
+        "cpu_rss_after": rss_after,
+        "cpu_rss_delta": rss_delta,
+        "cuda_peak_alloc": cuda_peak,
+        "input_features": list(policy.config.input_features.keys()),
+        "output_features": list(policy.config.output_features.keys()),
+    }
+
+    # Format and write results to log file
+    log_content = f"""
+=== LeRobot Policy Inference Benchmark ===
+Timestamp: {results["timestamp"]}
+Policy: {results["policy_type"]} ({results["policy_id"]})
+Device: {results["device"]}
+Seed: {results["seed"]}
+
+=== Model Information ===
+Parameters: {results["num_params"]:,}
+Input Features: {", ".join(results["input_features"])}
+Output Features: {", ".join(results["output_features"])}
+
+=== Benchmark Configuration ===
+Samples: {results["num_trials"]}
+Warmup: {results["warmup"]}
+Total Measurements: {len(per_forward_ms)}
+Timeout: {results["timeout_ms"]:.1f}ms
+Timeouts: {results["timeout_count"]} / {results["num_trials"]}
+
+=== Latency Results (ms) ===
+Mean:     {results["latency_mean_ms"]:.3f}
+Std Dev:  {results["latency_std_ms"]:.3f}
+Min:      {results["latency_min_ms"]:.3f}
+Max:      {results["latency_max_ms"]:.3f}
+P50:      {results["latency_p50_ms"]:.3f}
+P95:      {results["latency_p95_ms"]:.3f}
+
+=== Memory Footprint ===
+CPU RSS Before: {bytes_to_human(results["cpu_rss_before"])}
+CPU RSS After:  {bytes_to_human(results["cpu_rss_after"])} (Δ {bytes_to_human(results["cpu_rss_delta"])})
+"""
+
+    if use_cuda:
+        log_content += f"CUDA Peak:      {bytes_to_human(results['cuda_peak_alloc'])} (reset before timing)\n"
+
+    log_content += f"""
+=== Raw Timing Data (first 20 measurements, ms) ===
+{", ".join(f"{t:.3f}" for t in per_forward_ms[:20])}
+{"..." if len(per_forward_ms) > 20 else ""}
+
+=== Summary Statistics ===
+Timing Method: {"CUDA Events" if use_cuda else "torch.utils.benchmark.Timer"}
+Device Available: {torch.cuda.is_available() if device == "cuda" else torch.backends.mps.is_available() if device == "mps" else True}
+PyTorch Version: {torch.__version__}
+
+Benchmark completed successfully at {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
+"""
+
+    # Write to log file
+    with open(log_file, "w") as f:
+        f.write(log_content)
+
+    # Print to console (shorter version)
+    print("\n=== Inference Benchmark Results ===")
+    print(f"Policy: {args.policy_type} ({args.policy_id})")
+    print(f"Device: {device}")
+    print(f"Samples: {args.num_samples} | Warmup: {args.warmup}")
+    print(f"Model params: {num_params:,}")
+    print(f"Timeout percentage: {timeout_count / args.num_samples * 100:.1f}%")
+
+    print("\nLatency per forward (ms):")
+    print(f"  mean: {mean_ms:.3f}  std: {std_ms:.3f}")
+    print(f"  min:  {min_ms:.3f}   max: {max_ms:.3f}")
+    print(f"  p50:  {p50_ms:.3f}   p95: {p95_ms:.3f}")
+
+    print("\nMemory footprint:")
+    print(f"  CPU RSS before: {bytes_to_human(rss_before)}")
+    print(f"  CPU RSS after : {bytes_to_human(rss_after)}  (Δ {bytes_to_human(rss_delta)})")
+    if use_cuda:
+        print(
+            f"  CUDA peak allocated: {bytes_to_human(cuda_peak)} "
+            f"(reset by reset_peak_memory_stats before timing)"
+        )
+
+    print(f"\nResults saved to: {log_file}")
+    print("Benchmark completed successfully!")
+
+
+if __name__ == "__main__":
+    main()
@@ -35,13 +35,12 @@ import torch
 from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity
 from tqdm import tqdm

-from benchmarks.video.benchmark import TimeBenchmark
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.video_utils import (
    decode_video_frames_torchvision,
    encode_video_frames,
 )
-from lerobot.utils.constants import OBS_IMAGE
+from lerobot.utils.benchmark import TimeBenchmark

 BASE_ENCODING = OrderedDict(
    [
@@ -118,7 +117,7 @@ def save_first_episode(imgs_dir: Path, dataset: LeRobotDataset) -> None:
    hf_dataset = dataset.hf_dataset.with_format(None)

    # We only save images from the first camera
-    img_keys = [key for key in hf_dataset.features if key.startswith(OBS_IMAGE)]
+    img_keys = [key for key in hf_dataset.features if key.startswith("observation.image")]
    imgs_dataset = hf_dataset.select_columns(img_keys[0])

    for i, item in enumerate(
@@ -31,7 +31,7 @@ Then, spin up a policy server (in one terminal, or in a separate machine) specif
 You can spin up a policy server running:

 ```shell
-python src/lerobot/async_inference/policy_server.py \
+python src/lerobot/scripts/server/policy_server.py \
    --host=127.0.0.1 \
    --port=8080 \
 ```
@@ -39,7 +39,7 @@ python src/lerobot/async_inference/policy_server.py \
 This will start a policy server listening on `127.0.0.1:8080` (`localhost`, port 8080). At this stage, the policy server is empty, as all information related to which policy to run and with which parameters are specified during the first handshake with the client. Spin up a client with:

 ```shell
-python src/lerobot/async_inference/robot_client.py \
+python src/lerobot/scripts/server/robot_client.py \
    --server_address=127.0.0.1:8080 \ # SERVER: the host address and port of the policy server
    --robot.type=so100_follower \ # ROBOT: your robot type
    --robot.port=/dev/tty.usbmodem585A0076841 \ # ROBOT: your robot port
@@ -122,8 +122,8 @@ python -m lerobot.scripts.server.policy_server \

 <!-- prettier-ignore-start -->
 ```python
-from lerobot.async_inference.configs import PolicyServerConfig
-from lerobot.async_inference.policy_server import serve
+from lerobot.scripts.server.configs import PolicyServerConfig
+from lerobot.scripts.server.policy_server import serve

 config = PolicyServerConfig(
    host="localhost",
@@ -148,7 +148,7 @@ The `RobotClient` streams observations to the `PolicyServer`, and receives actio
 <hfoptions id="start_robot_client">
 <hfoption id="Command">
 ```bash
-python src/lerobot/async_inference/robot_client.py \
+python src/lerobot/scripts/server/robot_client.py \
    --server_address=127.0.0.1:8080 \ # SERVER: the host address and port of the policy server
    --robot.type=so100_follower \ # ROBOT: your robot type
    --robot.port=/dev/tty.usbmodem585A0076841 \ # ROBOT: your robot port
@@ -171,9 +171,9 @@ python src/lerobot/async_inference/robot_client.py \
 import threading
 from lerobot.robots.so100_follower import SO100FollowerConfig
 from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.async_inference.configs import RobotClientConfig
-from lerobot.async_inference.robot_client import RobotClient
-from lerobot.async_inference.helpers import visualize_action_queue_size
+from lerobot.scripts.server.configs import RobotClientConfig
+from lerobot.scripts.server.robot_client import RobotClient
+from lerobot.scripts.server.helpers import visualize_action_queue_size

 # 1. Create the robot instance
 """Check out the cameras available in your setup by running `python lerobot/find_cameras.py`"""
@@ -304,19 +304,19 @@ Before collecting demonstrations, you need to determine the appropriate operatio

 This helps simplify the problem of learning on the real robot in two ways: 1) by limiting the robot's operational space to a specific region that solves the task and avoids unnecessary or unsafe exploration, and 2) by allowing training in end-effector space rather than joint space. Empirically, learning in joint space for reinforcement learning in manipulation is often a harder problem - some tasks are nearly impossible to learn in joint space but become learnable when the action space is transformed to end-effector coordinates.

-**Using lerobot-find-joint-limits**
+**Using find_joint_limits.py**

 This script helps you find the safe operational bounds for your robot's end-effector. Given that you have a follower and leader arm, you can use the script to find the bounds for the follower arm that will be applied during training.
 Bounding the action space will reduce the redundant exploration of the agent and guarantees safety.

 ```bash
-lerobot-find-joint-limits \
-  --robot.type=so100_follower \
-  --robot.port=/dev/tty.usbmodem58760431541 \
-  --robot.id=black \
-  --teleop.type=so100_leader \
-  --teleop.port=/dev/tty.usbmodem58760431551 \
-  --teleop.id=blue
+python -m lerobot.scripts.find_joint_limits \
+    --robot.type=so100_follower \
+    --robot.port=/dev/tty.usbmodem58760431541 \
+    --robot.id=black \
+    --teleop.type=so100_leader \
+    --teleop.port=/dev/tty.usbmodem58760431551 \
+    --teleop.id=blue
 ```

 **Workflow**
@@ -200,7 +200,7 @@ from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderCo
 from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 from lerobot.record import record_loop

 NUM_EPISODES = 5
@@ -237,7 +237,7 @@ dataset = LeRobotDataset.create(

 # Initialize the keyboard listener and rerun visualization
 _, events = init_keyboard_listener()
-init_rerun(session_name="recording")
+_init_rerun(session_name="recording")

 # Connect the robot and teleoperator
 robot.connect()
@@ -517,7 +517,7 @@ from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerCon
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun
 from lerobot.record import record_loop
 from lerobot.policies.factory import make_processor

@@ -557,7 +557,7 @@ dataset = LeRobotDataset.create(

 # Initialize the keyboard listener and rerun visualization
 _, events = init_keyboard_listener()
-init_rerun(session_name="recording")
+_init_rerun(session_name="recording")

 # Connect the robot
 robot.connect()
@@ -277,7 +277,7 @@ leader.disconnect()
 </hfoption>
 </hfoptions>

-Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./il_robots)
+Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./getting_started_real_world_robot)

 > [!TIP]
 > If you have any questions or need help, please reach out on [Discord](https://discord.com/invite/s3KuuzsPFb).
@@ -323,7 +323,7 @@ To replay an episode run the API example below, make sure to change `remote_ip`,
 python examples/lekiwi/replay.py
 ```

-Congrats 🎉, your robot is all set to learn a task on its own. Start training it by the training part of this tutorial: [Getting started with real-world robots](./il_robots)
+Congrats 🎉, your robot is all set to learn a task on its own. Start training it by the training part of this tutorial: [Getting started with real-world robots](./getting_started_real_world_robot)

 ## Evaluate your policy

@@ -33,7 +33,7 @@ To Install LIBERO, after following LeRobot official instructions, just do:
 Evaluate a policy on one LIBERO suite:

 ```bash
-lerobot-eval \
+python src/lerobot/scripts/eval.py \
  --policy.path="your-policy-id" \
  --env.type=libero \
  --env.task=libero_object \
@@ -52,7 +52,7 @@ lerobot-eval \
 Benchmark a policy across multiple suites at once:

 ```bash
-lerobot-eval \
+python src/lerobot/scripts/eval.py \
  --policy.path="your-policy-id" \
  --env.type=libero \
  --env.task=libero_object,libero_spatial \
@@ -103,11 +103,10 @@ For reference, here is the **original dataset** published by Physical Intelligen
 ### Example training command

 ```bash
-lerobot-train \
+python src/lerobot/scripts/train.py \
  --policy.type=smolvla \
  --policy.repo_id=${HF_USER}/libero-test \
-  --policy.load_vlm_weights=true \
-  --dataset.repo_id=HuggingFaceVLA/libero \
+  --dataset.repo_id=jadechoghari/smol-libero3 \
  --env.type=libero \
  --env.task=libero_10 \
  --output_dir=./outputs/ \
@@ -29,7 +29,7 @@ SmolVLA is Hugging Face’s lightweight foundation model for robotics. Designed
 ## Collect a dataset

 SmolVLA is a base model, so fine-tuning on your own data is required for optimal performance in your setup.
-We recommend recording ~50 episodes of your task as a starting point. Follow our guide to get started: [Recording a Dataset](./il_robots)
+We recommend recording ~50 episodes of your task as a starting point. Follow our guide to get started: [Recording a Dataset](https://huggingface.co/docs/lerobot/getting_started_real_world_robot#record-a-dataset)

 <Tip>

@@ -93,7 +93,7 @@ lerobot-train --help

 ## Evaluate the finetuned model and run it in real-time

-Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
+Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./getting_started_real_world_robot#record-a-dataset).
 Once you are logged in, you can run inference in your setup by doing:

 ```bash
@@ -634,7 +634,7 @@ leader.disconnect()
 </hfoption>
 </hfoptions>

-Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./il_robots)
+Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./getting_started_real_world_robot)

 > [!TIP]
 > If you have any questions or need help, please reach out on [Discord](https://discord.com/invite/s3KuuzsPFb).
@@ -430,7 +430,7 @@ leader.disconnect()
 </hfoption>
 </hfoptions>

-Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./il_robots)
+Congrats 🎉, your robot is all set to learn a task on its own. Start training it by following this tutorial: [Getting started with real-world robots](./getting_started_real_world_robot)

 > [!TIP]
 > If you have any questions or need help, please reach out on [Discord](https://discord.com/invite/s3KuuzsPFb).
@@ -44,7 +44,6 @@ from lerobot.robots import (  # noqa: F401
    so100_follower,
    so101_follower,
 )
-from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import (
    init_logging,
@@ -79,16 +78,16 @@ def replay(cfg: ReplayConfig):

    robot = make_robot_from_config(cfg.robot)
    dataset = LeRobotDataset(cfg.dataset.repo_id, root=cfg.dataset.root, episodes=[cfg.dataset.episode])
-    actions = dataset.hf_dataset.select_columns(ACTION)
+    actions = dataset.hf_dataset.select_columns("action")
    robot.connect()

    log_say("Replaying episode", cfg.play_sounds, blocking=True)
    for idx in range(dataset.num_frames):
        start_episode_t = time.perf_counter()

-        action_array = actions[idx][ACTION]
+        action_array = actions[idx]["action"]
        action = {}
-        for i, name in enumerate(dataset.features[ACTION]["names"]):
+        for i, name in enumerate(dataset.features["action"]["names"]):
            key = f"{name.removeprefix('main_')}.pos"
            action[key] = action_array[i].item()

@@ -19,12 +19,11 @@ from lerobot.datasets.utils import hw_to_dataset_features
 from lerobot.policies.act.modeling_act import ACTPolicy
 from lerobot.policies.factory import make_pre_post_processors
 from lerobot.processor import make_default_processors
+from lerobot.record import record_loop
 from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
-from lerobot.scripts.lerobot_record import record_loop
-from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun

 NUM_EPISODES = 2
 FPS = 30
@@ -42,8 +41,8 @@ robot = LeKiwiClient(robot_config)
 policy = ACTPolicy.from_pretrained(HF_MODEL_ID)

 # Configure the dataset features
-action_features = hw_to_dataset_features(robot.action_features, ACTION)
-obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
+action_features = hw_to_dataset_features(robot.action_features, "action")
+obs_features = hw_to_dataset_features(robot.observation_features, "observation")
 dataset_features = {**action_features, **obs_features}

 # Create the dataset
@@ -74,7 +73,7 @@ teleop_action_processor, robot_action_processor, robot_observation_processor = m

 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="lekiwi_evaluate")
+_init_rerun(session_name="lekiwi_evaluate")

 if not robot.is_connected:
    raise ValueError("Robot is not connected!")
@@ -17,15 +17,14 @@
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.datasets.utils import hw_to_dataset_features
 from lerobot.processor import make_default_processors
+from lerobot.record import record_loop
 from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
 from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.keyboard import KeyboardTeleop, KeyboardTeleopConfig
 from lerobot.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
-from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun

 NUM_EPISODES = 2
 FPS = 30
@@ -48,8 +47,8 @@ keyboard = KeyboardTeleop(keyboard_config)
 teleop_action_processor, robot_action_processor, robot_observation_processor = make_default_processors()

 # Configure the dataset features
-action_features = hw_to_dataset_features(robot.action_features, ACTION)
-obs_features = hw_to_dataset_features(robot.observation_features, OBS_STR)
+action_features = hw_to_dataset_features(robot.action_features, "action")
+obs_features = hw_to_dataset_features(robot.observation_features, "observation")
 dataset_features = {**action_features, **obs_features}

 # Create the dataset
@@ -70,7 +69,7 @@ keyboard.connect()

 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="lekiwi_record")
+_init_rerun(session_name="lekiwi_record")

 if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
    raise ValueError("Robot or teleop is not connected!")
@@ -19,7 +19,6 @@ import time
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
 from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient
-from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import log_say

@@ -35,7 +34,7 @@ robot = LeKiwiClient(robot_config)
 dataset = LeRobotDataset("<hf_username>/<dataset_repo_id>", episodes=[EPISODE_IDX])
 # Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
 episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
-actions = episode_frames.select_columns(ACTION)
+actions = episode_frames.select_columns("action")

 # Connect to the robot
 robot.connect()
@@ -50,7 +49,7 @@ for idx in range(len(episode_frames)):

    # Get recorded action from dataset
    action = {
-        name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
    }

    # Send action to robot
@@ -20,7 +20,7 @@ from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
 from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop, KeyboardTeleopConfig
 from lerobot.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
 from lerobot.utils.robot_utils import busy_wait
-from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
+from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data

 FPS = 30

@@ -41,7 +41,7 @@ leader_arm.connect()
 keyboard.connect()

 # Init rerun viewer
-init_rerun(session_name="lekiwi_teleop")
+_init_rerun(session_name="lekiwi_teleop")

 if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
    raise ValueError("Robot or teleop is not connected!")
@@ -34,16 +34,16 @@ from lerobot.processor.converters import (
    transition_to_observation,
    transition_to_robot_action,
 )
+from lerobot.record import record_loop
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
    ForwardKinematicsJointsToEE,
    InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun

 NUM_EPISODES = 5
 FPS = 30
@@ -137,7 +137,7 @@ robot.connect()

 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="phone_so100_evaluate")
+_init_rerun(session_name="phone_so100_evaluate")

 if not robot.is_connected:
    raise ValueError("Robot is not connected!")
@@ -26,6 +26,7 @@ from lerobot.processor.converters import (
    transition_to_observation,
    transition_to_robot_action,
 )
+from lerobot.record import record_loop
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
    EEBoundsAndSafety,
@@ -35,13 +36,12 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
    InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
 from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
 from lerobot.teleoperators.phone.teleop_phone import Phone
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun

 NUM_EPISODES = 2
 FPS = 30
@@ -143,7 +143,7 @@ phone.connect()

 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="phone_so100_record")
+_init_rerun(session_name="phone_so100_record")

 if not robot.is_connected or not phone.is_connected:
    raise ValueError("Robot or teleop is not connected!")
@@ -28,7 +28,6 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
    InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import log_say

@@ -67,7 +66,7 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
 dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
 # Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
 episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
-actions = episode_frames.select_columns(ACTION)
+actions = episode_frames.select_columns("action")

 # Connect to the robot
 robot.connect()
@@ -82,7 +81,7 @@ for idx in range(len(episode_frames)):

    # Get recorded action from dataset
    ee_action = {
-        name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
    }

    # Get robot observation
@@ -33,7 +33,7 @@ from lerobot.teleoperators.phone.config_phone import PhoneConfig, PhoneOS
 from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
 from lerobot.teleoperators.phone.teleop_phone import Phone
 from lerobot.utils.robot_utils import busy_wait
-from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
+from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data

 FPS = 30

@@ -87,7 +87,7 @@ robot.connect()
 teleop_device.connect()

 # Init rerun viewer
-init_rerun(session_name="phone_so100_teleop")
+_init_rerun(session_name="phone_so100_teleop")

 if not robot.is_connected or not teleop_device.is_connected:
    raise ValueError("Robot or teleop is not connected!")
@@ -34,16 +34,16 @@ from lerobot.processor.converters import (
    transition_to_observation,
    transition_to_robot_action,
 )
+from lerobot.record import record_loop
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
    ForwardKinematicsJointsToEE,
    InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun

 NUM_EPISODES = 5
 FPS = 30
@@ -138,7 +138,7 @@ robot.connect()

 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="so100_so100_evaluate")
+_init_rerun(session_name="so100_so100_evaluate")

 if not robot.is_connected:
    raise ValueError("Robot is not connected!")
@@ -27,6 +27,7 @@ from lerobot.processor.converters import (
    transition_to_observation,
    transition_to_robot_action,
 )
+from lerobot.record import record_loop
 from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
 from lerobot.robots.so100_follower.robot_kinematic_processor import (
    EEBoundsAndSafety,
@@ -34,12 +35,11 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
    InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderConfig
 from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader
 from lerobot.utils.control_utils import init_keyboard_listener
 from lerobot.utils.utils import log_say
-from lerobot.utils.visualization_utils import init_rerun
+from lerobot.utils.visualization_utils import _init_rerun

 NUM_EPISODES = 2
 FPS = 30
@@ -143,7 +143,7 @@ follower.connect()

 # Initialize the keyboard listener and rerun visualization
 listener, events = init_keyboard_listener()
-init_rerun(session_name="recording_phone")
+_init_rerun(session_name="recording_phone")

 if not leader.is_connected or not follower.is_connected:
    raise ValueError("Robot or teleop is not connected!")
@@ -29,7 +29,6 @@ from lerobot.robots.so100_follower.robot_kinematic_processor import (
    InverseKinematicsEEToJoints,
 )
 from lerobot.robots.so100_follower.so100_follower import SO100Follower
-from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import log_say

@@ -68,7 +67,7 @@ robot_ee_to_joints_processor = RobotProcessorPipeline[tuple[RobotAction, RobotOb
 dataset = LeRobotDataset(HF_REPO_ID, episodes=[EPISODE_IDX])
 # Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
 episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == EPISODE_IDX)
-actions = episode_frames.select_columns(ACTION)
+actions = episode_frames.select_columns("action")

 # Connect to the robot
 robot.connect()
@@ -83,7 +82,7 @@ for idx in range(len(episode_frames)):

    # Get recorded action from dataset
    ee_action = {
-        name: float(actions[idx][ACTION][i]) for i, name in enumerate(dataset.features[ACTION]["names"])
+        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
    }

    # Get robot observation
@@ -33,7 +33,7 @@ from lerobot.robots.so100_follower.so100_follower import SO100Follower
 from lerobot.teleoperators.so100_leader.config_so100_leader import SO100LeaderConfig
 from lerobot.teleoperators.so100_leader.so100_leader import SO100Leader
 from lerobot.utils.robot_utils import busy_wait
-from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
+from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data

 FPS = 30

@@ -95,7 +95,7 @@ follower.connect()
 leader.connect()

 # Init rerun viewer
-init_rerun(session_name="so100_so100_EE_teleop")
+_init_rerun(session_name="so100_so100_EE_teleop")

 print("Starting teleop loop...")
 while True:
@@ -20,13 +20,13 @@ from pathlib import Path
 import torch

 from lerobot.configs.types import FeatureType
+from lerobot.constants import ACTION
 from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
 from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
 from lerobot.datasets.utils import dataset_to_policy_features
 from lerobot.policies.act.configuration_act import ACTConfig
 from lerobot.policies.act.modeling_act import ACTPolicy
 from lerobot.policies.factory import make_pre_post_processors
-from lerobot.utils.constants import ACTION


 def main():
@@ -162,18 +162,17 @@ all = [
 ]

 [project.scripts]
-lerobot-calibrate="lerobot.scripts.lerobot_calibrate:main"
-lerobot-find-cameras="lerobot.scripts.lerobot_find_cameras:main"
-lerobot-find-port="lerobot.scripts.lerobot_find_port:main"
-lerobot-record="lerobot.scripts.lerobot_record:main"
-lerobot-replay="lerobot.scripts.lerobot_replay:main"
-lerobot-setup-motors="lerobot.scripts.lerobot_setup_motors:main"
-lerobot-teleoperate="lerobot.scripts.lerobot_teleoperate:main"
-lerobot-eval="lerobot.scripts.lerobot_eval:main"
-lerobot-train="lerobot.scripts.lerobot_train:main"
+lerobot-calibrate="lerobot.calibrate:main"
+lerobot-find-cameras="lerobot.find_cameras:main"
+lerobot-find-port="lerobot.find_port:main"
+lerobot-record="lerobot.record:main"
+lerobot-replay="lerobot.replay:main"
+lerobot-setup-motors="lerobot.setup_motors:main"
+lerobot-teleoperate="lerobot.teleoperate:main"
+lerobot-eval="lerobot.scripts.eval:main"
+lerobot-train="lerobot.scripts.train:main"
 lerobot-dataset-viz="lerobot.scripts.lerobot_dataset_viz:main"
 lerobot-info="lerobot.scripts.lerobot_info:main"
-lerobot-find-joint-limits="lerobot.scripts.lerobot_find_joint_limits:main"
 lerobot-imgtransform-viz="lerobot.scripts.lerobot_imgtransform_viz:main"

 # ---------------- Tool Configurations ----------------
@@ -267,83 +266,8 @@ default.extend-ignore-identifiers-re = [
 # color = true
 # paths = ["src/lerobot"]

-# TODO: Enable mypy gradually module by module across multiple PRs
-# Uncomment [tool.mypy] first, then uncomment individual module overrides as they get proper type annotations
-
 # [tool.mypy]
 # python_version = "3.10"
 # warn_return_any = true
 # warn_unused_configs = true
 # ignore_missing_imports = false
-# strict = true
-# disallow_untyped_defs = true
-# disallow_incomplete_defs = true
-# check_untyped_defs = true
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.utils.*"
-# # include = "src/lerobot/utils/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.configs.*"
-# # include = "src/lerobot/configs/**/*.py"
-
-# # Data processing modules
-# [[tool.mypy.overrides]]
-# module = "lerobot.processor.*"
-# # include = "src/lerobot/processor/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.datasets.*"
-# # include = "src/lerobot/datasets/**/*.py"
-
-# # Core machine learning modules
-# [[tool.mypy.overrides]]
-# module = "lerobot.optim.*"
-# # include = "src/lerobot/optim/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.model.*"
-# # include = "src/lerobot/model/**/*.py"
-
-# # Hardware interfaces
-# [[tool.mypy.overrides]]
-# module = "lerobot.cameras.*"
-# # include = "src/lerobot/cameras/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.motors.*"
-# # include = "src/lerobot/motors/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.robots.*"
-# # include = "src/lerobot/robots/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.teleoperators.*"
-# # include = "src/lerobot/teleoperators/**/*.py"
-
-# # Complex modules (enable these last)
-# [[tool.mypy.overrides]]
-# module = "lerobot.policies.*"
-# # include = "src/lerobot/policies/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.rl.*"
-# # include = "src/lerobot/rl/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.envs.*"
-# # include = "src/lerobot/envs/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.async_inference.*"
-# # include = "src/lerobot/async_inference/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.transport.*"
-# # include = "src/lerobot/transport/**/*.py"
-
-# [[tool.mypy.overrides]]
-# module = "lerobot.scripts.*"
-# # include = "src/lerobot/scripts/**/*.py"
@@ -31,7 +31,7 @@ if platform.system() == "Windows" and "OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"
 import cv2
 import numpy as np

-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError

 from ..camera import Camera
 from ..utils import get_cv2_backend, get_cv2_rotation
@@ -31,7 +31,7 @@ import numpy as np
 from reachy2_sdk.media.camera import CameraView
 from reachy2_sdk.media.camera_manager import CameraManager

-from lerobot.utils.errors import DeviceNotConnectedError
+from lerobot.errors import DeviceNotConnectedError

 from ..camera import Camera
 from .configuration_reachy2_camera import ColorMode, Reachy2CameraConfig
@@ -29,7 +29,7 @@ try:
 except Exception as e:
    logging.info(f"Could not import realsense: {e}")

-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError

 from ..camera import Camera
 from ..configs import ColorMode
@@ -27,9 +27,9 @@ from huggingface_hub.constants import CONFIG_NAME
 from huggingface_hub.errors import HfHubHTTPError

 from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.constants import ACTION, OBS_STATE
 from lerobot.optim.optimizers import OptimizerConfig
 from lerobot.optim.schedulers import LRSchedulerConfig
-from lerobot.utils.constants import ACTION, OBS_STATE
 from lerobot.utils.hub import HubMixin
 from lerobot.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available

@@ -17,21 +17,19 @@ from pathlib import Path

 from huggingface_hub.constants import HF_HOME

-OBS_STR = "observation"
-OBS_PREFIX = OBS_STR + "."
-OBS_ENV_STATE = OBS_STR + ".environment_state"
-OBS_STATE = OBS_STR + ".state"
-OBS_IMAGE = OBS_STR + ".image"
-OBS_IMAGES = OBS_IMAGE + "s"
-OBS_LANGUAGE = OBS_STR + ".language"
-OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
-OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
-
+OBS_ENV_STATE = "observation.environment_state"
+OBS_STATE = "observation.state"
+OBS_IMAGE = "observation.image"
+OBS_IMAGES = "observation.images"
+OBS_LANGUAGE = "observation.language"
 ACTION = "action"
 REWARD = "next.reward"
 TRUNCATED = "next.truncated"
 DONE = "next.done"

+OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
+OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
+
 ROBOTS = "robots"
 ROBOT_TYPE = "robot_type"
 TELEOPERATORS = "teleoperators"
@@ -27,7 +27,6 @@ from lerobot.datasets.lerobot_dataset import (
 )
 from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
 from lerobot.datasets.transforms import ImageTransforms
-from lerobot.utils.constants import ACTION, OBS_PREFIX, REWARD

 IMAGENET_STATS = {
    "mean": [[[0.485]], [[0.456]], [[0.406]]],  # (c,1,1)
@@ -55,11 +54,11 @@ def resolve_delta_timestamps(
    """
    delta_timestamps = {}
    for key in ds_meta.features:
-        if key == REWARD and cfg.reward_delta_indices is not None:
+        if key == "next.reward" and cfg.reward_delta_indices is not None:
            delta_timestamps[key] = [i / ds_meta.fps for i in cfg.reward_delta_indices]
-        if key == ACTION and cfg.action_delta_indices is not None:
+        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(OBS_PREFIX) and cfg.observation_delta_indices is not None:
+        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]

    if len(delta_timestamps) == 0:
@@ -31,6 +31,7 @@ import torch.utils
 from huggingface_hub import HfApi, snapshot_download
 from huggingface_hub.errors import RevisionNotFoundError

+from lerobot.constants import HF_LEROBOT_HOME
 from lerobot.datasets.compute_stats import aggregate_stats, compute_episode_stats
 from lerobot.datasets.image_writer import AsyncImageWriter, write_image
 from lerobot.datasets.utils import (
@@ -78,7 +79,6 @@ from lerobot.datasets.video_utils import (
    get_video_duration_in_s,
    get_video_info,
 )
-from lerobot.utils.constants import HF_LEROBOT_HOME

 CODEBASE_VERSION = "v3.0"

@@ -17,9 +17,9 @@ from collections.abc import Sequence
 from typing import Any

 from lerobot.configs.types import PipelineFeatureType
+from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE
 from lerobot.datasets.utils import hw_to_dataset_features
 from lerobot.processor import DataProcessorPipeline
-from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR


 def create_initial_features(
@@ -92,8 +92,8 @@ def aggregate_pipeline_dataset_features(

    # Intermediate storage for categorized and filtered features.
    processed_features: dict[str, dict[str, Any]] = {
-        ACTION: {},
-        OBS_STR: {},
+        "action": {},
+        "observation": {},
    }
    images_token = OBS_IMAGES.split(".")[-1]

@@ -125,15 +125,17 @@ def aggregate_pipeline_dataset_features(
            # 3. Add the feature to the appropriate group with a clean name.
            name = strip_prefix(key, PREFIXES_TO_STRIP)
            if is_action:
-                processed_features[ACTION][name] = value
+                processed_features["action"][name] = value
            else:
-                processed_features[OBS_STR][name] = value
+                processed_features["observation"][name] = value

    # Convert the processed features into the final dataset format.
    dataset_features = {}
-    if processed_features[ACTION]:
-        dataset_features.update(hw_to_dataset_features(processed_features[ACTION], ACTION, use_videos))
-    if processed_features[OBS_STR]:
-        dataset_features.update(hw_to_dataset_features(processed_features[OBS_STR], OBS_STR, use_videos))
+    if processed_features["action"]:
+        dataset_features.update(hw_to_dataset_features(processed_features["action"], ACTION, use_videos))
+    if processed_features["observation"]:
+        dataset_features.update(
+            hw_to_dataset_features(processed_features["observation"], "observation", use_videos)
+        )

    return dataset_features
@@ -21,6 +21,7 @@ import numpy as np
 import torch
 from datasets import load_dataset

+from lerobot.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE
 from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDatasetMetadata
 from lerobot.datasets.utils import (
    Backtrackable,
@@ -37,7 +38,6 @@ from lerobot.datasets.video_utils import (
    VideoDecoderCache,
    decode_video_frames_torchcodec,
 )
-from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE


 class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
@@ -43,7 +43,6 @@ from lerobot.datasets.backward_compatibility import (
    BackwardCompatibilityError,
    ForwardCompatibilityError,
 )
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_STR
 from lerobot.utils.utils import is_valid_numpy_dtype_string

 DEFAULT_CHUNK_SIZE = 1000  # Max number of files per chunk
@@ -646,14 +645,14 @@ def hw_to_dataset_features(
    }
    cam_fts = {key: shape for key, shape in hw_features.items() if isinstance(shape, tuple)}

-    if joint_fts and prefix == ACTION:
+    if joint_fts and prefix == "action":
        features[prefix] = {
            "dtype": "float32",
            "shape": (len(joint_fts),),
            "names": list(joint_fts),
        }

-    if joint_fts and prefix == OBS_STR:
+    if joint_fts and prefix == "observation":
        features[f"{prefix}.state"] = {
            "dtype": "float32",
            "shape": (len(joint_fts),),
@@ -729,11 +728,11 @@ def dataset_to_policy_features(features: dict[str, dict]) -> dict[str, PolicyFea
            # Backward compatibility for "channel" which is an error introduced in LeRobotDataset v2.0 for ported datasets.
            if names[2] in ["channel", "channels"]:  # (h, w, c) -> (c, h, w)
                shape = (shape[2], shape[0], shape[1])
-        elif key == OBS_ENV_STATE:
+        elif key == "observation.environment_state":
            type = FeatureType.ENV
-        elif key.startswith(OBS_STR):
+        elif key.startswith("observation"):
            type = FeatureType.STATE
-        elif key.startswith(ACTION):
+        elif key.startswith("action"):
            type = FeatureType.ACTION
        else:
            continue
@@ -34,7 +34,6 @@ python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \
 """

 import argparse
-import logging
 import shutil
 from pathlib import Path
 from typing import Any
@@ -47,6 +46,7 @@ from datasets import Dataset, Features, Image
 from huggingface_hub import HfApi, snapshot_download
 from requests import HTTPError

+from lerobot.constants import HF_LEROBOT_HOME
 from lerobot.datasets.compute_stats import aggregate_stats
 from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
 from lerobot.datasets.utils import (
@@ -71,8 +71,6 @@ from lerobot.datasets.utils import (
    write_tasks,
 )
 from lerobot.datasets.video_utils import concatenate_video_files, get_video_duration_in_s
-from lerobot.utils.constants import HF_LEROBOT_HOME
-from lerobot.utils.utils import init_logging

 V21 = "v2.1"

@@ -146,7 +144,6 @@ def legacy_load_tasks(local_dir: Path) -> tuple[dict, dict]:


 def convert_tasks(root, new_root):
-    logging.info(f"Converting tasks from {root} to {new_root}")
    tasks, _ = legacy_load_tasks(root)
    task_indices = tasks.keys()
    task_strings = tasks.values()
@@ -188,10 +185,7 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):
    num_frames = 0
    paths_to_cat = []
    episodes_metadata = []
-
-    logging.info(f"Converting data files from {len(ep_paths)} episodes")
-
-    for ep_path in tqdm.tqdm(ep_paths, desc="convert data files"):
+    for ep_path in ep_paths:
        ep_size_in_mb = get_parquet_file_size_in_mb(ep_path)
        ep_num_frames = get_parquet_num_frames(ep_path)
        ep_metadata = {
@@ -215,6 +209,7 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int):

        # Reset for the next file
        size_in_mb = ep_size_in_mb
+        num_frames = ep_num_frames
        paths_to_cat = [ep_path]

        chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE)
@@ -241,8 +236,6 @@ def get_image_keys(root):


 def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int):
-    logging.info(f"Converting videos from {root} to {new_root}")
-
    video_keys = get_video_keys(root)
    if len(video_keys) == 0:
        return None
@@ -261,7 +254,7 @@ def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int):
    episods_metadata = []
    num_cameras = len(video_keys)
    num_episodes = num_eps_per_cam[0]
-    for ep_idx in tqdm.tqdm(range(num_episodes), desc="convert videos"):
+    for ep_idx in range(num_episodes):
        # Sanity check
        ep_ids = [eps_metadata_per_cam[cam_idx][ep_idx]["episode_index"] for cam_idx in range(num_cameras)]
        ep_ids += [ep_idx]
@@ -288,7 +281,6 @@ def convert_videos_of_camera(root: Path, new_root: Path, video_key: str, video_f
    duration_in_s = 0.0
    paths_to_cat = []
    episodes_metadata = []
-
    for ep_path in tqdm.tqdm(ep_paths, desc=f"convert videos of {video_key}"):
        ep_size_in_mb = get_video_size_in_mb(ep_path)
        ep_duration_in_s = get_video_duration_in_s(ep_path)
@@ -382,8 +374,6 @@ def generate_episode_metadata_dict(


 def convert_episodes_metadata(root, new_root, episodes_metadata, episodes_video_metadata=None):
-    logging.info(f"Converting episodes metadata from {root} to {new_root}")
-
    episodes_legacy_metadata = legacy_load_episodes(root)
    episodes_stats = legacy_load_episodes_stats(root)

@@ -415,7 +405,6 @@ def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb):
    info["data_path"] = DEFAULT_DATA_PATH
    info["video_path"] = DEFAULT_VIDEO_PATH
    info["fps"] = int(info["fps"])
-    logging.info(f"Converting info from {root} to {new_root}")
    for key in info["features"]:
        if info["features"][key]["dtype"] == "video":
            # already has fps in video_info
@@ -480,7 +469,6 @@ def convert_dataset(


 if __name__ == "__main__":
-    init_logging()
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--repo-id",
@@ -19,9 +19,9 @@ from typing import Any
 import draccus

 from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 from lerobot.robots import RobotConfig
 from lerobot.teleoperators.config import TeleoperatorConfig
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE


@dataclass
@@ -53,12 +53,12 @@ class AlohaEnv(EnvConfig):
    render_mode: str = "rgb_array"
    features: dict[str, PolicyFeature] = field(
        default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(14,)),
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(14,)),
        }
    )
    features_map: dict[str, str] = field(
        default_factory=lambda: {
-            ACTION: ACTION,
+            "action": ACTION,
            "agent_pos": OBS_STATE,
            "top": f"{OBS_IMAGE}.top",
            "pixels/top": f"{OBS_IMAGES}.top",
@@ -93,13 +93,13 @@ class PushtEnv(EnvConfig):
    visualization_height: int = 384
    features: dict[str, PolicyFeature] = field(
        default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(2,)),
            "agent_pos": PolicyFeature(type=FeatureType.STATE, shape=(2,)),
        }
    )
    features_map: dict[str, str] = field(
        default_factory=lambda: {
-            ACTION: ACTION,
+            "action": ACTION,
            "agent_pos": OBS_STATE,
            "environment_state": OBS_ENV_STATE,
            "pixels": OBS_IMAGE,
@@ -135,13 +135,13 @@ class XarmEnv(EnvConfig):
    visualization_height: int = 384
    features: dict[str, PolicyFeature] = field(
        default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(4,)),
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(4,)),
            "pixels": PolicyFeature(type=FeatureType.VISUAL, shape=(84, 84, 3)),
        }
    )
    features_map: dict[str, str] = field(
        default_factory=lambda: {
-            ACTION: ACTION,
+            "action": ACTION,
            "agent_pos": OBS_STATE,
            "pixels": OBS_IMAGE,
        }
@@ -259,12 +259,12 @@ class LiberoEnv(EnvConfig):
    camera_name_mapping: dict[str, str] | None = (None,)
    features: dict[str, PolicyFeature] = field(
        default_factory=lambda: {
-            ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
+            "action": PolicyFeature(type=FeatureType.ACTION, shape=(7,)),
        }
    )
    features_map: dict[str, str] = field(
        default_factory=lambda: {
-            ACTION: ACTION,
+            "action": ACTION,
            "agent_pos": OBS_STATE,
            "pixels/agentview_image": f"{OBS_IMAGES}.image",
            "pixels/robot0_eye_in_hand_image": f"{OBS_IMAGES}.image2",
@@ -26,7 +26,6 @@ from torch import Tensor

 from lerobot.configs.types import FeatureType, PolicyFeature
 from lerobot.envs.configs import EnvConfig
-from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
 from lerobot.utils.utils import get_channel_first_image_shape


@@ -42,9 +41,9 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten
    return_observations = {}
    if "pixels" in observations:
        if isinstance(observations["pixels"], dict):
-            imgs = {f"{OBS_IMAGES}.{key}": img for key, img in observations["pixels"].items()}
+            imgs = {f"observation.images.{key}": img for key, img in observations["pixels"].items()}
        else:
-            imgs = {OBS_IMAGE: observations["pixels"]}
+            imgs = {"observation.image": observations["pixels"]}

        for imgkey, img in imgs.items():
            # TODO(aliberts, rcadene): use transforms.ToTensor()?
@@ -73,13 +72,13 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten
        if env_state.dim() == 1:
            env_state = env_state.unsqueeze(0)

-        return_observations[OBS_ENV_STATE] = env_state
+        return_observations["observation.environment_state"] = env_state

    # TODO(rcadene): enable pixels only baseline with `obs_type="pixels"` in environment by removing
    agent_pos = torch.from_numpy(observations["agent_pos"]).float()
    if agent_pos.dim() == 1:
        agent_pos = agent_pos.unsqueeze(0)
-    return_observations[OBS_STATE] = agent_pos
+    return_observations["observation.state"] = agent_pos

    return return_observations

@@ -24,7 +24,7 @@ lerobot-find-cameras
 ```
 """

-# NOTE(Steven): RealSense can also be identified/opened as OpenCV cameras. If you know the camera is a RealSense, use the `lerobot-find-cameras realsense` flag to avoid confusion.
+# NOTE(Steven): RealSense can also be identified/opened as OpenCV cameras. If you know the camera is a RealSense, use the `lerobot.find_cameras realsense` flag to avoid confusion.
 # NOTE(Steven): macOS cameras sometimes report different FPS at init time, not an issue here as we don't specify FPS when opening the cameras, but the information displayed might not be truthful.

 import argparse
@@ -22,7 +22,7 @@ import logging
 from copy import deepcopy
 from enum import Enum

-from lerobot.motors.encoding_utils import decode_twos_complement, encode_twos_complement
+from lerobot.utils.encoding_utils import decode_twos_complement, encode_twos_complement

 from ..motors_bus import Motor, MotorCalibration, MotorsBus, NameOrID, Value, get_address
 from .tables import (
@@ -17,7 +17,7 @@ from copy import deepcopy
 from enum import Enum
 from pprint import pformat

-from lerobot.motors.encoding_utils import decode_sign_magnitude, encode_sign_magnitude
+from lerobot.utils.encoding_utils import decode_sign_magnitude, encode_sign_magnitude

 from ..motors_bus import Motor, MotorCalibration, MotorsBus, NameOrID, Value, get_address
 from .tables import (
@@ -32,7 +32,7 @@ import serial
 from deepdiff import DeepDiff
 from tqdm import tqdm

-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from lerobot.utils.utils import enter_pressed, move_cursor_up

 NameOrID: TypeAlias = str | int
@@ -22,11 +22,11 @@ import draccus
 import torch
 from safetensors.torch import load_file, save_file

-from lerobot.datasets.utils import flatten_dict, unflatten_dict, write_json
-from lerobot.utils.constants import (
+from lerobot.constants import (
    OPTIMIZER_PARAM_GROUPS,
    OPTIMIZER_STATE,
 )
+from lerobot.datasets.utils import flatten_dict, unflatten_dict, write_json
 from lerobot.utils.io_utils import deserialize_json_into_object


@@ -22,8 +22,8 @@ import draccus
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LambdaLR, LRScheduler

+from lerobot.constants import SCHEDULER_STATE
 from lerobot.datasets.utils import write_json
-from lerobot.utils.constants import SCHEDULER_STATE
 from lerobot.utils.io_utils import deserialize_json_into_object


@@ -33,9 +33,9 @@ from torch import Tensor, nn
 from torchvision.models._utils import IntermediateLayerGetter
 from torchvision.ops.misc import FrozenBatchNorm2d

+from lerobot.constants import ACTION, OBS_IMAGES
 from lerobot.policies.act.configuration_act import ACTConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE


 class ACTPolicy(PreTrainedPolicy):
@@ -394,25 +394,25 @@ class ACT(nn.Module):
            latent dimension.
        """
        if self.config.use_vae and self.training:
-            assert ACTION in batch, (
+            assert "action" in batch, (
                "actions must be provided when using the variational objective in training mode."
            )

-        if OBS_IMAGES in batch:
-            batch_size = batch[OBS_IMAGES][0].shape[0]
+        if "observation.images" in batch:
+            batch_size = batch["observation.images"][0].shape[0]
        else:
-            batch_size = batch[OBS_ENV_STATE].shape[0]
+            batch_size = batch["observation.environment_state"].shape[0]

        # Prepare the latent for input to the transformer encoder.
-        if self.config.use_vae and ACTION in batch and self.training:
+        if self.config.use_vae and "action" in batch and self.training:
            # Prepare the input to the VAE encoder: [cls, *joint_space_configuration, *action_sequence].
            cls_embed = einops.repeat(
                self.vae_encoder_cls_embed.weight, "1 d -> b 1 d", b=batch_size
            )  # (B, 1, D)
            if self.config.robot_state_feature:
-                robot_state_embed = self.vae_encoder_robot_state_input_proj(batch[OBS_STATE])
+                robot_state_embed = self.vae_encoder_robot_state_input_proj(batch["observation.state"])
                robot_state_embed = robot_state_embed.unsqueeze(1)  # (B, 1, D)
-            action_embed = self.vae_encoder_action_input_proj(batch[ACTION])  # (B, S, D)
+            action_embed = self.vae_encoder_action_input_proj(batch["action"])  # (B, S, D)

            if self.config.robot_state_feature:
                vae_encoder_input = [cls_embed, robot_state_embed, action_embed]  # (B, S+2, D)
@@ -430,7 +430,7 @@ class ACT(nn.Module):
            cls_joint_is_pad = torch.full(
                (batch_size, 2 if self.config.robot_state_feature else 1),
                False,
-                device=batch[OBS_STATE].device,
+                device=batch["observation.state"].device,
            )
            key_padding_mask = torch.cat(
                [cls_joint_is_pad, batch["action_is_pad"]], axis=1
@@ -454,7 +454,7 @@ class ACT(nn.Module):
            mu = log_sigma_x2 = None
            # TODO(rcadene, alexander-soare): remove call to `.to` to speedup forward ; precompute and use buffer
            latent_sample = torch.zeros([batch_size, self.config.latent_dim], dtype=torch.float32).to(
-                batch[OBS_STATE].device
+                batch["observation.state"].device
            )

        # Prepare transformer encoder inputs.
@@ -462,16 +462,18 @@ class ACT(nn.Module):
        encoder_in_pos_embed = list(self.encoder_1d_feature_pos_embed.weight.unsqueeze(1))
        # Robot state token.
        if self.config.robot_state_feature:
-            encoder_in_tokens.append(self.encoder_robot_state_input_proj(batch[OBS_STATE]))
+            encoder_in_tokens.append(self.encoder_robot_state_input_proj(batch["observation.state"]))
        # Environment state token.
        if self.config.env_state_feature:
-            encoder_in_tokens.append(self.encoder_env_state_input_proj(batch[OBS_ENV_STATE]))
+            encoder_in_tokens.append(
+                self.encoder_env_state_input_proj(batch["observation.environment_state"])
+            )

        if self.config.image_features:
            # For a list of images, the H and W may vary but H*W is constant.
            # NOTE: If modifying this section, verify on MPS devices that
            # gradients remain stable (no explosions or NaNs).
-            for img in batch[OBS_IMAGES]:
+            for img in batch["observation.images"]:
                cam_features = self.backbone(img)["feature_map"]
                cam_pos_embed = self.encoder_cam_feat_pos_embed(cam_features).to(dtype=cam_features.dtype)
                cam_features = self.encoder_img_feat_input_proj(cam_features)
@@ -17,6 +17,7 @@ from typing import Any

 import torch

+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.act.configuration_act import ACTConfig
 from lerobot.processor import (
    AddBatchDimensionProcessorStep,
@@ -28,7 +29,6 @@ from lerobot.processor import (
    UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME


 def make_act_pre_post_processors(
@@ -33,6 +33,7 @@ from diffusers.schedulers.scheduling_ddim import DDIMScheduler
 from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
 from torch import Tensor, nn

+from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
 from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import (
@@ -41,7 +42,6 @@ from lerobot.policies.utils import (
    get_output_shape,
    populate_queues,
 )
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE


 class DiffusionPolicy(PreTrainedPolicy):
@@ -81,13 +81,13 @@ class DiffusionPolicy(PreTrainedPolicy):
    def reset(self):
        """Clear observation and action queues. Should be called on `env.reset()`"""
        self._queues = {
-            OBS_STATE: deque(maxlen=self.config.n_obs_steps),
-            ACTION: deque(maxlen=self.config.n_action_steps),
+            "observation.state": deque(maxlen=self.config.n_obs_steps),
+            "action": deque(maxlen=self.config.n_action_steps),
        }
        if self.config.image_features:
-            self._queues[OBS_IMAGES] = deque(maxlen=self.config.n_obs_steps)
+            self._queues["observation.images"] = deque(maxlen=self.config.n_obs_steps)
        if self.config.env_state_feature:
-            self._queues[OBS_ENV_STATE] = deque(maxlen=self.config.n_obs_steps)
+            self._queues["observation.environment_state"] = deque(maxlen=self.config.n_obs_steps)

    @torch.no_grad()
    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
@@ -234,7 +234,7 @@ class DiffusionModel(nn.Module):
        if self.config.image_features:
            if self.config.use_separate_rgb_encoder_per_camera:
                # Combine batch and sequence dims while rearranging to make the camera index dimension first.
-                images_per_camera = einops.rearrange(batch[OBS_IMAGES], "b s n ... -> n (b s) ...")
+                images_per_camera = einops.rearrange(batch["observation.images"], "b s n ... -> n (b s) ...")
                img_features_list = torch.cat(
                    [
                        encoder(images)
@@ -249,7 +249,7 @@ class DiffusionModel(nn.Module):
            else:
                # Combine batch, sequence, and "which camera" dims before passing to shared encoder.
                img_features = self.rgb_encoder(
-                    einops.rearrange(batch[OBS_IMAGES], "b s n ... -> (b s n) ...")
+                    einops.rearrange(batch["observation.images"], "b s n ... -> (b s n) ...")
                )
                # Separate batch dim and sequence dim back out. The camera index dim gets absorbed into the
                # feature dim (effectively concatenating the camera features).
@@ -275,7 +275,7 @@ class DiffusionModel(nn.Module):
            "observation.environment_state": (B, n_obs_steps, environment_dim)
        }
        """
-        batch_size, n_obs_steps = batch[OBS_STATE].shape[:2]
+        batch_size, n_obs_steps = batch["observation.state"].shape[:2]
        assert n_obs_steps == self.config.n_obs_steps

        # Encode image features and concatenate them all together along with the state vector.
@@ -306,10 +306,10 @@ class DiffusionModel(nn.Module):
        }
        """
        # Input validation.
-        assert set(batch).issuperset({OBS_STATE, ACTION, "action_is_pad"})
-        assert OBS_IMAGES in batch or OBS_ENV_STATE in batch
-        n_obs_steps = batch[OBS_STATE].shape[1]
-        horizon = batch[ACTION].shape[1]
+        assert set(batch).issuperset({"observation.state", "action", "action_is_pad"})
+        assert "observation.images" in batch or "observation.environment_state" in batch
+        n_obs_steps = batch["observation.state"].shape[1]
+        horizon = batch["action"].shape[1]
        assert horizon == self.config.horizon
        assert n_obs_steps == self.config.n_obs_steps

@@ -317,7 +317,7 @@ class DiffusionModel(nn.Module):
        global_cond = self._prepare_global_conditioning(batch)  # (B, global_cond_dim)

        # Forward diffusion.
-        trajectory = batch[ACTION]
+        trajectory = batch["action"]
        # Sample noise to add to the trajectory.
        eps = torch.randn(trajectory.shape, device=trajectory.device)
        # Sample a random noising timestep for each item in the batch.
@@ -338,7 +338,7 @@ class DiffusionModel(nn.Module):
        if self.config.prediction_type == "epsilon":
            target = eps
        elif self.config.prediction_type == "sample":
-            target = batch[ACTION]
+            target = batch["action"]
        else:
            raise ValueError(f"Unsupported prediction type {self.config.prediction_type}")

@@ -18,6 +18,7 @@ from typing import Any

 import torch

+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
 from lerobot.processor import (
    AddBatchDimensionProcessorStep,
@@ -29,7 +30,6 @@ from lerobot.processor import (
    UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME


 def make_diffusion_pre_post_processors(
@@ -24,6 +24,7 @@ from typing_extensions import Unpack

 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
 from lerobot.datasets.utils import dataset_to_policy_features
 from lerobot.envs.configs import EnvConfig
@@ -45,7 +46,6 @@ from lerobot.processor.converters import (
    transition_to_batch,
    transition_to_policy_action,
 )
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME


 def get_policy_class(name: str) -> type[PreTrainedPolicy]:
@@ -20,7 +20,6 @@ from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import (
    CosineDecayWithWarmupSchedulerConfig,
 )
-from lerobot.utils.constants import OBS_IMAGES


@PreTrainedConfig.register_subclass("pi0")
@@ -114,7 +113,7 @@ class PI0Config(PreTrainedConfig):
        #     raise ValueError("You must provide at least one image or the environment state among the inputs.")

        for i in range(self.empty_cameras):
-            key = f"{OBS_IMAGES}.empty_camera_{i}"
+            key = f"observation.images.empty_camera_{i}"
            empty_camera = PolicyFeature(
                type=FeatureType.VISUAL,
                shape=(3, 480, 640),
@@ -21,7 +21,6 @@ import torch
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
 from lerobot.policies.factory import make_policy
-from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE


 def display(tensor: torch.Tensor):
@@ -61,26 +60,26 @@ def main():

    # Override stats
    dataset_meta = LeRobotDatasetMetadata(dataset_repo_id)
-    dataset_meta.stats[OBS_STATE]["mean"] = torch.tensor(
+    dataset_meta.stats["observation.state"]["mean"] = torch.tensor(
        norm_stats["norm_stats"]["state"]["mean"][:num_motors], dtype=torch.float32
    )
-    dataset_meta.stats[OBS_STATE]["std"] = torch.tensor(
+    dataset_meta.stats["observation.state"]["std"] = torch.tensor(
        norm_stats["norm_stats"]["state"]["std"][:num_motors], dtype=torch.float32
    )

    # Create LeRobot batch from Jax
    batch = {}
    for cam_key, uint_chw_array in example["images"].items():
-        batch[f"{OBS_IMAGES}.{cam_key}"] = torch.from_numpy(uint_chw_array) / 255.0
-    batch[OBS_STATE] = torch.from_numpy(example["state"])
-    batch[ACTION] = torch.from_numpy(outputs["actions"])
+        batch[f"observation.images.{cam_key}"] = torch.from_numpy(uint_chw_array) / 255.0
+    batch["observation.state"] = torch.from_numpy(example["state"])
+    batch["action"] = torch.from_numpy(outputs["actions"])
    batch["task"] = example["prompt"]

    if model_name == "pi0_aloha_towel":
-        del batch[f"{OBS_IMAGES}.cam_low"]
+        del batch["observation.images.cam_low"]
    elif model_name == "pi0_aloha_sim":
-        batch[f"{OBS_IMAGES}.top"] = batch[f"{OBS_IMAGES}.cam_high"]
-        del batch[f"{OBS_IMAGES}.cam_high"]
+        batch["observation.images.top"] = batch["observation.images.cam_high"]
+        del batch["observation.images.cam_high"]

    # Batchify
    for key in batch:
@@ -117,7 +116,7 @@ def main():
        actions.append(action)

    actions = torch.stack(actions, dim=1)
-    pi_actions = batch[ACTION]
+    pi_actions = batch["action"]
    print("actions")
    display(actions)
    print()
@@ -57,13 +57,13 @@ import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn

+from lerobot.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.policies.pi0.paligemma_with_expert import (
    PaliGemmaWithExpertConfig,
    PaliGemmaWithExpertModel,
 )
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.utils.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.utils.utils import get_safe_dtype


@@ -493,7 +493,7 @@ class PI0FlowMatching(nn.Module):
            img_mask,
        ) in zip(images, img_masks, strict=False):
            img_emb = self.paligemma_with_expert.embed_image(img)
-            img_emb = img_emb.to(dtype=torch.bfloat16)
+            img_emb = img_emb.to(dtype=torch.float32)

            # Normalize image embeddings
            img_emb_dim = img_emb.shape[-1]
@@ -536,7 +536,7 @@ class PI0FlowMatching(nn.Module):

        # Embed state
        state_emb = self.state_proj(state)
-        state_emb = state_emb.to(dtype=torch.bfloat16)
+        state_emb = state_emb.to(dtype=torch.float32)
        embs.append(state_emb[:, None, :])
        bsize = state_emb.shape[0]
        dtype = state_emb.dtype
@@ -202,7 +202,7 @@ class PaliGemmaWithExpertModel(PreTrainedModel):
            self.paligemma.eval()

    def to_bfloat16_like_physical_intelligence(self):
-        self.paligemma = self.paligemma.to(dtype=torch.bfloat16)
+        self.paligemma = self.paligemma.to(dtype=torch.float32)

        params_to_change_dtype = [
            "language_model.model.layers",
@@ -212,7 +212,7 @@ class PaliGemmaWithExpertModel(PreTrainedModel):
        ]
        for name, param in self.named_parameters():
            if any(selector in name for selector in params_to_change_dtype):
-                param.data = param.data.to(dtype=torch.bfloat16)
+                param.data = param.data.to(dtype=torch.float32)

    def embed_image(self, image: torch.Tensor):
        # Handle different transformers versions
@@ -262,7 +262,7 @@ class PaliGemmaWithExpertModel(PreTrainedModel):
                input_shape = hidden_states.shape[:-1]
                hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)

-                hidden_states = hidden_states.to(dtype=torch.bfloat16)
+                hidden_states = hidden_states.to(dtype=torch.float32)
                query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape)
                key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape)
                value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape)
@@ -303,7 +303,7 @@ class PaliGemmaWithExpertModel(PreTrainedModel):
            att_output = attention_interface(
                attention_mask, batch_size, head_dim, query_states, key_states, value_states
            )
-            att_output = att_output.to(dtype=torch.bfloat16)
+            att_output = att_output.to(dtype=torch.float32)

            # first part of att_output is prefix (up to sequence length, [:, 0:prefix_seq_len])
            outputs_embeds = []
@@ -19,6 +19,7 @@ from typing import Any
 import torch

 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.pi0.configuration_pi0 import PI0Config
 from lerobot.processor import (
    AddBatchDimensionProcessorStep,
@@ -34,7 +35,6 @@ from lerobot.processor import (
    UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME


@ProcessorStepRegistry.register(name="pi0_new_line_processor")
@@ -6,7 +6,6 @@ from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import (
    CosineDecayWithWarmupSchedulerConfig,
 )
-from lerobot.utils.constants import OBS_IMAGES


@PreTrainedConfig.register_subclass("pi0fast")
@@ -100,7 +99,7 @@ class PI0FASTConfig(PreTrainedConfig):

    def validate_features(self) -> None:
        for i in range(self.empty_cameras):
-            key = f"{OBS_IMAGES}.empty_camera_{i}"
+            key = f"observation.images.empty_camera_{i}"
            empty_camera = PolicyFeature(
                type=FeatureType.VISUAL,
                shape=(3, 480, 640),
@@ -57,9 +57,9 @@ from transformers import AutoProcessor, AutoTokenizer, PaliGemmaForConditionalGe
 from transformers.cache_utils import HybridCache, StaticCache
 from transformers.models.auto import CONFIG_MAPPING

+from lerobot.constants import ACTION, OBS_STATE
 from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.utils.constants import ACTION, OBS_STATE

 PRECISION = {
    "float16": torch.float16,
@@ -18,6 +18,7 @@ from typing import Any

 import torch

+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
 from lerobot.processor import (
    AddBatchDimensionProcessorStep,
@@ -29,7 +30,6 @@ from lerobot.processor import (
    UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME


 def make_pi0fast_pre_post_processors(
@@ -19,8 +19,8 @@ from dataclasses import dataclass, field

 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
+from lerobot.constants import ACTION, OBS_IMAGE, OBS_STATE
 from lerobot.optim.optimizers import MultiAdamConfig
-from lerobot.utils.constants import ACTION, OBS_IMAGE, OBS_STATE


 def is_image_feature(key: str) -> bool:
@@ -225,7 +225,7 @@ class SACConfig(PreTrainedConfig):
                "You must provide either 'observation.state' or an image observation (key starting with 'observation.image') in the input features"
            )

-        if ACTION not in self.output_features:
+        if "action" not in self.output_features:
            raise ValueError("You must provide 'action' in the output features")

    @property
@@ -31,7 +31,6 @@ from torch.distributions import MultivariateNormal, TanhTransform, Transform, Tr
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.configuration_sac import SACConfig, is_image_feature
 from lerobot.policies.utils import get_device_from_parameters
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_STATE

 DISCRETE_DIMENSION_INDEX = -1  # Gripper is always the last dimension

@@ -51,7 +50,7 @@ class SACPolicy(
        self.config = config

        # Determine action dimension and initialize all components
-        continuous_action_dim = config.output_features[ACTION].shape[0]
+        continuous_action_dim = config.output_features["action"].shape[0]
        self._init_encoders()
        self._init_critics(continuous_action_dim)
        self._init_actor(continuous_action_dim)
@@ -158,7 +157,7 @@ class SACPolicy(
            The computed loss tensor
        """
        # Extract common components from batch
-        actions: Tensor = batch[ACTION]
+        actions: Tensor = batch["action"]
        observations: dict[str, Tensor] = batch["state"]
        observation_features: Tensor = batch.get("observation_feature")

@@ -514,17 +513,17 @@ class SACObservationEncoder(nn.Module):
            )

    def _init_state_layers(self) -> None:
-        self.has_env = OBS_ENV_STATE in self.config.input_features
-        self.has_state = OBS_STATE in self.config.input_features
+        self.has_env = "observation.environment_state" in self.config.input_features
+        self.has_state = "observation.state" in self.config.input_features
        if self.has_env:
-            dim = self.config.input_features[OBS_ENV_STATE].shape[0]
+            dim = self.config.input_features["observation.environment_state"].shape[0]
            self.env_encoder = nn.Sequential(
                nn.Linear(dim, self.config.latent_dim),
                nn.LayerNorm(self.config.latent_dim),
                nn.Tanh(),
            )
        if self.has_state:
-            dim = self.config.input_features[OBS_STATE].shape[0]
+            dim = self.config.input_features["observation.state"].shape[0]
            self.state_encoder = nn.Sequential(
                nn.Linear(dim, self.config.latent_dim),
                nn.LayerNorm(self.config.latent_dim),
@@ -550,9 +549,9 @@ class SACObservationEncoder(nn.Module):
                cache = self.get_cached_image_features(obs)
            parts.append(self._encode_images(cache, detach))
        if self.has_env:
-            parts.append(self.env_encoder(obs[OBS_ENV_STATE]))
+            parts.append(self.env_encoder(obs["observation.environment_state"]))
        if self.has_state:
-            parts.append(self.state_encoder(obs[OBS_STATE]))
+            parts.append(self.state_encoder(obs["observation.state"]))
        if parts:
            return torch.cat(parts, dim=-1)

@@ -19,6 +19,7 @@ from typing import Any

 import torch

+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.sac.configuration_sac import SACConfig
 from lerobot.processor import (
    AddBatchDimensionProcessorStep,
@@ -30,7 +31,6 @@ from lerobot.processor import (
    UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME


 def make_sac_pre_post_processors(
@@ -19,7 +19,6 @@ from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
 from lerobot.optim.optimizers import AdamWConfig, OptimizerConfig
 from lerobot.optim.schedulers import LRSchedulerConfig
-from lerobot.utils.constants import OBS_IMAGE


@PreTrainedConfig.register_subclass(name="reward_classifier")
@@ -70,7 +69,7 @@ class RewardClassifierConfig(PreTrainedConfig):

    def validate_features(self) -> None:
        """Validate feature configurations."""
-        has_image = any(key.startswith(OBS_IMAGE) for key in self.input_features)
+        has_image = any(key.startswith("observation.image") for key in self.input_features)
        if not has_image:
            raise ValueError(
                "You must provide an image observation (key starting with 'observation.image') in the input features"
@@ -19,9 +19,9 @@ import logging
 import torch
 from torch import Tensor, nn

+from lerobot.constants import OBS_IMAGE, REWARD
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
-from lerobot.utils.constants import OBS_IMAGE, REWARD


 class ClassifierOutput:
@@ -59,9 +59,7 @@ class SpatialLearnedEmbeddings(nn.Module):
        super().__init__()
        self.height = height
        self.width = width
-        self.channel = (
-            channel  # TODO(fracapuano): this gives issues with non-square images bc is hardcoded to 4
-        )
+        self.channel = channel
        self.num_features = num_features

        self.kernel = nn.Parameter(torch.empty(channel, height, width, num_features))
@@ -20,7 +20,6 @@ from lerobot.optim.optimizers import AdamWConfig
 from lerobot.optim.schedulers import (
    CosineDecayWithWarmupSchedulerConfig,
 )
-from lerobot.utils.constants import OBS_IMAGES


@PreTrainedConfig.register_subclass("smolvla")
@@ -118,7 +117,7 @@ class SmolVLAConfig(PreTrainedConfig):

    def validate_features(self) -> None:
        for i in range(self.empty_cameras):
-            key = f"{OBS_IMAGES}.empty_camera_{i}"
+            key = f"observation.images.empty_camera_{i}"
            empty_camera = PolicyFeature(
                type=FeatureType.VISUAL,
                shape=(3, 480, 640),
@@ -59,13 +59,13 @@ import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn

+from lerobot.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.smolvla.smolvlm_with_expert import SmolVLMWithExpertModel
 from lerobot.policies.utils import (
    populate_queues,
 )
-from lerobot.utils.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.utils.utils import get_safe_dtype


@@ -19,6 +19,7 @@ from typing import Any
 import torch

 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.processor import (
    AddBatchDimensionProcessorStep,
@@ -33,7 +34,6 @@ from lerobot.processor import (
    UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME


 def make_smolvla_pre_post_processors(
@@ -35,10 +35,10 @@ import torch.nn as nn
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor

+from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_STATE, REWARD
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
 from lerobot.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
-from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_PREFIX, OBS_STATE, OBS_STR, REWARD


 class TDMPCPolicy(PreTrainedPolicy):
@@ -91,13 +91,13 @@ class TDMPCPolicy(PreTrainedPolicy):
        called on `env.reset()`
        """
        self._queues = {
-            OBS_STATE: deque(maxlen=1),
-            ACTION: deque(maxlen=max(self.config.n_action_steps, self.config.n_action_repeats)),
+            "observation.state": deque(maxlen=1),
+            "action": deque(maxlen=max(self.config.n_action_steps, self.config.n_action_repeats)),
        }
        if self.config.image_features:
-            self._queues[OBS_IMAGE] = deque(maxlen=1)
+            self._queues["observation.image"] = deque(maxlen=1)
        if self.config.env_state_feature:
-            self._queues[OBS_ENV_STATE] = deque(maxlen=1)
+            self._queues["observation.environment_state"] = deque(maxlen=1)
        # Previous mean obtained from the cross-entropy method (CEM) used during MPC. It is used to warm start
        # CEM for the next step.
        self._prev_mean: torch.Tensor | None = None
@@ -325,7 +325,7 @@ class TDMPCPolicy(PreTrainedPolicy):

        action = batch[ACTION]  # (t, b, action_dim)
        reward = batch[REWARD]  # (t, b)
-        observations = {k: v for k, v in batch.items() if k.startswith(OBS_PREFIX)}
+        observations = {k: v for k, v in batch.items() if k.startswith("observation.")}

        # Apply random image augmentations.
        if self.config.image_features and self.config.max_random_shift_ratio > 0:
@@ -387,10 +387,10 @@ class TDMPCPolicy(PreTrainedPolicy):
                temporal_loss_coeffs
                * F.mse_loss(z_preds[1:], z_targets, reduction="none").mean(dim=-1)
                # `z_preds` depends on the current observation and the actions.
-                * ~batch[f"{OBS_STR}.state_is_pad"][0]
+                * ~batch["observation.state_is_pad"][0]
                * ~batch["action_is_pad"]
                # `z_targets` depends on the next observation.
-                * ~batch[f"{OBS_STR}.state_is_pad"][1:]
+                * ~batch["observation.state_is_pad"][1:]
            )
            .sum(0)
            .mean()
@@ -403,7 +403,7 @@ class TDMPCPolicy(PreTrainedPolicy):
                * F.mse_loss(reward_preds, reward, reduction="none")
                * ~batch["next.reward_is_pad"]
                # `reward_preds` depends on the current observation and the actions.
-                * ~batch[f"{OBS_STR}.state_is_pad"][0]
+                * ~batch["observation.state_is_pad"][0]
                * ~batch["action_is_pad"]
            )
            .sum(0)
@@ -419,11 +419,11 @@ class TDMPCPolicy(PreTrainedPolicy):
                    reduction="none",
                ).sum(0)  # sum over ensemble
                # `q_preds_ensemble` depends on the first observation and the actions.
-                * ~batch[f"{OBS_STR}.state_is_pad"][0]
+                * ~batch["observation.state_is_pad"][0]
                * ~batch["action_is_pad"]
                # q_targets depends on the reward and the next observations.
                * ~batch["next.reward_is_pad"]
-                * ~batch[f"{OBS_STR}.state_is_pad"][1:]
+                * ~batch["observation.state_is_pad"][1:]
            )
            .sum(0)
            .mean()
@@ -441,7 +441,7 @@ class TDMPCPolicy(PreTrainedPolicy):
                temporal_loss_coeffs
                * raw_v_value_loss
                # `v_targets` depends on the first observation and the actions, as does `v_preds`.
-                * ~batch[f"{OBS_STR}.state_is_pad"][0]
+                * ~batch["observation.state_is_pad"][0]
                * ~batch["action_is_pad"]
            )
            .sum(0)
@@ -477,7 +477,7 @@ class TDMPCPolicy(PreTrainedPolicy):
            * mse
            * temporal_loss_coeffs
            # `action_preds` depends on the first observation and the actions.
-            * ~batch[f"{OBS_STR}.state_is_pad"][0]
+            * ~batch["observation.state_is_pad"][0]
            * ~batch["action_is_pad"]
        ).mean()

@@ -18,6 +18,7 @@ from typing import Any

 import torch

+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
 from lerobot.processor import (
    AddBatchDimensionProcessorStep,
@@ -29,7 +30,6 @@ from lerobot.processor import (
    UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME


 def make_tdmpc_pre_post_processors(
@@ -27,11 +27,11 @@ import torch.nn.functional as F  # noqa: N812
 import torchvision
 from torch import Tensor, nn

+from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
 from lerobot.policies.vqbet.vqbet_utils import GPT, ResidualVQ
-from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE

 # ruff: noqa: N806

@@ -133,7 +133,7 @@ class VQBeTPolicy(PreTrainedPolicy):
            batch.pop(ACTION)
        batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
        # NOTE: It's important that this happens after stacking the images into a single key.
-        batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
+        batch["observation.images"] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
        # NOTE: for offline evaluation, we have action in the batch, so we need to pop it out
        if ACTION in batch:
            batch.pop(ACTION)
@@ -340,12 +340,14 @@ class VQBeTModel(nn.Module):

    def forward(self, batch: dict[str, Tensor], rollout: bool) -> tuple[dict, dict]:
        # Input validation.
-        assert set(batch).issuperset({OBS_STATE, OBS_IMAGES})
-        batch_size, n_obs_steps = batch[OBS_STATE].shape[:2]
+        assert set(batch).issuperset({"observation.state", "observation.images"})
+        batch_size, n_obs_steps = batch["observation.state"].shape[:2]
        assert n_obs_steps == self.config.n_obs_steps

        # Extract image feature (first combine batch and sequence dims).
-        img_features = self.rgb_encoder(einops.rearrange(batch[OBS_IMAGES], "b s n ... -> (b s n) ..."))
+        img_features = self.rgb_encoder(
+            einops.rearrange(batch["observation.images"], "b s n ... -> (b s n) ...")
+        )
        # Separate batch and sequence dims.
        img_features = einops.rearrange(
            img_features, "(b s n) ... -> b s n ...", b=batch_size, s=n_obs_steps, n=self.num_images
@@ -357,7 +359,9 @@ class VQBeTModel(nn.Module):
            img_features
        )  # (batch, obs_step, number of different cameras, projection dims)
        input_tokens = [rgb_tokens[:, :, i] for i in range(rgb_tokens.size(2))]
-        input_tokens.append(self.state_projector(batch[OBS_STATE]))  # (batch, obs_step, projection dims)
+        input_tokens.append(
+            self.state_projector(batch["observation.state"])
+        )  # (batch, obs_step, projection dims)
        input_tokens.append(einops.repeat(self.action_token, "1 1 d -> b n d", b=batch_size, n=n_obs_steps))
        # Interleave tokens by stacking and rearranging.
        input_tokens = torch.stack(input_tokens, dim=2)
@@ -19,6 +19,7 @@ from typing import Any

 import torch

+from lerobot.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
 from lerobot.processor import (
    AddBatchDimensionProcessorStep,
@@ -30,7 +31,6 @@ from lerobot.processor import (
    UnnormalizerProcessorStep,
 )
 from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME


 def make_vqbet_pre_post_processors(
@@ -25,7 +25,7 @@ from dataclasses import dataclass, field
 from torch import Tensor

 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
-from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
+from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE

 from .core import EnvTransition, PolicyAction
 from .pipeline import (
@@ -23,8 +23,6 @@ from typing import Any
 import numpy as np
 import torch

-from lerobot.utils.constants import ACTION, DONE, OBS_PREFIX, REWARD, TRUNCATED
-
 from .core import EnvTransition, PolicyAction, RobotAction, RobotObservation, TransitionKey


@@ -344,20 +342,20 @@ def batch_to_transition(batch: dict[str, Any]) -> EnvTransition:
    if not isinstance(batch, dict):
        raise ValueError(f"EnvTransition must be a dictionary. Got {type(batch).__name__}")

-    action = batch.get(ACTION)
+    action = batch.get("action")
    if action is not None and not isinstance(action, PolicyAction):
        raise ValueError(f"Action should be a PolicyAction type got {type(action)}")

    # Extract observation and complementary data keys.
-    observation_keys = {k: v for k, v in batch.items() if k.startswith(OBS_PREFIX)}
+    observation_keys = {k: v for k, v in batch.items() if k.startswith("observation.")}
    complementary_data = _extract_complementary_data(batch)

    return create_transition(
        observation=observation_keys if observation_keys else None,
-        action=batch.get(ACTION),
-        reward=batch.get(REWARD, 0.0),
-        done=batch.get(DONE, False),
-        truncated=batch.get(TRUNCATED, False),
+        action=batch.get("action"),
+        reward=batch.get("next.reward", 0.0),
+        done=batch.get("next.done", False),
+        truncated=batch.get("next.truncated", False),
        info=batch.get("info", {}),
        complementary_data=complementary_data if complementary_data else None,
    )
@@ -379,10 +377,10 @@ def transition_to_batch(transition: EnvTransition) -> dict[str, Any]:
        raise ValueError(f"Transition should be a EnvTransition type (dict) got {type(transition)}")

    batch = {
-        ACTION: transition.get(TransitionKey.ACTION),
-        REWARD: transition.get(TransitionKey.REWARD, 0.0),
-        DONE: transition.get(TransitionKey.DONE, False),
-        TRUNCATED: transition.get(TransitionKey.TRUNCATED, False),
+        "action": transition.get(TransitionKey.ACTION),
+        "next.reward": transition.get(TransitionKey.REWARD, 0.0),
+        "next.done": transition.get(TransitionKey.DONE, False),
+        "next.truncated": transition.get(TransitionKey.TRUNCATED, False),
        "info": transition.get(TransitionKey.INFO, {}),
    }

@@ -20,12 +20,12 @@ from typing import Any
 import torch

 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
+from lerobot.constants import OBS_STATE
 from lerobot.processor.pipeline import (
    ObservationProcessorStep,
    ProcessorStepRegistry,
 )
 from lerobot.robots import Robot
-from lerobot.utils.constants import OBS_STATE


@dataclass
@@ -59,7 +59,6 @@ from safetensors.torch import load_file as load_safetensors

 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.policies.factory import get_policy_class, make_policy_config, make_pre_post_processors
-from lerobot.utils.constants import ACTION


 def extract_normalization_stats(state_dict: dict[str, torch.Tensor]) -> dict[str, dict[str, torch.Tensor]]:
@@ -197,7 +196,7 @@ def detect_features_and_norm_modes(
                feature_type = FeatureType.VISUAL
            elif "state" in key:
                feature_type = FeatureType.STATE
-            elif ACTION in key:
+            elif "action" in key:
                feature_type = FeatureType.ACTION
            else:
                feature_type = FeatureType.STATE  # Default
@@ -216,7 +215,7 @@ def detect_features_and_norm_modes(
                feature_type = FeatureType.VISUAL
            elif "state" in key or "joint" in key or "position" in key:
                feature_type = FeatureType.STATE
-            elif ACTION in key:
+            elif "action" in key:
                feature_type = FeatureType.ACTION
            else:
                feature_type = FeatureType.STATE
@@ -322,7 +321,7 @@ def convert_features_to_policy_features(features_dict: dict[str, dict]) -> dict[
            feature_type = FeatureType.VISUAL
        elif "state" in key:
            feature_type = FeatureType.STATE
-        elif ACTION in key:
+        elif "action" in key:
            feature_type = FeatureType.ACTION
        else:
            feature_type = FeatureType.STATE
@@ -26,7 +26,6 @@ from torch import Tensor

 from lerobot.configs.types import FeatureType, NormalizationMode, PipelineFeatureType, PolicyFeature
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.utils.constants import ACTION

 from .converters import from_tensor_to_numpy, to_tensor
 from .core import EnvTransition, PolicyAction, TransitionKey
@@ -273,7 +272,7 @@ class _NormalizationMixin:
        Returns:
            The transformed action tensor.
        """
-        processed_action = self._apply_transform(action, ACTION, FeatureType.ACTION, inverse=inverse)
+        processed_action = self._apply_transform(action, "action", FeatureType.ACTION, inverse=inverse)
        return processed_action

    def _apply_transform(
@@ -21,7 +21,7 @@ import torch
 from torch import Tensor

 from lerobot.configs.types import PipelineFeatureType, PolicyFeature
-from lerobot.utils.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE, OBS_STR
+from lerobot.constants import OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE

 from .pipeline import ObservationProcessorStep, ProcessorStepRegistry

@@ -171,7 +171,7 @@ class VanillaObservationProcessorStep(ObservationProcessorStep):

                # Prefix-based rules (e.g. pixels.cam1 -> OBS_IMAGES.cam1)
                for old_prefix, new_prefix in prefix_pairs.items():
-                    prefixed_old = f"{OBS_STR}.{old_prefix}"
+                    prefixed_old = f"observation.{old_prefix}"
                    if key.startswith(prefixed_old):
                        suffix = key[len(prefixed_old) :]
                        new_key = f"{new_prefix}{suffix}"
@@ -191,7 +191,7 @@ class VanillaObservationProcessorStep(ObservationProcessorStep):

                # Exact-name rules (pixels, environment_state, agent_pos)
                for old, new in exact_pairs.items():
-                    if key == old or key == f"{OBS_STR}.{old}":
+                    if key == old or key == f"observation.{old}":
                        new_key = new
                        new_features[src_ft][new_key] = feat
                        handled = True
@@ -422,7 +422,7 @@ class DataProcessorPipeline(HubMixin, Generic[TInput, TOutput]):
        """
        if save_directory is None:
            # Use default directory in HF_LEROBOT_HOME
-            from lerobot.utils.constants import HF_LEROBOT_HOME
+            from lerobot.constants import HF_LEROBOT_HOME

            sanitized_name = re.sub(r"[^a-zA-Z0-9_]", "_", self.name.lower())
            save_directory = HF_LEROBOT_HOME / "processors" / sanitized_name
@@ -5,7 +5,6 @@ import torch

 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
 from lerobot.processor import ActionProcessorStep, PolicyAction, ProcessorStepRegistry, RobotAction
-from lerobot.utils.constants import ACTION


@dataclass
@@ -24,7 +23,7 @@ class RobotActionToPolicyActionProcessorStep(ActionProcessorStep):
        return asdict(self)

    def transform_features(self, features):
-        features[PipelineFeatureType.ACTION][ACTION] = PolicyFeature(
+        features[PipelineFeatureType.ACTION]["action"] = PolicyFeature(
            type=FeatureType.ACTION, shape=(len(self.motor_names),)
        )
        return features
@@ -29,7 +29,7 @@ from typing import TYPE_CHECKING, Any
 import torch

 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
+from lerobot.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS
 from lerobot.utils.import_utils import _transformers_available

 from .core import EnvTransition, TransitionKey
@@ -109,7 +109,6 @@ from lerobot.teleoperators import (  # noqa: F401
    so101_leader,
 )
 from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop
-from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.control_utils import (
    init_keyboard_listener,
    is_headless,
@@ -123,7 +122,7 @@ from lerobot.utils.utils import (
    init_logging,
    log_say,
 )
-from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
+from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data


@dataclass
@@ -304,7 +303,7 @@ def record_loop(
        obs_processed = robot_observation_processor(obs)

        if policy is not None or dataset is not None:
-            observation_frame = build_dataset_frame(dataset.features, obs_processed, prefix=OBS_STR)
+            observation_frame = build_dataset_frame(dataset.features, obs_processed, prefix="observation")

        # Get action from either policy or teleop
        if policy is not None and preprocessor is not None and postprocessor is not None:
@@ -319,7 +318,7 @@ def record_loop(
                robot_type=robot.robot_type,
            )

-            action_names = dataset.features[ACTION]["names"]
+            action_names = dataset.features["action"]["names"]
            act_processed_policy: RobotAction = {
                f"{name}": float(action_values[i]) for i, name in enumerate(action_names)
            }
@@ -361,7 +360,7 @@ def record_loop(

        # Write to dataset
        if dataset is not None:
-            action_frame = build_dataset_frame(dataset.features, action_values, prefix=ACTION)
+            action_frame = build_dataset_frame(dataset.features, action_values, prefix="action")
            frame = {**observation_frame, **action_frame, "task": single_task}
            dataset.add_frame(frame)

@@ -379,7 +378,7 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
    init_logging()
    logging.info(pformat(asdict(cfg)))
    if cfg.display_data:
-        init_rerun(session_name="recording")
+        _init_rerun(session_name="recording")

    robot = make_robot_from_config(cfg.robot)
    teleop = make_teleoperator_from_config(cfg.teleop) if cfg.teleop is not None else None
@@ -60,7 +60,6 @@ from lerobot.robots import (  # noqa: F401
    so100_follower,
    so101_follower,
 )
-from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import (
    init_logging,
@@ -100,7 +99,7 @@ def replay(cfg: ReplayConfig):

    # Filter dataset to only include frames from the specified episode since episodes are chunked in dataset V3.0
    episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == cfg.dataset.episode)
-    actions = episode_frames.select_columns(ACTION)
+    actions = episode_frames.select_columns("action")

    robot.connect()

@@ -108,9 +107,9 @@ def replay(cfg: ReplayConfig):
    for idx in range(len(episode_frames)):
        start_episode_t = time.perf_counter()

-        action_array = actions[idx][ACTION]
+        action_array = actions[idx]["action"]
        action = {}
-        for i, name in enumerate(dataset.features[ACTION]["names"]):
+        for i, name in enumerate(dataset.features["action"]["names"]):
            action[name] = action_array[i]

        robot_obs = robot.get_observation()
@@ -35,7 +35,7 @@ gamepad to take control of the robot during training. Initially intervene freque
 reduce interventions as the policy improves.

 **WORKFLOW**:
-1. Determine robot workspace bounds using `lerobot-find-joint-limits`
+1. Determine robot workspace bounds using `find_joint_limits.py`
 2. Record demonstrations with `gym_manipulator.py` in record mode
 3. Process the dataset and determine camera crops with `crop_dataset_roi.py`
 4. Start the learner server with the training configuration
@@ -63,8 +63,6 @@ from lerobot.configs.train import TrainRLServerPipelineConfig
 from lerobot.policies.factory import make_policy
 from lerobot.policies.sac.modeling_sac import SACPolicy
 from lerobot.processor import TransitionKey
-from lerobot.rl.process import ProcessSignalHandler
-from lerobot.rl.queue import get_last_item_from_queue
 from lerobot.robots import so100_follower  # noqa: F401
 from lerobot.teleoperators import gamepad, so101_leader  # noqa: F401
 from lerobot.teleoperators.utils import TeleopEvents
@@ -77,6 +75,8 @@ from lerobot.transport.utils import (
    send_bytes_in_chunks,
    transitions_to_bytes,
 )
+from lerobot.utils.process import ProcessSignalHandler
+from lerobot.utils.queue import get_last_item_from_queue
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.transition import (
@@ -25,7 +25,6 @@ import torchvision.transforms.functional as F  # type: ignore  # noqa: N812
 from tqdm import tqdm  # type: ignore

 from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.utils.constants import DONE, REWARD


 def select_rect_roi(img):
@@ -213,7 +212,7 @@ def convert_lerobot_dataset_to_cropper_lerobot_dataset(
        for key, value in frame.items():
            if key in ("task_index", "timestamp", "episode_index", "frame_index", "index", "task"):
                continue
-            if key in (DONE, REWARD):
+            if key in ("next.done", "next.reward"):
                # if not isinstance(value, str) and len(value.shape) == 0:
                value = value.unsqueeze(0)

@@ -73,7 +73,6 @@ from lerobot.teleoperators import (
 )
 from lerobot.teleoperators.teleoperator import Teleoperator
 from lerobot.teleoperators.utils import TeleopEvents
-from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, REWARD
 from lerobot.utils.robot_utils import busy_wait
 from lerobot.utils.utils import log_say

@@ -181,7 +180,7 @@ class RobotEnv(gym.Env):

        # Define observation spaces for images and other states.
        if current_observation is not None and "pixels" in current_observation:
-            prefix = OBS_IMAGES
+            prefix = "observation.images"
            observation_spaces = {
                f"{prefix}.{key}": gym.spaces.Box(
                    low=0, high=255, shape=current_observation["pixels"][key].shape, dtype=np.uint8
@@ -191,7 +190,7 @@ class RobotEnv(gym.Env):

        if current_observation is not None:
            agent_pos = current_observation["agent_pos"]
-            observation_spaces[OBS_STATE] = gym.spaces.Box(
+            observation_spaces["observation.state"] = gym.spaces.Box(
                low=0,
                high=10,
                shape=agent_pos.shape,
@@ -601,9 +600,9 @@ def control_loop(
    if cfg.mode == "record":
        action_features = teleop_device.action_features
        features = {
-            ACTION: action_features,
-            REWARD: {"dtype": "float32", "shape": (1,), "names": None},
-            DONE: {"dtype": "bool", "shape": (1,), "names": None},
+            "action": action_features,
+            "next.reward": {"dtype": "float32", "shape": (1,), "names": None},
+            "next.done": {"dtype": "bool", "shape": (1,), "names": None},
        }
        if use_gripper:
            features["complementary_info.discrete_penalty"] = {
@@ -613,7 +612,7 @@ def control_loop(
            }

        for key, value in transition[TransitionKey.OBSERVATION].items():
-            if key == OBS_STATE:
+            if key == "observation.state":
                features[key] = {
                    "dtype": "float32",
                    "shape": value.squeeze(0).shape,
@@ -672,9 +671,9 @@ def control_loop(
            )
            frame = {
                **observations,
-                ACTION: action_to_record.cpu(),
-                REWARD: np.array([transition[TransitionKey.REWARD]], dtype=np.float32),
-                DONE: np.array([terminated or truncated], dtype=bool),
+                "action": action_to_record.cpu(),
+                "next.reward": np.array([transition[TransitionKey.REWARD]], dtype=np.float32),
+                "next.done": np.array([terminated or truncated], dtype=bool),
            }
            if use_gripper:
                discrete_penalty = transition[TransitionKey.COMPLEMENTARY_DATA].get("discrete_penalty", 0.0)
@@ -733,7 +732,7 @@ def replay_trajectory(
        download_videos=False,
    )
    episode_frames = dataset.hf_dataset.filter(lambda x: x["episode_index"] == cfg.dataset.replay_episode)
-    actions = episode_frames.select_columns(ACTION)
+    actions = episode_frames.select_columns("action")

    _, info = env.reset()

@@ -741,7 +740,7 @@ def replay_trajectory(
        start_time = time.perf_counter()
        transition = create_transition(
            observation=env.get_raw_joint_positions() if hasattr(env, "get_raw_joint_positions") else {},
-            action=action_data[ACTION],
+            action=action_data["action"],
        )
        transition = action_processor(transition)
        env.step(transition[TransitionKey.ACTION])
@@ -62,13 +62,16 @@ from torch.optim.optimizer import Optimizer
 from lerobot.cameras import opencv  # noqa: F401
 from lerobot.configs import parser
 from lerobot.configs.train import TrainRLServerPipelineConfig
+from lerobot.constants import (
+    CHECKPOINTS_DIR,
+    LAST_CHECKPOINT_LINK,
+    PRETRAINED_MODEL_DIR,
+    TRAINING_STATE_DIR,
+)
 from lerobot.datasets.factory import make_dataset
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
 from lerobot.policies.factory import make_policy
 from lerobot.policies.sac.modeling_sac import SACPolicy
-from lerobot.rl.buffer import ReplayBuffer, concatenate_batch_transitions
-from lerobot.rl.process import ProcessSignalHandler
-from lerobot.rl.wandb_utils import WandBLogger
 from lerobot.robots import so100_follower  # noqa: F401
 from lerobot.teleoperators import gamepad, so101_leader  # noqa: F401
 from lerobot.teleoperators.utils import TeleopEvents
@@ -79,13 +82,8 @@ from lerobot.transport.utils import (
    bytes_to_transitions,
    state_to_bytes,
 )
-from lerobot.utils.constants import (
-    ACTION,
-    CHECKPOINTS_DIR,
-    LAST_CHECKPOINT_LINK,
-    PRETRAINED_MODEL_DIR,
-    TRAINING_STATE_DIR,
-)
+from lerobot.utils.buffer import ReplayBuffer, concatenate_batch_transitions
+from lerobot.utils.process import ProcessSignalHandler
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.train_utils import (
    get_step_checkpoint_dir,
@@ -99,6 +97,7 @@ from lerobot.utils.utils import (
    get_safe_torch_device,
    init_logging,
 )
+from lerobot.utils.wandb_utils import WandBLogger

 from .learner_service import MAX_WORKERS, SHUTDOWN_TIMEOUT, LearnerService

@@ -154,7 +153,7 @@ def train(cfg: TrainRLServerPipelineConfig, job_name: str | None = None):

    # Setup WandB logging if enabled
    if cfg.wandb.enable and cfg.wandb.project:
-        from lerobot.rl.wandb_utils import WandBLogger
+        from lerobot.utils.wandb_utils import WandBLogger

        wandb_logger = WandBLogger(cfg)
    else:
@@ -403,7 +402,7 @@ def add_actor_information_and_train(
                    left_batch_transitions=batch, right_batch_transition=batch_offline
                )

-            actions = batch[ACTION]
+            actions = batch["action"]
            rewards = batch["reward"]
            observations = batch["state"]
            next_observations = batch["next_state"]
@@ -416,7 +415,7 @@ def add_actor_information_and_train(

            # Create a batch dictionary with all required elements for the forward method
            forward_batch = {
-                ACTION: actions,
+                "action": actions,
                "reward": rewards,
                "state": observations,
                "next_state": next_observations,
@@ -461,7 +460,7 @@ def add_actor_information_and_train(
                left_batch_transitions=batch, right_batch_transition=batch_offline
            )

-        actions = batch[ACTION]
+        actions = batch["action"]
        rewards = batch["reward"]
        observations = batch["state"]
        next_observations = batch["next_state"]
@@ -475,7 +474,7 @@ def add_actor_information_and_train(

        # Create a batch dictionary with all required elements for the forward method
        forward_batch = {
-            ACTION: actions,
+            "action": actions,
            "reward": rewards,
            "state": observations,
            "next_state": next_observations,
@@ -1156,7 +1155,7 @@ def process_transitions(
            # Skip transitions with NaN values
            if check_nan_in_transition(
                observations=transition["state"],
-                actions=transition[ACTION],
+                actions=transition["action"],
                next_state=transition["next_state"],
            ):
                logging.warning("[LEARNER] NaN detected in transition, skipping")
@@ -19,9 +19,9 @@ import logging
 import time
 from multiprocessing import Event, Queue

-from lerobot.rl.queue import get_last_item_from_queue
 from lerobot.transport import services_pb2, services_pb2_grpc
 from lerobot.transport.utils import receive_bytes_in_chunks, send_bytes_in_chunks
+from lerobot.utils.queue import get_last_item_from_queue

 MAX_WORKERS = 3  # Stream parameters, send transitions and interactions
 SHUTDOWN_TIMEOUT = 10
@@ -20,12 +20,12 @@ from functools import cached_property
 from typing import Any

 from lerobot.cameras.utils import make_cameras_from_configs
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from lerobot.motors import Motor, MotorNormMode
 from lerobot.motors.calibration_gui import RangeFinderGUI
 from lerobot.motors.feetech import (
    FeetechMotorsBus,
 )
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError

 from ..robot import Robot
 from ..utils import ensure_safe_goal_position
@@ -20,12 +20,12 @@ from functools import cached_property
 from typing import Any

 from lerobot.cameras.utils import make_cameras_from_configs
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from lerobot.motors import Motor, MotorNormMode
 from lerobot.motors.calibration_gui import RangeFinderGUI
 from lerobot.motors.feetech import (
    FeetechMotorsBus,
 )
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError

 from ..robot import Robot
 from .config_hope_jr import HopeJrHandConfig
@@ -20,12 +20,12 @@ from functools import cached_property
 from typing import Any

 from lerobot.cameras.utils import make_cameras_from_configs
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.dynamixel import (
    DynamixelMotorsBus,
    OperatingMode,
 )
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError

 from ..robot import Robot
 from ..utils import ensure_safe_goal_position
@@ -23,12 +23,12 @@ from typing import Any
 import numpy as np

 from lerobot.cameras.utils import make_cameras_from_configs
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.feetech import (
    FeetechMotorsBus,
    OperatingMode,
 )
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError

 from ..robot import Robot
 from ..utils import ensure_safe_goal_position
@@ -23,8 +23,7 @@ from typing import Any
 import cv2
 import numpy as np

-from lerobot.utils.constants import ACTION, OBS_STATE
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError

 from ..robot import Robot
 from .config_lekiwi import LeKiwiClientConfig
@@ -204,7 +203,7 @@ class LeKiwiClient(Robot):

        state_vec = np.array([flat_state[key] for key in self._state_order], dtype=np.float32)

-        obs_dict: dict[str, Any] = {**flat_state, OBS_STATE: state_vec}
+        obs_dict: dict[str, Any] = {**flat_state, "observation.state": state_vec}

        # Decode images
        current_frames: dict[str, np.ndarray] = {}
@@ -330,7 +329,7 @@ class LeKiwiClient(Robot):
        actions = np.array([action.get(k, 0.0) for k in self._state_order], dtype=np.float32)

        action_sent = {key: actions[i] for i, key in enumerate(self._state_order)}
-        action_sent[ACTION] = actions
+        action_sent["action"] = actions
        return action_sent

    def disconnect(self):
@@ -19,8 +19,8 @@ from typing import Any

 import draccus

+from lerobot.constants import HF_LEROBOT_CALIBRATION, ROBOTS
 from lerobot.motors import MotorCalibration
-from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, ROBOTS

 from .config import RobotConfig

@@ -20,12 +20,12 @@ from functools import cached_property
 from typing import Any

 from lerobot.cameras.utils import make_cameras_from_configs
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.feetech import (
    FeetechMotorsBus,
    OperatingMode,
 )
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError

 from ..robot import Robot
 from ..utils import ensure_safe_goal_position
@@ -20,12 +20,12 @@ from functools import cached_property
 from typing import Any

 from lerobot.cameras.utils import make_cameras_from_configs
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.feetech import (
    FeetechMotorsBus,
    OperatingMode,
 )
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError

 from ..robot import Robot
 from ..utils import ensure_safe_goal_position
@@ -22,8 +22,8 @@ from stretch_body.robot import Robot as StretchAPI
 from stretch_body.robot_params import RobotParams

 from lerobot.cameras.utils import make_cameras_from_configs
+from lerobot.constants import OBS_IMAGES, OBS_STATE
 from lerobot.datasets.utils import get_nested_item
-from lerobot.utils.constants import OBS_IMAGES, OBS_STATE

 from ..robot import Robot
 from .configuration_stretch3 import Stretch3RobotConfig
@@ -18,13 +18,13 @@ from functools import cached_property
 from typing import Any

 from lerobot.cameras.utils import make_cameras_from_configs
+from lerobot.constants import OBS_STATE
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 from lerobot.motors import Motor, MotorCalibration, MotorNormMode
 from lerobot.motors.dynamixel import (
    DynamixelMotorsBus,
    OperatingMode,
 )
-from lerobot.utils.constants import OBS_STATE
-from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError

 from ..robot import Robot
 from ..utils import ensure_safe_goal_position
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
AdilZouitine	6986b46fa8	refactor(inference): improve timeout handling and report timeout percentage - Commented out the timeout handling logic to prevent appending timeout values to the results. - Added a print statement to display the percentage of timeouts during inference.	2025-09-24 14:50:58 +02:00
AdilZouitine	b1d72ac29c	refactor(model): change tensor data type from bfloat16 to float32 - Updated image and state embeddings to use float32 for improved compatibility. - Adjusted model parameters and hidden states to ensure consistent data type usage.	2025-09-24 14:33:11 +02:00
AdilZouitine	cffd545527	refactor(inference): improve timeout handling and enhance dummy observation generation - Renamed TimeoutException to TimeoutExceptionError for clarity. - Updated dummy observation generation to include a task string. - Integrated pre-processing and post-processing steps in the main function. - Added deep copy of dummy observations to prevent mutation during processing. - Enhanced timeout handling to provide percentage of timeouts during inference.	2025-09-24 14:32:47 +02:00
Francesco Capuano	6eaf6a861a	fix: single level loop	2025-09-24 01:06:13 +02:00
Francesco Capuano	cdd6cb606c	add: inference benchmark	2025-09-23 22:34:52 +02:00
Jade Choghari	f6cd24be17	update Signed-off-by: Jade Choghari <chogharijade@gmail.com>	2025-09-23 21:52:15 +02:00
Jade Choghari	54c6b8ae52	add file Signed-off-by: Jade Choghari <chogharijade@gmail.com>	2025-09-23 21:52:14 +02:00