diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
index 2e6fd4490..df2e2db29 100644
--- a/.github/PULL_REQUEST_TEMPLATE.md
+++ b/.github/PULL_REQUEST_TEMPLATE.md
@@ -24,7 +24,7 @@ Examples:
 pytest -sx tests/test_stuff.py::test_something
 ```
 ```bash
-python lerobot/scripts/train.py --some.option=true
+python -m lerobot.scripts.train --some.option=true
 ```
 
 ## SECTION TO REMOVE BEFORE SUBMITTING YOUR PR
diff --git a/.github/workflows/nightly-tests.yml b/.github/workflows/nightly-tests.yml
index be248b335..728016915 100644
--- a/.github/workflows/nightly-tests.yml
+++ b/.github/workflows/nightly-tests.yml
@@ -44,7 +44,7 @@ jobs:
         working-directory: /lerobot
     steps:
       - name: Tests
-        run: pytest -v --cov=./lerobot --disable-warnings tests
+        run: pytest -v --cov=./src/lerobot --disable-warnings tests
 
       - name: Tests end-to-end
         run: make test-end-to-end
@@ -74,7 +74,7 @@ jobs:
         run: nvidia-smi
 
       - name: Test
-        run: pytest -v --cov=./lerobot --cov-report=xml --disable-warnings tests
+        run: pytest -v --cov=./src/lerobot --cov-report=xml --disable-warnings tests
       #   TODO(aliberts): Link with HF Codecov account
       # - name: Upload coverage reports to Codecov with GitHub Action
       #   uses: codecov/codecov-action@v4
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 8822956cf..d6ea1d404 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -17,7 +17,7 @@ name: Tests
 on:
   pull_request:
     paths:
-      - "lerobot/**"
+      - "src/**"
       - "tests/**"
       - "examples/**"
       - ".github/**"
@@ -29,7 +29,7 @@ on:
     branches:
       - main
     paths:
-      - "lerobot/**"
+      - "src/**"
       - "tests/**"
       - "examples/**"
       - ".github/**"
@@ -73,7 +73,7 @@ jobs:
 
       - name: Test with pytest
         run: |
-          uv run pytest tests -v --cov=./lerobot --durations=0 \
+          uv run pytest tests -v --cov=./src/lerobot --durations=0 \
             -W ignore::DeprecationWarning:imageio_ffmpeg._utils:7 \
             -W ignore::UserWarning:torch.utils.data.dataloader:558 \
             -W ignore::UserWarning:gymnasium.utils.env_checker:247 \
@@ -105,7 +105,7 @@ jobs:
 
       - name: Test with pytest
         run: |
-          uv run pytest tests -v --cov=./lerobot --durations=0 \
+          uv run pytest tests -v --cov=./src/lerobot --durations=0 \
             -W ignore::DeprecationWarning:imageio_ffmpeg._utils:7 \
             -W ignore::UserWarning:torch.utils.data.dataloader:558 \
             -W ignore::UserWarning:gymnasium.utils.env_checker:247 \
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 23a180046..a354e1346 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -67,7 +67,7 @@ post it.
 
 ## Adding new policies, datasets or environments
 
-Look at our implementations for [datasets](./lerobot/common/datasets/), [policies](./lerobot/common/policies/),
+Look at our implementations for [datasets](./src/lerobot/datasets/), [policies](./src/lerobot/policies/),
 environments ([aloha](https://github.com/huggingface/gym-aloha),
 [xarm](https://github.com/huggingface/gym-xarm),
 [pusht](https://github.com/huggingface/gym-pusht))
diff --git a/MANIFEST.in b/MANIFEST.in
new file mode 100644
index 000000000..c1fb2ea75
--- /dev/null
+++ b/MANIFEST.in
@@ -0,0 +1,2 @@
+include src/lerobot/templates/lerobot_modelcard_template.md
+include src/lerobot/datasets/card_template.md
diff --git a/Makefile b/Makefile
index c82483cc3..ca1495fac 100644
--- a/Makefile
+++ b/Makefile
@@ -40,14 +40,17 @@ test-end-to-end:
 	${MAKE} DEVICE=$(DEVICE) test-diffusion-ete-eval
 	${MAKE} DEVICE=$(DEVICE) test-tdmpc-ete-train
 	${MAKE} DEVICE=$(DEVICE) test-tdmpc-ete-eval
+	${MAKE} DEVICE=$(DEVICE) test-smolvla-ete-train
+	${MAKE} DEVICE=$(DEVICE) test-smolvla-ete-eval
 
 test-act-ete-train:
-	python lerobot/scripts/train.py \
+	python -m lerobot.scripts.train \
 		--policy.type=act \
 		--policy.dim_model=64 \
 		--policy.n_action_steps=20 \
 		--policy.chunk_size=20 \
 		--policy.device=$(DEVICE) \
+		--policy.push_to_hub=false \
 		--env.type=aloha \
 		--env.episode_length=5 \
 		--dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
@@ -65,12 +68,12 @@ test-act-ete-train:
 		--output_dir=tests/outputs/act/
 
 test-act-ete-train-resume:
-	python lerobot/scripts/train.py \
+	python -m lerobot.scripts.train \
 		--config_path=tests/outputs/act/checkpoints/000002/pretrained_model/train_config.json \
 		--resume=true
 
 test-act-ete-eval:
-	python lerobot/scripts/eval.py \
+	python -m lerobot.scripts.eval \
 		--policy.path=tests/outputs/act/checkpoints/000004/pretrained_model \
 		--policy.device=$(DEVICE) \
 		--env.type=aloha \
@@ -79,12 +82,13 @@ test-act-ete-eval:
 		--eval.batch_size=1
 
 test-diffusion-ete-train:
-	python lerobot/scripts/train.py \
+	python -m lerobot.scripts.train \
 		--policy.type=diffusion \
 		--policy.down_dims='[64,128,256]' \
 		--policy.diffusion_step_embed_dim=32 \
 		--policy.num_inference_steps=10 \
 		--policy.device=$(DEVICE) \
+		--policy.push_to_hub=false \
 		--env.type=pusht \
 		--env.episode_length=5 \
 		--dataset.repo_id=lerobot/pusht \
@@ -102,7 +106,7 @@ test-diffusion-ete-train:
 		--output_dir=tests/outputs/diffusion/
 
 test-diffusion-ete-eval:
-	python lerobot/scripts/eval.py \
+	python -m lerobot.scripts.eval \
 		--policy.path=tests/outputs/diffusion/checkpoints/000002/pretrained_model \
 		--policy.device=$(DEVICE) \
 		--env.type=pusht \
@@ -111,9 +115,10 @@ test-diffusion-ete-eval:
 		--eval.batch_size=1
 
 test-tdmpc-ete-train:
-	python lerobot/scripts/train.py \
+	python -m lerobot.scripts.train \
 		--policy.type=tdmpc \
 		--policy.device=$(DEVICE) \
+		--policy.push_to_hub=false \
 		--env.type=xarm \
 		--env.task=XarmLift-v0 \
 		--env.episode_length=5 \
@@ -132,7 +137,7 @@ test-tdmpc-ete-train:
 		--output_dir=tests/outputs/tdmpc/
 
 test-tdmpc-ete-eval:
-	python lerobot/scripts/eval.py \
+	python -m lerobot.scripts.eval \
 		--policy.path=tests/outputs/tdmpc/checkpoints/000002/pretrained_model \
 		--policy.device=$(DEVICE) \
 		--env.type=xarm \
@@ -140,3 +145,36 @@ test-tdmpc-ete-eval:
 		--env.task=XarmLift-v0 \
 		--eval.n_episodes=1 \
 		--eval.batch_size=1
+
+
+test-smolvla-ete-train:
+	python -m lerobot.scripts.train \
+		--policy.type=smolvla \
+		--policy.n_action_steps=20 \
+		--policy.chunk_size=20 \
+		--policy.device=$(DEVICE) \
+		--policy.push_to_hub=false \
+		--env.type=aloha \
+		--env.episode_length=5 \
+		--dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
+		--dataset.image_transforms.enable=true \
+		--dataset.episodes="[0]" \
+		--batch_size=2 \
+		--steps=4 \
+		--eval_freq=2 \
+		--eval.n_episodes=1 \
+		--eval.batch_size=1 \
+		--save_freq=2 \
+		--save_checkpoint=true \
+		--log_freq=1 \
+		--wandb.enable=false \
+		--output_dir=tests/outputs/smolvla/
+
+test-smolvla-ete-eval:
+	python -m lerobot.scripts.eval \
+		--policy.path=tests/outputs/smolvla/checkpoints/000004/pretrained_model \
+		--policy.device=$(DEVICE) \
+		--env.type=aloha \
+		--env.episode_length=5 \
+		--eval.n_episodes=1 \
+		--eval.batch_size=1
diff --git a/README.md b/README.md
index e98f35663..153a3a215 100644
--- a/README.md
+++ b/README.md
@@ -130,7 +130,7 @@ pip install -e .
 ```
 
 > **NOTE:** If you encounter build errors, you may need to install additional dependencies (`cmake`, `build-essential`, and `ffmpeg libs`). On Linux, run:
-`sudo apt-get install cmake build-essential python3-dev pkg-config libavformat-dev libavcodec-dev libavdevice-dev libavutil-dev libswscale-dev libswresample-dev libavfilter-dev pkg-config`. For other systems, see: [Compiling PyAV](https://pyav.org/docs/develop/overview/installation.html#bring-your-own-ffmpeg)
+`sudo apt-get install cmake build-essential python3-dev pkg-config libavformat-dev libavcodec-dev libavdevice-dev libavutil-dev libswscale-dev libswresample-dev libavfilter-dev`. For other systems, see: [Compiling PyAV](https://pyav.org/docs/develop/overview/installation.html#bring-your-own-ffmpeg)
 
 For simulations, 🤗 LeRobot comes with gymnasium environments that can be installed as extras:
 - [aloha](https://github.com/huggingface/gym-aloha)
@@ -149,44 +149,20 @@ wandb login
 
 (note: you will also need to enable WandB in the configuration. See below.)
 
-## Walkthrough
-
-```
-.
-├── examples             # contains demonstration examples, start here to learn about LeRobot
-|   └── advanced         # contains even more examples for those who have mastered the basics
-├── lerobot
-|   ├── configs          # contains config classes with all options that you can override in the command line
-|   ├── common           # contains classes and utilities
-|   |   ├── datasets       # various datasets of human demonstrations: aloha, pusht, xarm
-|   |   ├── envs           # various sim environments: aloha, pusht, xarm
-|   |   ├── policies       # various policies: act, diffusion, tdmpc
-|   |   ├── robot_devices  # various real devices: dynamixel motors, opencv cameras, koch robots
-|   |   └── utils          # various utilities
-|   └── scripts          # contains functions to execute via command line
-|       ├── eval.py                 # load policy and evaluate it on an environment
-|       ├── train.py                # train a policy via imitation learning and/or reinforcement learning
-|       ├── control_robot.py        # teleoperate a real robot, record data, run a policy
-|       ├── push_dataset_to_hub.py  # convert your dataset into LeRobot dataset format and upload it to the Hugging Face hub
-|       └── visualize_dataset.py    # load a dataset and render its demonstrations
-├── outputs               # contains results of scripts execution: logs, videos, model checkpoints
-└── tests                 # contains pytest utilities for continuous integration
-```
-
 ### Visualize datasets
 
 Check out [example 1](./examples/1_load_lerobot_dataset.py) that illustrates how to use our dataset class which automatically downloads data from the Hugging Face hub.
 
 You can also locally visualize episodes from a dataset on the hub by executing our script from the command line:
 ```bash
-python lerobot/scripts/visualize_dataset.py \
+python -m lerobot.scripts.visualize_dataset \
     --repo-id lerobot/pusht \
     --episode-index 0
 ```
 
 or from a dataset in a local folder with the `root` option and the `--local-files-only` (in the following case the dataset will be searched for in `./my_local_data_dir/lerobot/pusht`)
 ```bash
-python lerobot/scripts/visualize_dataset.py \
+python -m lerobot.scripts.visualize_dataset \
     --repo-id lerobot/pusht \
     --root ./my_local_data_dir \
     --local-files-only 1 \
@@ -199,7 +175,7 @@ It will open `rerun.io` and display the camera streams, robot states and actions
 https://github-production-user-asset-6210df.s3.amazonaws.com/4681518/328035972-fd46b787-b532-47e2-bb6f-fd536a55a7ed.mov?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAVCODYLSA53PQK4ZA%2F20240505%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20240505T172924Z&X-Amz-Expires=300&X-Amz-Signature=d680b26c532eeaf80740f08af3320d22ad0b8a4e4da1bcc4f33142c15b509eda&X-Amz-SignedHeaders=host&actor_id=24889239&key_id=0&repo_id=748713144
 
 
-Our script can also visualize datasets stored on a distant server. See `python lerobot/scripts/visualize_dataset.py --help` for more instructions.
+Our script can also visualize datasets stored on a distant server. See `python -m lerobot.scripts.visualize_dataset --help` for more instructions.
 
 ### The `LeRobotDataset` format
 
@@ -252,7 +228,7 @@ Check out [example 2](./examples/2_evaluate_pretrained_policy.py) that illustrat
 
 We also provide a more capable script to parallelize the evaluation over multiple environments during the same rollout. Here is an example with a pretrained model hosted on [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht):
 ```bash
-python lerobot/scripts/eval.py \
+python -m lerobot.scripts.eval \
     --policy.path=lerobot/diffusion_pusht \
     --env.type=pusht \
     --eval.batch_size=10 \
@@ -264,10 +240,10 @@ python lerobot/scripts/eval.py \
 Note: After training your own policy, you can re-evaluate the checkpoints with:
 
 ```bash
-python lerobot/scripts/eval.py --policy.path={OUTPUT_DIR}/checkpoints/last/pretrained_model
+python -m lerobot.scripts.eval --policy.path={OUTPUT_DIR}/checkpoints/last/pretrained_model
 ```
 
-See `python lerobot/scripts/eval.py --help` for more instructions.
+See `python -m lerobot.scripts.eval --help` for more instructions.
 
 ### Train your own policy
 
@@ -279,14 +255,14 @@ A link to the wandb logs for the run will also show up in yellow in your termina
 
 ![](media/wandb.png)
 
-Note: For efficiency, during training every checkpoint is evaluated on a low number of episodes. You may use `--eval.n_episodes=500` to evaluate on more episodes than the default. Or, after training, you may want to re-evaluate your best checkpoints on more episodes or change the evaluation settings. See `python lerobot/scripts/eval.py --help` for more instructions.
+Note: For efficiency, during training every checkpoint is evaluated on a low number of episodes. You may use `--eval.n_episodes=500` to evaluate on more episodes than the default. Or, after training, you may want to re-evaluate your best checkpoints on more episodes or change the evaluation settings. See `python -m lerobot.scripts.eval --help` for more instructions.
 
 #### Reproduce state-of-the-art (SOTA)
 
 We provide some pretrained policies on our [hub page](https://huggingface.co/lerobot) that can achieve state-of-the-art performances.
 You can reproduce their training by loading the config from their run. Simply running:
 ```bash
-python lerobot/scripts/train.py --config_path=lerobot/diffusion_pusht
+python -m lerobot.scripts.train --config_path=lerobot/diffusion_pusht
 ```
 reproduces SOTA results for Diffusion Policy on the PushT task.
 
@@ -312,7 +288,7 @@ python lerobot/scripts/push_dataset_to_hub.py \
 
 See `python lerobot/scripts/push_dataset_to_hub.py --help` for more instructions.
 
-If your dataset format is not supported, implement your own in `lerobot/common/datasets/push_dataset_to_hub/${raw_format}_format.py` by copying examples like [pusht_zarr](https://github.com/huggingface/lerobot/blob/main/lerobot/common/datasets/push_dataset_to_hub/pusht_zarr_format.py), [umi_zarr](https://github.com/huggingface/lerobot/blob/main/lerobot/common/datasets/push_dataset_to_hub/umi_zarr_format.py), [aloha_hdf5](https://github.com/huggingface/lerobot/blob/main/lerobot/common/datasets/push_dataset_to_hub/aloha_hdf5_format.py), or [xarm_pkl](https://github.com/huggingface/lerobot/blob/main/lerobot/common/datasets/push_dataset_to_hub/xarm_pkl_format.py). -->
+If your dataset format is not supported, implement your own in `lerobot/datasets/push_dataset_to_hub/${raw_format}_format.py` by copying examples like [pusht_zarr](https://github.com/huggingface/lerobot/blob/main/lerobot/datasets/push_dataset_to_hub/pusht_zarr_format.py), [umi_zarr](https://github.com/huggingface/lerobot/blob/main/lerobot/datasets/push_dataset_to_hub/umi_zarr_format.py), [aloha_hdf5](https://github.com/huggingface/lerobot/blob/main/lerobot/datasets/push_dataset_to_hub/aloha_hdf5_format.py), or [xarm_pkl](https://github.com/huggingface/lerobot/blob/main/lerobot/datasets/push_dataset_to_hub/xarm_pkl_format.py). -->
 
 
 ### Add a pretrained policy
diff --git a/benchmarks/video/run_video_benchmark.py b/benchmarks/video/run_video_benchmark.py
index 9d587ee9f..bababf636 100644
--- a/benchmarks/video/run_video_benchmark.py
+++ b/benchmarks/video/run_video_benchmark.py
@@ -35,12 +35,12 @@ import torch
 from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity
 from tqdm import tqdm
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.datasets.video_utils import (
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.video_utils import (
     decode_video_frames_torchvision,
     encode_video_frames,
 )
-from lerobot.common.utils.benchmark import TimeBenchmark
+from lerobot.utils.benchmark import TimeBenchmark
 
 BASE_ENCODING = OrderedDict(
     [
diff --git a/docker/lerobot-cpu/Dockerfile b/docker/lerobot-cpu/Dockerfile
index 3bc9cb260..85c31ac1a 100644
--- a/docker/lerobot-cpu/Dockerfile
+++ b/docker/lerobot-cpu/Dockerfile
@@ -22,7 +22,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 COPY . /lerobot
 WORKDIR /lerobot
 RUN /opt/venv/bin/pip install --upgrade --no-cache-dir pip \
-    && /opt/venv/bin/pip install --no-cache-dir ".[test, aloha, xarm, pusht]" \
+    && /opt/venv/bin/pip install --no-cache-dir ".[test, aloha, xarm, pusht, smolvla]" \
         --extra-index-url https://download.pytorch.org/whl/cpu
 
 # Execute in bash shell rather than python
diff --git a/docker/lerobot-gpu/Dockerfile b/docker/lerobot-gpu/Dockerfile
index 642a8ded6..746ea29b7 100644
--- a/docker/lerobot-gpu/Dockerfile
+++ b/docker/lerobot-gpu/Dockerfile
@@ -21,4 +21,4 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 COPY . /lerobot
 WORKDIR /lerobot
 RUN /opt/venv/bin/pip install --upgrade --no-cache-dir pip \
-    && /opt/venv/bin/pip install --no-cache-dir ".[test, aloha, xarm, pusht, dynamixel]"
+    && /opt/venv/bin/pip install --no-cache-dir ".[test, aloha, xarm, pusht, dynamixel, smolvla]"
diff --git a/docs/source/cameras.mdx b/docs/source/cameras.mdx
index d8a49c1ee..313d5a7cd 100644
--- a/docs/source/cameras.mdx
+++ b/docs/source/cameras.mdx
@@ -8,7 +8,7 @@ To instantiate a camera, you need a camera identifier. This identifier might cha
 
 To find the camera indices of the cameras plugged into your system, run the following script:
 ```bash
-python lerobot/find_cameras.py opencv # or realsense for Intel Realsense cameras
+python -m lerobot.find_cameras opencv # or realsense for Intel Realsense cameras
 ```
 
 The output will look something like this if you have two cameras connected:
@@ -44,9 +44,9 @@ Below are two examples, demonstrating how to work with the API.
 <hfoption id="Open CV Camera">
 
 ```python
-from lerobot.common.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.common.cameras.opencv.camera_opencv import OpenCVCamera
-from lerobot.common.cameras.configs import ColorMode, Cv2Rotation
+from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
+from lerobot.cameras.opencv.camera_opencv import OpenCVCamera
+from lerobot.cameras.configs import ColorMode, Cv2Rotation
 
 # Construct an `OpenCVCameraConfig` with your desired FPS, resolution, color mode, and rotation.
 config = OpenCVCameraConfig(
@@ -75,9 +75,9 @@ finally:
 <hfoption id="Intel Realsense Camera">
 
 ```python
-from lerobot.common.cameras.realsense.configuration_realsense import RealSenseCameraConfig
-from lerobot.common.cameras.realsense.camera_realsense import RealSenseCamera
-from lerobot.common.cameras.configs import ColorMode, Cv2Rotation
+from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig
+from lerobot.cameras.realsense.camera_realsense import RealSenseCamera
+from lerobot.cameras.configs import ColorMode, Cv2Rotation
 
 # Create a `RealSenseCameraConfig` specifying your camera’s serial number and enabling depth.
 config = RealSenseCameraConfig(
diff --git a/docs/source/hilserl.mdx b/docs/source/hilserl.mdx
index 149b25c68..b3ab40c89 100644
--- a/docs/source/hilserl.mdx
+++ b/docs/source/hilserl.mdx
@@ -24,6 +24,7 @@ This guide provides step-by-step instructions for training a robot policy using
 - A gamepad (recommended) or keyboard to control the robot
 - A Nvidia GPU
 - A real robot with a follower and leader arm (optional if you use the keyboard or the gamepad)
+- A URDF file for the robot for the kinematics package (check `lerobot/common/model/kinematics.py`)
 
 ## What kind of tasks can I train?
 
@@ -50,12 +51,12 @@ pip install -e ".[hilserl]"
 
 ### Understanding Configuration
 
-The training process begins with proper configuration for the HILSerl environment. The configuration class of interest is `HILSerlRobotEnvConfig` in `lerobot/common/envs/configs.py`. Which is defined as:
+The training process begins with proper configuration for the HILSerl environment. The configuration class of interest is `HILSerlRobotEnvConfig` in `lerobot/envs/configs.py`. Which is defined as:
 
 ```python
 class HILSerlRobotEnvConfig(EnvConfig):
-    robot: RobotConfig | None = None    # Main robot agent (defined in `lerobot/common/robots`)
-    teleop: TeleoperatorConfig | None = None    # Teleoperator agent, e.g., gamepad or leader arm, (defined in `lerobot/common/teleoperators`)
+    robot: RobotConfig | None = None    # Main robot agent (defined in `lerobot/robots`)
+    teleop: TeleoperatorConfig | None = None    # Teleoperator agent, e.g., gamepad or leader arm, (defined in `lerobot/teleoperators`)
     wrapper: EnvTransformConfig | None = None    # Environment wrapper settings; check `lerobot/scripts/server/gym_manipulator.py`
     fps: int = 10    # Control frequency
     name: str = "real_robot"    # Environment name
@@ -172,7 +173,7 @@ class SO100FollowerEndEffectorConfig(SO100FollowerConfig):
     )
 ```
 
-The `Teleoperator` defines the teleoperation device. You can check the list of available teleoperators in `lerobot/common/teleoperators`.
+The `Teleoperator` defines the teleoperation device. You can check the list of available teleoperators in `lerobot/teleoperators`.
 
 **Setting up the Gamepad**
 
@@ -226,7 +227,7 @@ During the online training, press `space` to take over the policy and `space` ag
 Start the recording process, an example of the config file can be found [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/env_config_so100.json):
 
 ```bash
-python lerobot/scripts/rl/gym_manipulator.py --config_path lerobot/configs/env_config_so100.json
+python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config_so100.json
 ```
 
 During recording:
@@ -256,7 +257,7 @@ Note: If you already know the crop parameters, you can skip this step and just s
 Use the `crop_dataset_roi.py` script to interactively select regions of interest in your camera images:
 
 ```bash
-python lerobot/scripts/rl/crop_dataset_roi.py --repo-id username/pick_lift_cube
+python -m lerobot.scripts.rl.crop_dataset_roi --repo-id username/pick_lift_cube
 ```
 
 1. For each camera view, the script will display the first frame
@@ -313,7 +314,7 @@ Before training, you need to collect a dataset with labeled examples. The `recor
 To collect a dataset, you need to modify some parameters in the environment configuration based on HILSerlRobotEnvConfig.
 
 ```bash
-python lerobot/scripts/rl/gym_manipulator.py --config_path lerobot/configs/reward_classifier_train_config.json
+python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/reward_classifier_train_config.json
 ```
 
 **Key Parameters for Data Collection**
@@ -387,7 +388,7 @@ Example configuration for training the [reward classifier](https://huggingface.c
 To train the classifier, use the `train.py` script with your configuration:
 
 ```bash
-python lerobot/scripts/train.py --config_path path/to/reward_classifier_train_config.json
+python -m lerobot.scripts.train --config_path path/to/reward_classifier_train_config.json
 ```
 
 **Deploying and Testing the Model**
@@ -410,7 +411,7 @@ or set the argument in the json config file.
 
 Run `gym_manipulator.py` to test the model.
 ```bash
-python lerobot/scripts/rl/gym_manipulator.py --config_path path/to/env_config.json
+python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config.json
 ```
 
 The reward classifier will automatically provide rewards based on the visual input from the robot's cameras.
@@ -422,17 +423,17 @@ The reward classifier will automatically provide rewards based on the visual inp
 
 2. **Collect a dataset**:
    ```bash
-   python lerobot/scripts/rl/gym_manipulator.py --config_path lerobot/configs/env_config.json
+   python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
    ```
 
 3. **Train the classifier**:
    ```bash
-   python lerobot/scripts/train.py --config_path lerobot/configs/reward_classifier_train_config.json
+   python -m lerobot.scripts.train --config_path src/lerobot/configs/reward_classifier_train_config.json
    ```
 
 4. **Test the classifier**:
    ```bash
-   python lerobot/scripts/rl/gym_manipulator.py --config_path lerobot/configs/env_config.json
+   python -m lerobot.scripts.rl.gym_manipulator --config_path src/lerobot/configs/env_config.json
    ```
 
 ### Training with Actor-Learner
@@ -446,7 +447,7 @@ Create a training configuration file (example available [here](https://huggingfa
 1. Configure the policy settings (`type="sac"`, `device`, etc.)
 2. Set `dataset` to your cropped dataset
 3. Configure environment settings with crop parameters
-4. Check the other parameters related to SAC in [configuration_sac.py](https://github.com/huggingface/lerobot/blob/19bb621a7d0a31c20cd3cc08b1dbab68d3031454/lerobot/common/policies/sac/configuration_sac.py#L79).
+4. Check the other parameters related to SAC in [configuration_sac.py](https://github.com/huggingface/lerobot/blob/19bb621a7d0a31c20cd3cc08b1dbab68d3031454/lerobot/policies/sac/configuration_sac.py#L79).
 5. Verify that the `policy` config is correct with the right `input_features` and `output_features` for your task.
 
 **Starting the Learner**
@@ -454,7 +455,7 @@ Create a training configuration file (example available [here](https://huggingfa
 First, start the learner server process:
 
 ```bash
-python lerobot/scripts/rl/learner.py --config_path lerobot/configs/train_config_hilserl_so100.json
+python -m lerobot.scripts.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
 ```
 
 The learner:
@@ -468,7 +469,7 @@ The learner:
 In a separate terminal, start the actor process with the same configuration:
 
 ```bash
-python lerobot/scripts/rl/actor.py --config_path lerobot/configs/train_config_hilserl_so100.json
+python -m lerobot.scripts.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
 ```
 
 The actor:
diff --git a/docs/source/hilserl_sim.mdx b/docs/source/hilserl_sim.mdx
index 3239ba91a..ad7a9584a 100644
--- a/docs/source/hilserl_sim.mdx
+++ b/docs/source/hilserl_sim.mdx
@@ -77,7 +77,7 @@ Important parameters:
 To run the environment, set mode to null:
 
 ```python
-python lerobot/scripts/rl/gym_manipulator.py --config_path path/to/gym_hil_env.json
+python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/gym_hil_env.json
 ```
 
 ### Recording a Dataset
@@ -85,7 +85,7 @@ python lerobot/scripts/rl/gym_manipulator.py --config_path path/to/gym_hil_env.j
 To collect a dataset, set the mode to `record` whilst defining the repo_id and number of episodes to record:
 
 ```python
-python lerobot/scripts/rl/gym_manipulator.py --config_path path/to/gym_hil_env.json
+python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/gym_hil_env.json
 ```
 
 ### Training a Policy
@@ -93,13 +93,13 @@ python lerobot/scripts/rl/gym_manipulator.py --config_path path/to/gym_hil_env.j
 To train a policy, checkout the configuration example available [here](https://huggingface.co/datasets/aractingi/lerobot-example-config-files/blob/main/train_gym_hil_env.json) and run the actor and learner servers:
 
 ```python
-python lerobot/scripts/rl/actor.py --config_path path/to/train_gym_hil_env.json
+python -m lerobot.scripts.rl.actor --config_path path/to/train_gym_hil_env.json
 ```
 
 In a different terminal, run the learner server:
 
 ```python
-python lerobot/scripts/rl/learner.py --config_path path/to/train_gym_hil_env.json
+python -m lerobot.scripts.rl.learner --config_path path/to/train_gym_hil_env.json
 ```
 
 The simulation environment provides a safe and repeatable way to develop and test your Human-In-the-Loop reinforcement learning components before deploying to real robots.
diff --git a/docs/source/il_robots.mdx b/docs/source/il_robots.mdx
index f3b4b1a25..cfa0a2809 100644
--- a/docs/source/il_robots.mdx
+++ b/docs/source/il_robots.mdx
@@ -52,8 +52,8 @@ python -m lerobot.teleoperate \
 </hfoption>
 <hfoption id="API example">
 ```python
-from lerobot.common.teleoperators.so101_leader import SO101LeaderConfig, SO101Leader
-from lerobot.common.robots.so101_follower import SO101FollowerConfig, SO101Follower
+from lerobot.teleoperators.so101_leader import SO101LeaderConfig, SO101Leader
+from lerobot.robots.so101_follower import SO101FollowerConfig, SO101Follower
 
 robot_config = SO101FollowerConfig(
     port="/dev/tty.usbmodem58760431541",
@@ -105,9 +105,9 @@ python -m lerobot.teleoperate \
 </hfoption>
 <hfoption id="API example">
 ```python
-from lerobot.common.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.common.teleoperators.koch_leader import KochLeaderConfig, KochLeader
-from lerobot.common.robots.koch_follower import KochFollowerConfig, KochFollower
+from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
+from lerobot.teleoperators.koch_leader import KochLeaderConfig, KochLeader
+from lerobot.robots.koch_follower import KochFollowerConfig, KochFollower
 
 camera_config = {
     "front": OpenCVCameraConfig(index_or_path=0, width=1920, height=1080, fps=30)
@@ -154,7 +154,10 @@ HF_USER=$(huggingface-cli whoami | head -n 1)
 echo $HF_USER
 ```
 
-Now you can record a dataset. To record 2 episodes and upload your dataset to the hub, execute this command tailored to the SO101.
+Now you can record a dataset. To record 5 episodes and upload your dataset to the hub, adapt the code below for your robot and execute the command or API example.
+
+<hfoptions id="record">
+<hfoption id="Command">
 ```bash
 python -m lerobot.record \
     --robot.type=so101_follower \
@@ -166,9 +169,109 @@ python -m lerobot.record \
     --teleop.id=my_awesome_leader_arm \
     --display_data=true \
     --dataset.repo_id=${HF_USER}/record-test \
-    --dataset.num_episodes=2 \
+    --dataset.num_episodes=5 \
     --dataset.single_task="Grab the black cube"
 ```
+</hfoption>
+<hfoption id="API example">
+```python
+from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import hw_to_dataset_features
+from lerobot.robots.so100_follower import SO100Follower, SO100FollowerConfig
+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.record import record_loop
+
+NUM_EPISODES = 5
+FPS = 30
+EPISODE_TIME_SEC = 60
+RESET_TIME_SEC = 10
+TASK_DESCRIPTION = "My task description"
+
+# Create the robot and teleoperator configurations
+camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
+robot_config = SO100FollowerConfig(
+    port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", cameras=camera_config
+)
+teleop_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm")
+
+# Initialize the robot and teleoperator
+robot = SO100Follower(robot_config)
+teleop = SO100Leader(teleop_config)
+
+# Configure the dataset features
+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
+dataset = LeRobotDataset.create(
+    repo_id="<hf_username>/<dataset_repo_id>",
+    fps=FPS,
+    features=dataset_features,
+    robot_type=robot.name,
+    use_videos=True,
+    image_writer_threads=4,
+)
+
+# Initialize the keyboard listener and rerun visualization
+_, events = init_keyboard_listener()
+_init_rerun(session_name="recording")
+
+# Connect the robot and teleoperator
+robot.connect()
+teleop.connect()
+
+episode_idx = 0
+while episode_idx < NUM_EPISODES and not events["stop_recording"]:
+    log_say(f"Recording episode {episode_idx + 1} of {NUM_EPISODES}")
+
+    record_loop(
+        robot=robot,
+        events=events,
+        fps=FPS,
+        teleop=teleop,
+        dataset=dataset,
+        control_time_s=EPISODE_TIME_SEC,
+        single_task=TASK_DESCRIPTION,
+        display_data=True,
+    )
+
+    # Reset the environment if not stopping or re-recording
+    if not events["stop_recording"] and (episode_idx < NUM_EPISODES - 1 or events["rerecord_episode"]):
+        log_say("Reset the environment")
+        record_loop(
+            robot=robot,
+            events=events,
+            fps=FPS,
+            teleop=teleop,
+            control_time_s=RESET_TIME_SEC,
+            single_task=TASK_DESCRIPTION,
+            display_data=True,
+        )
+
+    if events["rerecord_episode"]:
+        log_say("Re-recording episode")
+        events["rerecord_episode"] = False
+        events["exit_early"] = False
+        dataset.clear_episode_buffer()
+        continue
+
+    dataset.save_episode()
+    episode_idx += 1
+
+# Clean up
+log_say("Stop recording")
+robot.disconnect()
+teleop.disconnect()
+dataset.push_to_hub()
+```
+</hfoption>
+</hfoptions>
 
 #### Dataset upload
 Locally, your dataset is stored in this folder: `~/.cache/huggingface/lerobot/{repo-id}`. At the end of data recording, your dataset will be uploaded on your Hugging Face page (e.g. https://huggingface.co/datasets/cadene/so101_test) that you can obtain by running:
@@ -233,7 +336,10 @@ echo ${HF_USER}/so101_test
 
 A useful feature is the `replay` function, which allows you to replay any episode that you've recorded or episodes from any dataset out there. This function helps you test the repeatability of your robot's actions and assess transferability across robots of the same model.
 
-You can replay the first episode on your robot with:
+You can replay the first episode on your robot with either the command below or with the API example:
+
+<hfoptions id="replay">
+<hfoption id="Command">
 ```bash
 python -m lerobot.replay \
     --robot.type=so101_follower \
@@ -242,25 +348,62 @@ python -m lerobot.replay \
     --dataset.repo_id=${HF_USER}/record-test \
     --dataset.episode=0 # choose the episode you want to replay
 ```
+</hfoption>
+<hfoption id="API example">
+```python
+import time
+
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
+from lerobot.robots.so100_follower.so100_follower import SO100Follower
+from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.utils import log_say
+
+episode_idx = 0
+
+robot_config = SO100FollowerConfig(port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm")
+
+robot = SO100Follower(robot_config)
+robot.connect()
+
+dataset = LeRobotDataset("<hf_username>/<dataset_repo_id>", episodes=[episode_idx])
+actions = dataset.hf_dataset.select_columns("action")
+
+log_say(f"Replaying episode {episode_idx}")
+for idx in range(dataset.num_frames):
+    t0 = time.perf_counter()
+
+    action = {
+        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
+    }
+    robot.send_action(action)
+
+    busy_wait(1.0 / dataset.fps - (time.perf_counter() - t0))
+
+robot.disconnect()
+```
+</hfoption>
+</hfoptions>
 
 Your robot should replicate movements similar to those you recorded. For example, check out [this video](https://x.com/RemiCadene/status/1793654950905680090) where we use `replay` on a Aloha robot from [Trossen Robotics](https://www.trossenrobotics.com).
 
 ## Train a policy
 
-To train a policy to control your robot, use the [`python lerobot/scripts/train.py`](../lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
+To train a policy to control your robot, use the [`python -m lerobot.scripts.train`](../src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
   --dataset.repo_id=${HF_USER}/so101_test \
   --policy.type=act \
   --output_dir=outputs/train/act_so101_test \
   --job_name=act_so101_test \
   --policy.device=cuda \
-  --wandb.enable=true
+  --wandb.enable=true \
+  --policy.repo_id=${HF_USER}/my_policy
 ```
 
 Let's explain the command:
 1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/so101_test`.
-2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../lerobot/common/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor states, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
+2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../src/lerobot/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor states, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
 4. We provided `policy.device=cuda` since we are training on a Nvidia GPU, but you could use `policy.device=mps` to train on Apple silicon.
 5. We provided `wandb.enable=true` to use [Weights and Biases](https://docs.wandb.ai/quickstart) for visualizing training plots. This is optional but if you use it, make sure you are logged in by running `wandb login`.
 
@@ -268,11 +411,15 @@ Training should take several hours. You will find checkpoints in `outputs/train/
 
 To resume training from a checkpoint, below is an example command to resume from `last` checkpoint of the `act_so101_test` policy:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
   --config_path=outputs/train/act_so101_test/checkpoints/last/pretrained_model/train_config.json \
   --resume=true
 ```
 
+If you do not want to push your model to the hub after training use `--policy.push_to_hub=false`.
+
+Additionally you can provide extra `tags` or specify a `license` for your model or make the model repo `private` by adding this: `--policy.private=true --policy.tags=\[ppo,rl\] --policy.license=mit`
+
 #### Train using Collab
 If your local computer doesn't have a powerful GPU you could utilize Google Collab to train your model by following the [ACT training notebook](./notebooks#training-act).
 
@@ -291,9 +438,12 @@ huggingface-cli upload ${HF_USER}/act_so101_test${CKPT} \
   outputs/train/act_so101_test/checkpoints/${CKPT}/pretrained_model
 ```
 
-## Evaluate your policy
+## Run inference and evaluate your policy
 
-You can use the `record` script from [`lerobot/record.py`](https://github.com/huggingface/lerobot/blob/main/lerobot/record.py) but with a policy checkpoint as input. For instance, run this command to record 10 evaluation episodes:
+You can use the `record` script from [`lerobot/record.py`](https://github.com/huggingface/lerobot/blob/main/lerobot/record.py) with a policy checkpoint as input, to run inference and evaluate your policy. For instance, run this command or API example to run inference and record 10 evaluation episodes:
+
+<hfoptions id="eval">
+<hfoption id="Command">
 ```bash
 python -m lerobot.record  \
   --robot.type=so100_follower \
@@ -309,6 +459,82 @@ python -m lerobot.record  \
   # --teleop.id=my_awesome_leader_arm \
   --policy.path=${HF_USER}/my_policy
 ```
+</hfoption>
+<hfoption id="API example">
+```python
+from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import hw_to_dataset_features
+from lerobot.policies.act.modeling_act import ACTPolicy
+from lerobot.robots.so100_follower.config_so100_follower import SO100FollowerConfig
+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.record import record_loop
+
+NUM_EPISODES = 5
+FPS = 30
+EPISODE_TIME_SEC = 60
+TASK_DESCRIPTION = "My task description"
+
+# Create the robot configuration
+camera_config = {"front": OpenCVCameraConfig(index_or_path=0, width=640, height=480, fps=FPS)}
+robot_config = SO100FollowerConfig(
+    port="/dev/tty.usbmodem58760434471", id="my_awesome_follower_arm", cameras=camera_config
+)
+
+# Initialize the robot
+robot = SO100Follower(robot_config)
+
+# Initialize the policy
+policy = ACTPolicy.from_pretrained("<hf_username>/<my_policy_repo_id>")
+
+# Configure the dataset features
+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
+dataset = LeRobotDataset.create(
+    repo_id="<hf_username>/eval_<dataset_repo_id>",
+    fps=FPS,
+    features=dataset_features,
+    robot_type=robot.name,
+    use_videos=True,
+    image_writer_threads=4,
+)
+
+# Initialize the keyboard listener and rerun visualization
+_, events = init_keyboard_listener()
+_init_rerun(session_name="recording")
+
+# Connect the robot
+robot.connect()
+
+for episode_idx in range(NUM_EPISODES):
+    log_say(f"Running inference, recording eval episode {episode_idx + 1} of {NUM_EPISODES}")
+
+    # Run the policy inference loop
+    record_loop(
+        robot=robot,
+        events=events,
+        fps=FPS,
+        policy=policy,
+        dataset=dataset,
+        control_time_s=EPISODE_TIME_SEC,
+        single_task=TASK_DESCRIPTION,
+        display_data=True,
+    )
+
+    dataset.save_episode()
+
+# Clean up
+robot.disconnect()
+dataset.push_to_hub()
+```
+</hfoption>
+</hfoptions>
 
 As you can see, it's almost the same command as previously used to record your training dataset. Two things changed:
 1. There is an additional `--control.policy.path` argument which indicates the path to your policy checkpoint with  (e.g. `outputs/train/eval_act_so101_test/checkpoints/last/pretrained_model`). You can also use the model repository if you uploaded a model checkpoint to the hub (e.g. `${HF_USER}/act_so101_test`).
diff --git a/docs/source/il_sim.mdx b/docs/source/il_sim.mdx
index 625b2fc00..048d3147e 100644
--- a/docs/source/il_sim.mdx
+++ b/docs/source/il_sim.mdx
@@ -35,14 +35,14 @@ Then we can run this command to start:
 <hfoption id="Linux">
 
 ```bash
-python lerobot/scripts/rl/gym_manipulator.py --config_path path/to/env_config_gym_hil_il.json
+python -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
 ```
 
 </hfoption>
 <hfoption id="MacOS">
 
 ```bash
-mjpython lerobot/scripts/rl/gym_manipulator.py --config_path path/to/env_config_gym_hil_il.json
+mjpython -m lerobot.scripts.rl.gym_manipulator --config_path path/to/env_config_gym_hil_il.json
 ```
 
 </hfoption>
@@ -81,9 +81,9 @@ If you uploaded your dataset to the hub you can [visualize your dataset online](
 
 ## Train a policy
 
-To train a policy to control your robot, use the [`python lerobot/scripts/train.py`](../lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
+To train a policy to control your robot, use the [`python -m lerobot.scripts.train`](../src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
   --dataset.repo_id=${HF_USER}/il_gym \
   --policy.type=act \
   --output_dir=outputs/train/il_sim_test \
@@ -94,7 +94,7 @@ python lerobot/scripts/train.py \
 
 Let's explain the command:
 1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/il_gym`.
-2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../lerobot/common/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor states, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
+2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../src/lerobot/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor states, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
 4. We provided `policy.device=cuda` since we are training on a Nvidia GPU, but you could use `policy.device=mps` to train on Apple silicon.
 5. We provided `wandb.enable=true` to use [Weights and Biases](https://docs.wandb.ai/quickstart) for visualizing training plots. This is optional but if you use it, make sure you are logged in by running `wandb login`.
 
@@ -130,14 +130,14 @@ Then you can run this command to visualize your trained policy
 <hfoption id="Linux">
 
 ```bash
-python lerobot/scripts/rl/eval_policy.py --config_path=path/to/eval_config_gym_hil.json
+python -m lerobot.scripts.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
 ```
 
 </hfoption>
 <hfoption id="MacOS">
 
 ```bash
-mjpython lerobot/scripts/rl/eval_policy.py --config_path=path/to/eval_config_gym_hil.json
+mjpython -m lerobot.scripts.rl.eval_policy --config_path=path/to/eval_config_gym_hil.json
 ```
 
 </hfoption>
diff --git a/docs/source/integrate_hardware.mdx b/docs/source/integrate_hardware.mdx
index f7de1cece..18d73d3cd 100644
--- a/docs/source/integrate_hardware.mdx
+++ b/docs/source/integrate_hardware.mdx
@@ -2,7 +2,7 @@
 
 This tutorial will explain how to integrate your own robot design into the LeRobot ecosystem and have it access all of our tools (data collection, control pipelines, policy training and inference).
 
-To that end, we provide the [`Robot`](https://github.com/huggingface/lerobot/blob/main/lerobot/common/robots/robot.py) base class in the LeRobot which specifies a standard interface for physical robot integration. Let's see how to implement it.
+To that end, we provide the [`Robot`](https://github.com/huggingface/lerobot/blob/main/lerobot/robots/robot.py) base class in the LeRobot which specifies a standard interface for physical robot integration. Let's see how to implement it.
 
 ## Prerequisites
 
@@ -14,11 +14,11 @@ To that end, we provide the [`Robot`](https://github.com/huggingface/lerobot/blo
 
 If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus interfaces:
 
-- [`FeetechMotorsBus`](https://github.com/huggingface/lerobot/blob/main/lerobot/common/motors/feetech/feetech.py) – for controlling Feetech servos
-- [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/lerobot/common/motors/dynamixel/dynamixel.py) – for controlling Dynamixel servos
+- [`FeetechMotorsBus`](https://github.com/huggingface/lerobot/blob/main/lerobot/motors/feetech/feetech.py) – for controlling Feetech servos
+- [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/lerobot/motors/dynamixel/dynamixel.py) – for controlling Dynamixel servos
 
-Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/lerobot/common/motors/motors_bus.py) abstract class to learn about its API.
-For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/lerobot/common/robots/so101_follower/so101_follower.py)
+Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/lerobot/motors/motors_bus.py) abstract class to learn about its API.
+For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/lerobot/robots/so101_follower/so101_follower.py)
 
 Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):
 - Find an existing SDK in Python (or use bindings to C/C++)
@@ -32,7 +32,7 @@ For Feetech and Dynamixel, we currently support these servos:
         - SCS series (protocol 1): `scs0009`
     - Dynamixel (protocol 2.0 only): `xl330-m077`, `xl330-m288`, `xl430-w250`, `xm430-w350`, `xm540-w270`, `xc430-w150`
 
-If you are using Feetech or Dynamixel servos that are not in this list, you can add those in the [Feetech table](https://github.com/huggingface/lerobot/blob/main/lerobot/common/motors/feetech/tables.py) or [Dynamixel table](https://github.com/huggingface/lerobot/blob/main/lerobot/common/motors/dynamixel/tables.py). Depending on the model, this will require you to add model-specific information. In most cases though, there shouldn't be a lot of additions to do.
+If you are using Feetech or Dynamixel servos that are not in this list, you can add those in the [Feetech table](https://github.com/huggingface/lerobot/blob/main/lerobot/motors/feetech/tables.py) or [Dynamixel table](https://github.com/huggingface/lerobot/blob/main/lerobot/motors/dynamixel/tables.py). Depending on the model, this will require you to add model-specific information. In most cases though, there shouldn't be a lot of additions to do.
 
 In the next sections, we'll use a `FeetechMotorsBus` as the motors interface for the examples. Replace it and adapt to your motors if necessary.
 
@@ -44,9 +44,9 @@ Here, we'll add the port name and one camera by default for our robot:
 ```python
 from dataclasses import dataclass, field
 
-from lerobot.common.cameras import CameraConfig
-from lerobot.common.cameras.opencv import OpenCVCameraConfig
-from lerobot.common.robots import RobotConfig
+from lerobot.cameras import CameraConfig
+from lerobot.cameras.opencv import OpenCVCameraConfig
+from lerobot.robots import RobotConfig
 
 
 @RobotConfig.register_subclass("my_cool_robot")
@@ -72,10 +72,10 @@ Next, we'll create our actual robot class which inherits from `Robot`. This abst
 Here we'll create a simple 5-DoF robot with one camera. It could be a simple arm but notice that the `Robot` abstract class does not assume anything on your robot's form factor. You can let you imagination run wild when designing new robots!
 
 ```python
-from lerobot.common.cameras import make_cameras_from_configs
-from lerobot.common.motors import Motor, MotorNormMode
-from lerobot.common.motors.feetech import FeetechMotorsBus
-from lerobot.common.robots import Robot
+from lerobot.cameras import make_cameras_from_configs
+from lerobot.motors import Motor, MotorNormMode
+from lerobot.motors.feetech import FeetechMotorsBus
+from lerobot.robots import Robot
 
 class MyCoolRobot(Robot):
     config_class = MyCoolRobotConfig
@@ -303,7 +303,7 @@ def send_action(self, action: dict[str, Any]) -> dict[str, Any]:
 
 ## Adding a Teleoperator
 
-For implementing teleoperation devices, we also provide a [`Teleoperator`](https://github.com/huggingface/lerobot/blob/main/lerobot/common/teleoperators/teleoperator.py) base class. This class is very similar to the `Robot` base class and also doesn't assume anything on form factor.
+For implementing teleoperation devices, we also provide a [`Teleoperator`](https://github.com/huggingface/lerobot/blob/main/lerobot/teleoperators/teleoperator.py) base class. This class is very similar to the `Robot` base class and also doesn't assume anything on form factor.
 
 The main differences are in the I/O functions: a teleoperator allows you to produce action via `get_action` and can receive feedback actions via `send_feedback`. Feedback could be anything controllable on the teleoperation device that could help the person controlling it understand the consequences of the actions sent. Think motion/force feedback on a leader arm, vibrations on a gamepad controller for example. To implement a teleoperator, you can follow this same tutorial and adapt it for these two methods.
 
diff --git a/docs/source/koch.mdx b/docs/source/koch.mdx
index b2399ae62..5383518b3 120000
--- a/docs/source/koch.mdx
+++ b/docs/source/koch.mdx
@@ -1 +1 @@
-../../lerobot/common/robots/koch_follower/koch.mdx
\ No newline at end of file
+../../src/lerobot/robots/koch_follower/koch.mdx
\ No newline at end of file
diff --git a/docs/source/lekiwi.mdx b/docs/source/lekiwi.mdx
index e2b4ff552..afc43077e 120000
--- a/docs/source/lekiwi.mdx
+++ b/docs/source/lekiwi.mdx
@@ -1 +1 @@
-../../lerobot/common/robots/lekiwi/lekiwi.mdx
\ No newline at end of file
+../../src/lerobot/robots/lekiwi/lekiwi.mdx
\ No newline at end of file
diff --git a/docs/source/smolvla.mdx b/docs/source/smolvla.mdx
index 1d6596f65..17a2bdf18 100644
--- a/docs/source/smolvla.mdx
+++ b/docs/source/smolvla.mdx
@@ -44,7 +44,7 @@ If you don't have a gpu device, you can train using our notebook on [![Google Co
 Pass your dataset to the training script using `--dataset.repo_id`. If you want to test your installation, run the following command where we use one of the datasets we collected for the [SmolVLA Paper](https://huggingface.co/papers/2506.01844).
 
 ```bash
-cd lerobot && python lerobot/scripts/train.py \
+cd lerobot && python -m lerobot.scripts.train \
   --policy.path=lerobot/smolvla_base \
   --dataset.repo_id=${HF_USER}/mydataset \
   --batch_size=64 \
@@ -62,7 +62,7 @@ You can start with a small batch size and increase it incrementally, if the GPU
 Fine-tuning is an art. For a complete overview of the options for finetuning, run
 
 ```bash
-python lerobot/scripts/train.py --help
+python -m lerobot.scripts.train --help
 ```
 
 <p align="center">
diff --git a/docs/source/so100.mdx b/docs/source/so100.mdx
index 65849e950..0a71dc307 120000
--- a/docs/source/so100.mdx
+++ b/docs/source/so100.mdx
@@ -1 +1 @@
-../../lerobot/common/robots/so100_follower/so100.mdx
\ No newline at end of file
+../../src/lerobot/robots/so100_follower/so100.mdx
\ No newline at end of file
diff --git a/docs/source/so101.mdx b/docs/source/so101.mdx
index dc4720c28..ab6d0ac61 120000
--- a/docs/source/so101.mdx
+++ b/docs/source/so101.mdx
@@ -1 +1 @@
-../../lerobot/common/robots/so101_follower/so101.mdx
\ No newline at end of file
+../../src/lerobot/robots/so101_follower/so101.mdx
\ No newline at end of file
diff --git a/examples/1_load_lerobot_dataset.py b/examples/1_load_lerobot_dataset.py
index 07db38a15..3d357dd19 100644
--- a/examples/1_load_lerobot_dataset.py
+++ b/examples/1_load_lerobot_dataset.py
@@ -32,7 +32,7 @@ import torch
 from huggingface_hub import HfApi
 
 import lerobot
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
 
 # We ported a number of existing datasets ourselves, use this to see the list:
 print("List of available datasets:")
diff --git a/examples/2_evaluate_pretrained_policy.py b/examples/2_evaluate_pretrained_policy.py
index 4e6154c2e..c0c7845e8 100644
--- a/examples/2_evaluate_pretrained_policy.py
+++ b/examples/2_evaluate_pretrained_policy.py
@@ -30,7 +30,7 @@ import imageio
 import numpy
 import torch
 
-from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy
+from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
 
 # Create a directory to store the video of the evaluation
 output_directory = Path("outputs/eval/example_pusht_diffusion")
diff --git a/examples/3_train_policy.py b/examples/3_train_policy.py
index f9c251a02..f2de79db8 100644
--- a/examples/3_train_policy.py
+++ b/examples/3_train_policy.py
@@ -22,11 +22,11 @@ from pathlib import Path
 
 import torch
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
-from lerobot.common.datasets.utils import dataset_to_policy_features
-from lerobot.common.policies.diffusion.configuration_diffusion import DiffusionConfig
-from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy
 from lerobot.configs.types import FeatureType
+from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.datasets.utils import dataset_to_policy_features
+from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
+from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
 
 
 def main():
diff --git a/examples/4_train_policy_with_script.md b/examples/4_train_policy_with_script.md
index cb4cc6268..f17411b75 100644
--- a/examples/4_train_policy_with_script.md
+++ b/examples/4_train_policy_with_script.md
@@ -4,7 +4,7 @@ This tutorial will explain the training script, how to use it, and particularly
 
 ## The training script
 
-LeRobot offers a training script at [`lerobot/scripts/train.py`](../lerobot/scripts/train.py). At a high level it does the following:
+LeRobot offers a training script at [`lerobot/scripts/train.py`](../src/lerobot/scripts/train.py). At a high level it does the following:
 
 - Initialize/load a configuration for the following steps using.
 - Instantiates a dataset.
@@ -21,7 +21,7 @@ In the training script, the main function `train` expects a `TrainPipelineConfig
 def train(cfg: TrainPipelineConfig):
 ```
 
-You can inspect the `TrainPipelineConfig` defined in [`lerobot/configs/train.py`](../lerobot/configs/train.py) (which is heavily commented and meant to be a reference to understand any option)
+You can inspect the `TrainPipelineConfig` defined in [`lerobot/configs/train.py`](../src/lerobot/configs/train.py) (which is heavily commented and meant to be a reference to understand any option)
 
 When running the script, inputs for the command line are parsed thanks to the `@parser.wrap()` decorator and an instance of this class is automatically generated. Under the hood, this is done with [Draccus](https://github.com/dlwh/draccus) which is a tool dedicated to this purpose. If you're familiar with Hydra, Draccus can similarly load configurations from config files (.json, .yaml) and also override their values through command line inputs. Unlike Hydra, these configurations are pre-defined in the code through dataclasses rather than being defined entirely in config files. This allows for more rigorous serialization/deserialization, typing, and to manipulate configuration as objects directly in the code and not as dictionaries or namespaces (which enables nice features in an IDE such as autocomplete, jump-to-def, etc.)
 
@@ -50,9 +50,9 @@ By default, every field takes its default value specified in the dataclass. If a
 
 ## Specifying values from the CLI
 
-Let's say that we want to train [Diffusion Policy](../lerobot/common/policies/diffusion) on the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using the [gym_pusht](https://github.com/huggingface/gym-pusht) environment for evaluation. The command to do so would look like this:
+Let's say that we want to train [Diffusion Policy](../src/lerobot/policies/diffusion) on the [pusht](https://huggingface.co/datasets/lerobot/pusht) dataset, using the [gym_pusht](https://github.com/huggingface/gym-pusht) environment for evaluation. The command to do so would look like this:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
     --dataset.repo_id=lerobot/pusht \
     --policy.type=diffusion \
     --env.type=pusht
@@ -60,12 +60,12 @@ python lerobot/scripts/train.py \
 
 Let's break this down:
 - To specify the dataset, we just need to specify its `repo_id` on the hub which is the only required argument in the `DatasetConfig`. The rest of the fields have default values and in this case we are fine with those so we can just add the option `--dataset.repo_id=lerobot/pusht`.
-- To specify the policy, we can just select diffusion policy using `--policy` appended with `.type`. Here, `.type` is a special argument which allows us to select config classes inheriting from `draccus.ChoiceRegistry` and that have been decorated with the `register_subclass()` method. To have a better explanation of this feature, have a look at this [Draccus demo](https://github.com/dlwh/draccus?tab=readme-ov-file#more-flexible-configuration-with-choice-types). In our code, we use this mechanism mainly to select policies, environments, robots, and some other components like optimizers. The policies available to select are located in [lerobot/common/policies](../lerobot/common/policies)
-- Similarly, we select the environment with `--env.type=pusht`. The different environment configs are available in [`lerobot/common/envs/configs.py`](../lerobot/common/envs/configs.py)
+- To specify the policy, we can just select diffusion policy using `--policy` appended with `.type`. Here, `.type` is a special argument which allows us to select config classes inheriting from `draccus.ChoiceRegistry` and that have been decorated with the `register_subclass()` method. To have a better explanation of this feature, have a look at this [Draccus demo](https://github.com/dlwh/draccus?tab=readme-ov-file#more-flexible-configuration-with-choice-types). In our code, we use this mechanism mainly to select policies, environments, robots, and some other components like optimizers. The policies available to select are located in [lerobot/policies](../src/lerobot/policies)
+- Similarly, we select the environment with `--env.type=pusht`. The different environment configs are available in [`lerobot/envs/configs.py`](../src/lerobot/envs/configs.py)
 
-Let's see another example. Let's say you've been training [ACT](../lerobot/common/policies/act) on [lerobot/aloha_sim_insertion_human](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_human) using the [gym-aloha](https://github.com/huggingface/gym-aloha) environment for evaluation with:
+Let's see another example. Let's say you've been training [ACT](../src/lerobot/policies/act) on [lerobot/aloha_sim_insertion_human](https://huggingface.co/datasets/lerobot/aloha_sim_insertion_human) using the [gym-aloha](https://github.com/huggingface/gym-aloha) environment for evaluation with:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
     --policy.type=act \
     --dataset.repo_id=lerobot/aloha_sim_insertion_human \
     --env.type=aloha \
@@ -74,9 +74,9 @@ python lerobot/scripts/train.py \
 > Notice we added `--output_dir` to explicitly tell where to write outputs from this run (checkpoints, training state, configs etc.). This is not mandatory and if you don't specify it, a default directory will be created from the current date and time, env.type and policy.type. This will typically look like `outputs/train/2025-01-24/16-10-05_aloha_act`.
 
 We now want to train a different policy for aloha on another task. We'll change the dataset and use [lerobot/aloha_sim_transfer_cube_human](https://huggingface.co/datasets/lerobot/aloha_sim_transfer_cube_human) instead. Of course, we also need to change the task of the environment as well to match this other task.
-Looking at the [`AlohaEnv`](../lerobot/common/envs/configs.py) config, the task is `"AlohaInsertion-v0"` by default, which corresponds to the task we trained on in the command above. The [gym-aloha](https://github.com/huggingface/gym-aloha?tab=readme-ov-file#description) environment also has the `AlohaTransferCube-v0` task which corresponds to this other task we want to train on. Putting this together, we can train this new policy on this different task using:
+Looking at the [`AlohaEnv`](../src/lerobot/envs/configs.py) config, the task is `"AlohaInsertion-v0"` by default, which corresponds to the task we trained on in the command above. The [gym-aloha](https://github.com/huggingface/gym-aloha?tab=readme-ov-file#description) environment also has the `AlohaTransferCube-v0` task which corresponds to this other task we want to train on. Putting this together, we can train this new policy on this different task using:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
     --policy.type=act \
     --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
     --env.type=aloha \
@@ -111,7 +111,7 @@ Now, let's assume that we want to reproduce the run just above. That run has pro
 
 We can then simply load the config values from this file using:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
     --config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
     --output_dir=outputs/train/act_aloha_transfer_2
 ```
@@ -119,7 +119,7 @@ python lerobot/scripts/train.py \
 
 Similarly to Hydra, we can still override some parameters in the CLI if we want to, e.g.:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
     --config_path=outputs/train/act_aloha_transfer/checkpoints/last/pretrained_model/ \
     --output_dir=outputs/train/act_aloha_transfer_2
     --policy.n_action_steps=80
@@ -128,7 +128,7 @@ python lerobot/scripts/train.py \
 
 `--config_path` can also accept the repo_id of a repo on the hub that contains a `train_config.json` file, e.g. running:
 ```bash
-python lerobot/scripts/train.py --config_path=lerobot/diffusion_pusht
+python -m lerobot.scripts.train --config_path=lerobot/diffusion_pusht
 ```
 will start a training run with the same configuration used for training [lerobot/diffusion_pusht](https://huggingface.co/lerobot/diffusion_pusht)
 
@@ -139,7 +139,7 @@ Being able to resume a training run is important in case it crashed or aborted f
 
 Let's reuse the command from the previous run and add a few more options:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
     --policy.type=act \
     --dataset.repo_id=lerobot/aloha_sim_transfer_cube_human \
     --env.type=aloha \
@@ -155,7 +155,7 @@ INFO 2025-01-24 16:10:56 ts/train.py:263 Checkpoint policy after step 100
 ```
 Now let's simulate a crash by killing the process (hit `ctrl`+`c`). We can then simply resume this run from the last checkpoint available with:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
     --config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
     --resume=true
 ```
@@ -164,7 +164,7 @@ You should see from the logging that your training picks up from where it left o
 Another reason for which you might want to resume a run is simply to extend training and add more training steps. The number of training steps is set by the option `--steps`, which is 100 000 by default.
 You could double the number of steps of the previous run with:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
     --config_path=outputs/train/run_resumption/checkpoints/last/pretrained_model/ \
     --resume=true \
     --steps=200000
@@ -195,7 +195,7 @@ In addition to the features currently in Draccus, we've added a special `.path`
 
 For example, we could fine-tune a [policy pre-trained on the aloha transfer task](https://huggingface.co/lerobot/act_aloha_sim_transfer_cube_human) on the aloha insertion task. We can achieve this with:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
     --policy.path=lerobot/act_aloha_sim_transfer_cube_human \
     --dataset.repo_id=lerobot/aloha_sim_insertion_human \
     --env.type=aloha \
@@ -236,7 +236,7 @@ We'll summarize here the main use cases to remember from this tutorial.
 
 #### Train a policy from scratch – CLI
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
     --policy.type=act \  # <- select 'act' policy
     --env.type=pusht \  # <- select 'pusht' environment
     --dataset.repo_id=lerobot/pusht  # <- train on this dataset
@@ -244,14 +244,14 @@ python lerobot/scripts/train.py \
 
 #### Train a policy from scratch - config file + CLI
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
     --config_path=path/to/pretrained_model \  # <- can also be a repo_id
     --policy.n_action_steps=80  # <- you may still override values
 ```
 
 #### Resume/continue a training run
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
     --config_path=checkpoint/pretrained_model/ \
     --resume=true \
     --steps=200000  # <- you can change some training parameters
@@ -259,7 +259,7 @@ python lerobot/scripts/train.py \
 
 #### Fine-tuning
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
     --policy.path=lerobot/act_aloha_sim_transfer_cube_human \  # <- can also be a local path to a checkpoint
     --dataset.repo_id=lerobot/aloha_sim_insertion_human \
     --env.type=aloha \
diff --git a/examples/advanced/1_add_image_transforms.py b/examples/advanced/1_add_image_transforms.py
index f14609261..3760feabb 100644
--- a/examples/advanced/1_add_image_transforms.py
+++ b/examples/advanced/1_add_image_transforms.py
@@ -22,7 +22,7 @@ from pathlib import Path
 
 from torchvision.transforms import ToPILImage, v2
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
 
 dataset_repo_id = "lerobot/aloha_static_screw_driver"
 
diff --git a/examples/advanced/2_calculate_validation_loss.py b/examples/advanced/2_calculate_validation_loss.py
index aac8e2e4e..9eeb1a2d9 100644
--- a/examples/advanced/2_calculate_validation_loss.py
+++ b/examples/advanced/2_calculate_validation_loss.py
@@ -26,8 +26,8 @@ import math
 
 import torch
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
-from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy
+from lerobot.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
 
 
 def main():
diff --git a/examples/backward_compatibility/replay.py b/examples/backward_compatibility/replay.py
index 11684d064..cc3397543 100644
--- a/examples/backward_compatibility/replay.py
+++ b/examples/backward_compatibility/replay.py
@@ -35,8 +35,8 @@ from pprint import pformat
 
 import draccus
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.robots import (  # noqa: F401
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.robots import (  # noqa: F401
     Robot,
     RobotConfig,
     koch_follower,
@@ -44,8 +44,8 @@ from lerobot.common.robots import (  # noqa: F401
     so100_follower,
     so101_follower,
 )
-from lerobot.common.utils.robot_utils import busy_wait
-from lerobot.common.utils.utils import (
+from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.utils import (
     init_logging,
     log_say,
 )
diff --git a/examples/lekiwi/evaluate.py b/examples/lekiwi/evaluate.py
index 2a41440a3..57fb62e10 100644
--- a/examples/lekiwi/evaluate.py
+++ b/examples/lekiwi/evaluate.py
@@ -1,32 +1,90 @@
-from lerobot.common.datasets.utils import build_dataset_frame, hw_to_dataset_features
-from lerobot.common.policies.act.modeling_act import ACTPolicy
-from lerobot.common.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
-from lerobot.common.utils.control_utils import predict_action
-from lerobot.common.utils.utils import get_safe_torch_device
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import hw_to_dataset_features
+from lerobot.policies.act.modeling_act import ACTPolicy
+from lerobot.record import record_loop
+from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
+from lerobot.utils.control_utils import init_keyboard_listener
+from lerobot.utils.utils import log_say
+from lerobot.utils.visualization_utils import _init_rerun
 
-NB_CYCLES_CLIENT_CONNECTION = 1000
+NUM_EPISODES = 2
+FPS = 30
+EPISODE_TIME_SEC = 60
+TASK_DESCRIPTION = "My task description"
 
+# Create the robot and teleoperator configurations
 robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
 robot = LeKiwiClient(robot_config)
 
+policy = ACTPolicy.from_pretrained("<hf_username>/<policy_repo_id>")
+
+# Configure the dataset features
+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
+dataset = LeRobotDataset.create(
+    repo_id="<hf_username>/<eval_dataset_repo_id>",
+    fps=FPS,
+    features=dataset_features,
+    robot_type=robot.name,
+    use_videos=True,
+    image_writer_threads=4,
+)
+
+# To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
 robot.connect()
 
-policy = ACTPolicy.from_pretrained("pepijn223/act_lekiwi_circle")
-policy.reset()
+_init_rerun(session_name="recording")
 
-obs_features = hw_to_dataset_features(robot.observation_features, "observation")
+listener, events = init_keyboard_listener()
 
-print("Running inference")
-i = 0
-while i < NB_CYCLES_CLIENT_CONNECTION:
-    obs = robot.get_observation()
+if not robot.is_connected:
+    raise ValueError("Robot is not connected!")
 
-    observation_frame = build_dataset_frame(obs_features, obs, prefix="observation")
-    action_values = predict_action(
-        observation_frame, policy, get_safe_torch_device(policy.config.device), policy.config.use_amp
+recorded_episodes = 0
+while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
+    log_say(f"Running inference, recording eval episode {recorded_episodes} of {NUM_EPISODES}")
+
+    # Run the policy inference loop
+    record_loop(
+        robot=robot,
+        events=events,
+        fps=FPS,
+        policy=policy,
+        dataset=dataset,
+        control_time_s=EPISODE_TIME_SEC,
+        single_task=TASK_DESCRIPTION,
+        display_data=True,
     )
-    action = {key: action_values[i].item() for i, key in enumerate(robot.action_features)}
-    robot.send_action(action)
-    i += 1
+
+    # Logic for reset env
+    if not events["stop_recording"] and (
+        (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"]
+    ):
+        log_say("Reset the environment")
+        record_loop(
+            robot=robot,
+            events=events,
+            fps=FPS,
+            control_time_s=EPISODE_TIME_SEC,
+            single_task=TASK_DESCRIPTION,
+            display_data=True,
+        )
+
+    if events["rerecord_episode"]:
+        log_say("Re-record episode")
+        events["rerecord_episode"] = False
+        events["exit_early"] = False
+        dataset.clear_episode_buffer()
+        continue
+
+    dataset.save_episode()
+    recorded_episodes += 1
+
+# Upload to hub and clean up
+dataset.push_to_hub()
 
 robot.disconnect()
+listener.stop()
diff --git a/examples/lekiwi/record.py b/examples/lekiwi/record.py
index 405a41bd3..11a716761 100644
--- a/examples/lekiwi/record.py
+++ b/examples/lekiwi/record.py
@@ -1,67 +1,101 @@
-import time
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import hw_to_dataset_features
+from lerobot.record import record_loop
+from lerobot.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
+from lerobot.robots.lekiwi.lekiwi_client import LeKiwiClient
+from lerobot.teleoperators.keyboard import KeyboardTeleop, KeyboardTeleopConfig
+from lerobot.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
+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.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.datasets.utils import hw_to_dataset_features
-from lerobot.common.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
-from lerobot.common.robots.lekiwi.lekiwi_client import LeKiwiClient
-from lerobot.common.teleoperators.keyboard import KeyboardTeleop, KeyboardTeleopConfig
-from lerobot.common.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
-
-NB_CYCLES_CLIENT_CONNECTION = 250
-
-leader_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem58760431551")
-leader_arm = SO100Leader(leader_arm_config)
+NUM_EPISODES = 3
+FPS = 30
+EPISODE_TIME_SEC = 30
+RESET_TIME_SEC = 10
+TASK_DESCRIPTION = "My task description"
 
+# Create the robot and teleoperator configurations
+robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
+leader_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm")
 keyboard_config = KeyboardTeleopConfig()
+
+robot = LeKiwiClient(robot_config)
+leader_arm = SO100Leader(leader_arm_config)
 keyboard = KeyboardTeleop(keyboard_config)
 
-robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
-robot = LeKiwiClient(robot_config)
-
+# Configure the dataset features
 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
 dataset = LeRobotDataset.create(
-    repo_id="pepijn223/lekiwi" + str(int(time.time())),
-    fps=10,
+    repo_id="<hf_username>/<dataset_repo_id>",
+    fps=FPS,
     features=dataset_features,
     robot_type=robot.name,
+    use_videos=True,
+    image_writer_threads=4,
 )
 
+# To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
+robot.connect()
 leader_arm.connect()
 keyboard.connect()
-robot.connect()
+
+_init_rerun(session_name="lekiwi_record")
+
+listener, events = init_keyboard_listener()
 
 if not robot.is_connected or not leader_arm.is_connected or not keyboard.is_connected:
-    exit()
+    raise ValueError("Robot, leader arm of keyboard is not connected!")
 
-print("Starting LeKiwi recording")
-i = 0
-while i < NB_CYCLES_CLIENT_CONNECTION:
-    arm_action = leader_arm.get_action()
-    arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
+recorded_episodes = 0
+while recorded_episodes < NUM_EPISODES and not events["stop_recording"]:
+    log_say(f"Recording episode {recorded_episodes}")
 
-    keyboard_keys = keyboard.get_action()
+    # Run the record loop
+    record_loop(
+        robot=robot,
+        events=events,
+        fps=FPS,
+        dataset=dataset,
+        teleop=[leader_arm, keyboard],
+        control_time_s=EPISODE_TIME_SEC,
+        single_task=TASK_DESCRIPTION,
+        display_data=True,
+    )
 
-    base_action = robot._from_keyboard_to_base_action(keyboard_keys)
+    # Logic for reset env
+    if not events["stop_recording"] and (
+        (recorded_episodes < NUM_EPISODES - 1) or events["rerecord_episode"]
+    ):
+        log_say("Reset the environment")
+        record_loop(
+            robot=robot,
+            events=events,
+            fps=FPS,
+            teleop=[leader_arm, keyboard],
+            control_time_s=RESET_TIME_SEC,
+            single_task=TASK_DESCRIPTION,
+            display_data=True,
+        )
 
-    action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
+    if events["rerecord_episode"]:
+        log_say("Re-record episode")
+        events["rerecord_episode"] = False
+        events["exit_early"] = False
+        dataset.clear_episode_buffer()
+        continue
 
-    action_sent = robot.send_action(action)
-    observation = robot.get_observation()
+    dataset.save_episode()
+    recorded_episodes += 1
 
-    frame = {**action_sent, **observation}
-    task = "Dummy Example Task Dataset"
+# Upload to hub and clean up
+dataset.push_to_hub()
 
-    dataset.add_frame(frame, task)
-    i += 1
-
-print("Disconnecting Teleop Devices and LeKiwi Client")
 robot.disconnect()
 leader_arm.disconnect()
 keyboard.disconnect()
-
-print("Uploading dataset to the hub")
-dataset.save_episode()
-dataset.push_to_hub()
+listener.stop()
diff --git a/examples/lekiwi/replay.py b/examples/lekiwi/replay.py
index f69092de0..248354df9 100644
--- a/examples/lekiwi/replay.py
+++ b/examples/lekiwi/replay.py
@@ -1,25 +1,33 @@
 import time
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.robots.lekiwi.config_lekiwi import LeKiwiClientConfig
-from lerobot.common.robots.lekiwi.lekiwi_client import LeKiwiClient
-from lerobot.common.utils.robot_utils import busy_wait
+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.robot_utils import busy_wait
+from lerobot.utils.utils import log_say
+
+EPISODE_IDX = 0
 
 robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="lekiwi")
 robot = LeKiwiClient(robot_config)
 
-dataset = LeRobotDataset("pepijn223/lekiwi1749025613", episodes=[0])
+dataset = LeRobotDataset("<hf_username>/<dataset_repo_id>", episodes=[EPISODE_IDX])
+actions = dataset.hf_dataset.select_columns("action")
 
 robot.connect()
 
-print("Replaying episode…")
-for _, action_array in enumerate(dataset.hf_dataset["action"]):
+if not robot.is_connected:
+    raise ValueError("Robot is not connected!")
+
+log_say(f"Replaying episode {EPISODE_IDX}")
+for idx in range(dataset.num_frames):
     t0 = time.perf_counter()
 
-    action = {name: float(action_array[i]) for i, name in enumerate(dataset.features["action"]["names"])}
+    action = {
+        name: float(actions[idx]["action"][i]) for i, name in enumerate(dataset.features["action"]["names"])
+    }
     robot.send_action(action)
 
     busy_wait(max(1.0 / dataset.fps - (time.perf_counter() - t0), 0.0))
 
-print("Disconnecting LeKiwi Client")
 robot.disconnect()
diff --git a/examples/lekiwi/teleoperate.py b/examples/lekiwi/teleoperate.py
index 2fe85d94e..8358a2b93 100644
--- a/examples/lekiwi/teleoperate.py
+++ b/examples/lekiwi/teleoperate.py
@@ -1,32 +1,47 @@
-from lerobot.common.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig
-from lerobot.common.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop, KeyboardTeleopConfig
-from lerobot.common.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
+import time
 
+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
+
+FPS = 30
+
+# Create the robot and teleoperator configurations
 robot_config = LeKiwiClientConfig(remote_ip="172.18.134.136", id="my_lekiwi")
-
-teleop__arm_config = SO100LeaderConfig(
-    port="/dev/tty.usbmodem58760431551",
-    id="my_awesome_leader_arm",
-)
-
-teleop_keyboard_config = KeyboardTeleopConfig(
-    id="my_laptop_keyboard",
-)
+teleop_arm_config = SO100LeaderConfig(port="/dev/tty.usbmodem585A0077581", id="my_awesome_leader_arm")
+keyboard_config = KeyboardTeleopConfig(id="my_laptop_keyboard")
 
 robot = LeKiwiClient(robot_config)
-teleop_arm = SO100Leader(teleop__arm_config)
-telep_keyboard = KeyboardTeleop(teleop_keyboard_config)
+leader_arm = SO100Leader(teleop_arm_config)
+keyboard = KeyboardTeleop(keyboard_config)
+
+# To connect you already should have this script running on LeKiwi: `python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi`
 robot.connect()
-teleop_arm.connect()
-telep_keyboard.connect()
+leader_arm.connect()
+keyboard.connect()
+
+_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, leader arm of keyboard is not connected!")
 
 while True:
+    t0 = time.perf_counter()
+
     observation = robot.get_observation()
 
-    arm_action = teleop_arm.get_action()
+    arm_action = leader_arm.get_action()
     arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
 
-    keyboard_keys = telep_keyboard.get_action()
+    keyboard_keys = keyboard.get_action()
     base_action = robot._from_keyboard_to_base_action(keyboard_keys)
 
-    robot.send_action(arm_action | base_action)
+    log_rerun_data(observation, {**arm_action, **base_action})
+
+    action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
+
+    robot.send_action(action)
+
+    busy_wait(max(1.0 / FPS - (time.perf_counter() - t0), 0.0))
diff --git a/lerobot/common/datasets/lerobot_dataset.py b/lerobot/common/datasets/lerobot_dataset.py
deleted file mode 100644
index a8b98b127..000000000
--- a/lerobot/common/datasets/lerobot_dataset.py
+++ /dev/null
@@ -1,1333 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import contextlib
-import logging
-import shutil
-import os
-from pathlib import Path
-from typing import Callable
-
-import datasets
-import numpy as np
-import packaging.version
-import PIL.Image
-import torch
-import torch.utils
-from datasets import concatenate_datasets, load_dataset
-from huggingface_hub import HfApi, snapshot_download
-from huggingface_hub.constants import REPOCARD_NAME
-from huggingface_hub.errors import RevisionNotFoundError
-
-from lerobot.common.constants import HF_LEROBOT_HOME
-from lerobot.common.datasets.compute_stats import aggregate_stats, compute_episode_stats #aggregate_stats_per_robot_type,
-from lerobot.common.datasets.image_writer import AsyncImageWriter, write_image
-from lerobot.common.datasets.utils import (
-    DEFAULT_FEATURES,
-    DEFAULT_IMAGE_PATH,
-    INFO_PATH,
-    TASKS_PATH,
-    _validate_feature_names,
-    append_jsonlines,
-    backward_compatible_episodes_stats,
-    check_delta_timestamps,
-    check_timestamps_sync,
-    check_version_compatibility,
-    create_empty_dataset_info,
-    create_lerobot_dataset_card,
-    embed_images,
-    get_delta_indices,
-    get_episode_data_index,
-    get_features_from_robot,
-    get_hf_features_from_features,
-    get_safe_version,
-    hf_transform_to_torch,
-    is_valid_version,
-    load_episodes,
-    load_episodes_stats,
-    load_info,
-    load_stats,
-    load_tasks,
-    map_dict_keys,
-    validate_episode_buffer,
-    validate_frame,
-    write_episode,
-    write_episode_stats,
-    write_info,
-    write_json,
-    #keep_datasets_with_the_same_features_per_robot_type,
-    #map_dict_pad_keys,
-    #keep_datasets_with_valid_fps,
-    #find_start_of_motion,
-)
-from lerobot.common.datasets.video_utils import (
-    VideoFrame,
-    decode_video_frames,
-    encode_video_frames,
-    get_safe_default_codec,
-    get_video_info,
-)
-
-#from lerobot.common.robot_devices.robots.utils import Robot
-from lerobot.configs.datasets import ROBOT_TYPE_KEYS_MAPPING, TASKS_KEYS_MAPPING
-#FIXME: remove this import
-from lerobot.common.datasets.collators import pad_tensor
-
-CODEBASE_VERSION = "v2.1"
-LEROBOT_HOME = Path(os.getenv("LEROBOT_HOME", "~/.cache/huggingface/lerobot")).expanduser()
-
-
-def find_start_of_motion(velocities, window_size, threshold, motion_buffer):
-    for t in range(len(velocities) - window_size):
-        window_mean = velocities[t:t+window_size].mean()
-        if window_mean > threshold:
-            return max(0, t - motion_buffer)  # include slight context before motion
-    return 0
-
-class LeRobotDatasetMetadata:
-    def __init__(
-        self,
-        repo_id: str,
-        root: str | Path | None = None,
-        local_files_only: bool = False,
-        feature_keys_mapping: dict[str, str] | None = None,
-        revision: str | None = None,
-        force_cache_sync: bool = False,
-    ):
-        self.repo_id = repo_id
-        self.local_files_only = local_files_only
-        self.revision = revision if revision else CODEBASE_VERSION
-        self.root = Path(root) if root is not None else HF_LEROBOT_HOME / repo_id
-
-        try:
-            if force_cache_sync:
-                raise FileNotFoundError
-            self.load_metadata()
-        except (FileNotFoundError, NotADirectoryError):
-            if is_valid_version(self.revision):
-                self.revision = get_safe_version(self.repo_id, self.revision)
-
-            (self.root / "meta").mkdir(exist_ok=True, parents=True)
-            self.pull_from_repo(allow_patterns="meta/")
-            self.load_metadata()
-        self.feature_keys_mapping = feature_keys_mapping.get(repo_id, None) if feature_keys_mapping else None
-        self.inverse_feature_keys_mapping = (
-            {v: k for k, v in self.feature_keys_mapping.items() if v} if self.feature_keys_mapping else {}
-        )
-        self.info["features"] = map_dict_keys(
-            self.info["features"], feature_keys_mapping=self.feature_keys_mapping
-        )
-    def load_metadata(self):
-        self.info = load_info(self.root)
-        check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
-        self.tasks, self.task_to_task_index = load_tasks(self.root)
-        self.episodes = load_episodes(self.root)
-        if self._version < packaging.version.parse("v2.1"):
-            self.stats = load_stats(self.root)
-            self.episodes_stats = backward_compatible_episodes_stats(self.stats, self.episodes)
-        else:
-            self.episodes_stats = load_episodes_stats(self.root)
-            self.stats = aggregate_stats(list(self.episodes_stats.values()))
-
-    def pull_from_repo(
-        self,
-        allow_patterns: list[str] | str | None = None,
-        ignore_patterns: list[str] | str | None = None,
-    ) -> None:
-        snapshot_download(
-            self.repo_id,
-            repo_type="dataset",
-            revision=self.revision,
-            local_dir=self.root,
-            allow_patterns=allow_patterns,
-            ignore_patterns=ignore_patterns,
-        )
-
-    @property
-    def _version(self) -> packaging.version.Version:
-        """Codebase version used to create this dataset."""
-        return packaging.version.parse(self.info["codebase_version"])
-
-    def get_data_file_path(self, ep_index: int) -> Path:
-        ep_chunk = self.get_episode_chunk(ep_index)
-        fpath = self.data_path.format(episode_chunk=ep_chunk, episode_index=ep_index)
-        return Path(fpath)
-
-    def get_video_file_path(self, ep_index: int, vid_key: str) -> Path:
-        ep_chunk = self.get_episode_chunk(ep_index)
-        fpath = self.video_path.format(episode_chunk=ep_chunk, video_key=vid_key, episode_index=ep_index)
-        return Path(fpath)
-
-    def get_episode_chunk(self, ep_index: int) -> int:
-        return ep_index // self.chunks_size
-
-    @property
-    def data_path(self) -> str:
-        """Formattable string for the parquet files."""
-        return self.info["data_path"]
-
-    @property
-    def video_path(self) -> str | None:
-        """Formattable string for the video files."""
-        return self.info["video_path"]
-
-    @property
-    def robot_type(self) -> str | None:
-        """Robot type used in recording this dataset."""
-        return self.info["robot_type"]
-
-    @property
-    def fps(self) -> int:
-        """Frames per second used during data collection."""
-        return self.info["fps"]
-
-    @property
-    def features(self) -> dict[str, dict]:
-        """All features contained in the dataset."""
-        return self.info["features"]
-
-    @property
-    def image_keys(self) -> list[str]:
-        """Keys to access visual modalities stored as images."""
-        return [key for key, ft in self.features.items() if ft["dtype"] == "image"]
-
-    @property
-    def video_keys(self) -> list[str]:
-        """Keys to access visual modalities stored as videos."""
-        return [key for key, ft in self.features.items() if ft["dtype"] == "video"]
-
-    @property
-    def camera_keys(self) -> list[str]:
-        """Keys to access visual modalities (regardless of their storage method)."""
-        return [key for key, ft in self.features.items() if ft["dtype"] in ["video", "image"]]
-
-    @property
-    def names(self) -> dict[str, list | dict]:
-        """Names of the various dimensions of vector modalities."""
-        return {key: ft["names"] for key, ft in self.features.items()}
-
-    @property
-    def shapes(self) -> dict:
-        """Shapes for the different features."""
-        return {key: tuple(ft["shape"]) for key, ft in self.features.items()}
-
-    @property
-    def total_episodes(self) -> int:
-        """Total number of episodes available."""
-        return self.info["total_episodes"]
-
-    @property
-    def total_frames(self) -> int:
-        """Total number of frames saved in this dataset."""
-        return self.info["total_frames"]
-
-    @property
-    def total_tasks(self) -> int:
-        """Total number of different tasks performed in this dataset."""
-        return self.info["total_tasks"]
-
-    @property
-    def total_chunks(self) -> int:
-        """Total number of chunks (groups of episodes)."""
-        return self.info["total_chunks"]
-
-    @property
-    def chunks_size(self) -> int:
-        """Max number of episodes per chunk."""
-        return self.info["chunks_size"]
-
-    def get_task_index(self, task: str) -> int | None:
-        """
-        Given a task in natural language, returns its task_index if the task already exists in the dataset,
-        otherwise return None.
-        """
-        return self.task_to_task_index.get(task, None)
-
-    def add_task(self, task: str):
-        """
-        Given a task in natural language, add it to the dictionary of tasks.
-        """
-        if task in self.task_to_task_index:
-            raise ValueError(f"The task '{task}' already exists and can't be added twice.")
-
-        task_index = self.info["total_tasks"]
-        self.task_to_task_index[task] = task_index
-        self.tasks[task_index] = task
-        self.info["total_tasks"] += 1
-
-        task_dict = {
-            "task_index": task_index,
-            "task": task,
-        }
-        append_jsonlines(task_dict, self.root / TASKS_PATH)
-
-    def save_episode(
-        self,
-        episode_index: int,
-        episode_length: int,
-        episode_tasks: list[str],
-        episode_stats: dict[str, dict],
-    ) -> None:
-        self.info["total_episodes"] += 1
-        self.info["total_frames"] += episode_length
-
-        chunk = self.get_episode_chunk(episode_index)
-        if chunk >= self.total_chunks:
-            self.info["total_chunks"] += 1
-
-        self.info["splits"] = {"train": f"0:{self.info['total_episodes']}"}
-        self.info["total_videos"] += len(self.video_keys)
-        if len(self.video_keys) > 0:
-            self.update_video_info()
-
-        write_info(self.info, self.root)
-
-        episode_dict = {
-            "episode_index": episode_index,
-            "tasks": episode_tasks,
-            "length": episode_length,
-        }
-        self.episodes[episode_index] = episode_dict
-        write_episode(episode_dict, self.root)
-
-        self.episodes_stats[episode_index] = episode_stats
-        self.stats = aggregate_stats([self.stats, episode_stats]) if self.stats else episode_stats
-        write_episode_stats(episode_index, episode_stats, self.root)
-
-    def update_video_info(self) -> None:
-        """
-        Warning: this function writes info from first episode videos, implicitly assuming that all videos have
-        been encoded the same way. Also, this means it assumes the first episode exists.
-        """
-        for key in self.video_keys:
-            if not self.features[key].get("info", None):
-                video_path = self.root / self.get_video_file_path(ep_index=0, vid_key=key)
-                self.info["features"][key]["info"] = get_video_info(video_path)
-
-    def __repr__(self):
-        feature_keys = list(self.features)
-        return (
-            f"{self.__class__.__name__}({{\n"
-            f"    Repository ID: '{self.repo_id}',\n"
-            f"    Total episodes: '{self.total_episodes}',\n"
-            f"    Total frames: '{self.total_frames}',\n"
-            f"    Features: '{feature_keys}',\n"
-            "})',\n"
-        )
-
-    @classmethod
-    def create(
-        cls,
-        repo_id: str,
-        fps: int,
-        features: dict,
-        robot_type: str | None = None,
-        root: str | Path | None = None,
-        use_videos: bool = True,
-    ) -> "LeRobotDatasetMetadata":
-        """Creates metadata for a LeRobotDataset."""
-        obj = cls.__new__(cls)
-        obj.repo_id = repo_id
-        obj.root = Path(root) if root is not None else HF_LEROBOT_HOME / repo_id
-
-        obj.root.mkdir(parents=True, exist_ok=False)
-
-        # TODO(aliberts, rcadene): implement sanity check for features
-        features = {**features, **DEFAULT_FEATURES}
-        _validate_feature_names(features)
-
-        obj.tasks, obj.task_to_task_index = {}, {}
-        obj.episodes_stats, obj.stats, obj.episodes = {}, {}, {}
-        obj.info = create_empty_dataset_info(CODEBASE_VERSION, fps, features, use_videos, robot_type)
-        if len(obj.video_keys) > 0 and not use_videos:
-            raise ValueError()
-        write_json(obj.info, obj.root / INFO_PATH)
-        obj.revision = None
-        return obj
-
-
-class LeRobotDataset(torch.utils.data.Dataset):
-    def __init__(
-        self,
-        repo_id: str,
-        root: str | Path | None = None,
-        episodes: list[int] | None = None,
-        image_transforms: Callable | None = None,
-        delta_timestamps: dict[list[float]] | None = None,
-        tolerance_s: float = 1e-4,
-        revision: str | None = None,
-        force_cache_sync: bool = False,
-        download_videos: bool = True,
-        video_backend: str | None = None,
-    ):
-        """
-        2 modes are available for instantiating this class, depending on 2 different use cases:
-
-        1. Your dataset already exists:
-            - On your local disk in the 'root' folder. This is typically the case when you recorded your
-              dataset locally and you may or may not have pushed it to the hub yet. Instantiating this class
-              with 'root' will load your dataset directly from disk. This can happen while you're offline (no
-              internet connection).
-
-            - On the Hugging Face Hub at the address https://huggingface.co/datasets/{repo_id} and not on
-              your local disk in the 'root' folder. Instantiating this class with this 'repo_id' will download
-              the dataset from that address and load it, pending your dataset is compliant with
-              codebase_version v2.0. If your dataset has been created before this new format, you will be
-              prompted to convert it using our conversion script from v1.6 to v2.0, which you can find at
-              lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py.
-
-
-        2. Your dataset doesn't already exists (either on local disk or on the Hub): you can create an empty
-           LeRobotDataset with the 'create' classmethod. This can be used for recording a dataset or port an
-           existing dataset to the LeRobotDataset format.
-
-
-        In terms of files, LeRobotDataset encapsulates 3 main things:
-            - metadata:
-                - info contains various information about the dataset like shapes, keys, fps etc.
-                - stats stores the dataset statistics of the different modalities for normalization
-                - tasks contains the prompts for each task of the dataset, which can be used for
-                  task-conditioned training.
-            - hf_dataset (from datasets.Dataset), which will read any values from parquet files.
-            - videos (optional) from which frames are loaded to be synchronous with data from parquet files.
-
-        A typical LeRobotDataset looks like this from its root path:
-        .
-        ├── data
-        │   ├── chunk-000
-        │   │   ├── episode_000000.parquet
-        │   │   ├── episode_000001.parquet
-        │   │   ├── episode_000002.parquet
-        │   │   └── ...
-        │   ├── chunk-001
-        │   │   ├── episode_001000.parquet
-        │   │   ├── episode_001001.parquet
-        │   │   ├── episode_001002.parquet
-        │   │   └── ...
-        │   └── ...
-        ├── meta
-        │   ├── episodes.jsonl
-        │   ├── info.json
-        │   ├── stats.json
-        │   └── tasks.jsonl
-        └── videos
-            ├── chunk-000
-            │   ├── observation.images.laptop
-            │   │   ├── episode_000000.mp4
-            │   │   ├── episode_000001.mp4
-            │   │   ├── episode_000002.mp4
-            │   │   └── ...
-            │   ├── observation.images.phone
-            │   │   ├── episode_000000.mp4
-            │   │   ├── episode_000001.mp4
-            │   │   ├── episode_000002.mp4
-            │   │   └── ...
-            ├── chunk-001
-            └── ...
-
-        Note that this file-based structure is designed to be as versatile as possible. The files are split by
-        episodes which allows a more granular control over which episodes one wants to use and download. The
-        structure of the dataset is entirely described in the info.json file, which can be easily downloaded
-        or viewed directly on the hub before downloading any actual data. The type of files used are very
-        simple and do not need complex tools to be read, it only uses .parquet, .json and .mp4 files (and .md
-        for the README).
-
-        Args:
-            repo_id (str): This is the repo id that will be used to fetch the dataset. Locally, the dataset
-                will be stored under root/repo_id.
-            root (Path | None, optional): Local directory to use for downloading/writing files. You can also
-                set the LEROBOT_HOME environment variable to point to a different location. Defaults to
-                '~/.cache/huggingface/lerobot'.
-            episodes (list[int] | None, optional): If specified, this will only load episodes specified by
-                their episode_index in this list. Defaults to None.
-            image_transforms (Callable | None, optional): You can pass standard v2 image transforms from
-                torchvision.transforms.v2 here which will be applied to visual modalities (whether they come
-                from videos or images). Defaults to None.
-            delta_timestamps (dict[list[float]] | None, optional): _description_. Defaults to None.
-            tolerance_s (float, optional): Tolerance in seconds used to ensure data timestamps are actually in
-                sync with the fps value. It is used at the init of the dataset to make sure that each
-                timestamps is separated to the next by 1/fps +/- tolerance_s. This also applies to frames
-                decoded from video files. It is also used to check that `delta_timestamps` (when provided) are
-                multiples of 1/fps. Defaults to 1e-4.
-            revision (str, optional): An optional Git revision id which can be a branch name, a tag, or a
-                commit hash. Defaults to current codebase version tag.
-            sync_cache_first (bool, optional): Flag to sync and refresh local files first. If True and files
-                are already present in the local cache, this will be faster. However, files loaded might not
-                be in sync with the version on the hub, especially if you specified 'revision'. Defaults to
-                False.
-            download_videos (bool, optional): Flag to download the videos. Note that when set to True but the
-                video files are already present on local disk, they won't be downloaded again. Defaults to
-                True.
-            video_backend (str | None, optional): Video backend to use for decoding videos. Defaults to torchcodec when available int the platform; otherwise, defaults to 'pyav'.
-                You can also use the 'pyav' decoder used by Torchvision, which used to be the default option, or 'video_reader' which is another decoder of Torchvision.
-        """
-        super().__init__()
-        self.repo_id = repo_id
-        self.root = Path(root) if root else HF_LEROBOT_HOME / repo_id
-        self.image_transforms = image_transforms
-        self.delta_timestamps = delta_timestamps
-        self.episodes = episodes
-        self.tolerance_s = tolerance_s
-        self.revision = revision if revision else CODEBASE_VERSION
-        self.video_backend = video_backend if video_backend else get_safe_default_codec()
-        self.delta_indices = None
-
-        # Unused attributes
-        self.image_writer = None
-        self.episode_buffer = None
-
-        self.root.mkdir(exist_ok=True, parents=True)
-
-        # Load metadata
-        self.meta = LeRobotDatasetMetadata(
-            self.repo_id, self.root, self.revision, force_cache_sync=force_cache_sync
-        )
-        if self.episodes is not None and self.meta._version >= packaging.version.parse("v2.1"):
-            episodes_stats = [self.meta.episodes_stats[ep_idx] for ep_idx in self.episodes]
-            self.stats = aggregate_stats(episodes_stats)
-
-        # Load actual data
-        try:
-            if force_cache_sync:
-                raise FileNotFoundError
-            assert all((self.root / fpath).is_file() for fpath in self.get_episodes_file_paths())
-            self.hf_dataset = self.load_hf_dataset()
-        except (AssertionError, FileNotFoundError, NotADirectoryError):
-            self.revision = get_safe_version(self.repo_id, self.revision)
-            self.download_episodes(download_videos)
-            self.hf_dataset = self.load_hf_dataset()
-
-        self.episode_data_index = get_episode_data_index(self.meta.episodes, self.episodes)
-
-        # Check timestamps
-        timestamps = torch.stack(self.hf_dataset["timestamp"]).numpy()
-        episode_indices = torch.stack(self.hf_dataset["episode_index"]).numpy()
-        ep_data_index_np = {k: t.numpy() for k, t in self.episode_data_index.items()}
-        check_timestamps_sync(timestamps, episode_indices, ep_data_index_np, self.fps, self.tolerance_s)
-
-        # Setup delta_indices
-        if self.delta_timestamps is not None:
-            check_delta_timestamps(self.delta_timestamps, self.fps, self.tolerance_s)
-            self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps)
-
-    def push_to_hub(
-        self,
-        branch: str | None = None,
-        tags: list | None = None,
-        license: str | None = "apache-2.0",
-        tag_version: bool = True,
-        push_videos: bool = True,
-        private: bool = False,
-        allow_patterns: list[str] | str | None = None,
-        upload_large_folder: bool = False,
-        **card_kwargs,
-    ) -> None:
-        ignore_patterns = ["images/"]
-        if not push_videos:
-            ignore_patterns.append("videos/")
-
-        hub_api = HfApi()
-        hub_api.create_repo(
-            repo_id=self.repo_id,
-            private=private,
-            repo_type="dataset",
-            exist_ok=True,
-        )
-        if branch:
-            hub_api.create_branch(
-                repo_id=self.repo_id,
-                branch=branch,
-                revision=self.revision,
-                repo_type="dataset",
-                exist_ok=True,
-            )
-
-        upload_kwargs = {
-            "repo_id": self.repo_id,
-            "folder_path": self.root,
-            "repo_type": "dataset",
-            "revision": branch,
-            "allow_patterns": allow_patterns,
-            "ignore_patterns": ignore_patterns,
-        }
-        if upload_large_folder:
-            hub_api.upload_large_folder(**upload_kwargs)
-        else:
-            hub_api.upload_folder(**upload_kwargs)
-
-        if not hub_api.file_exists(self.repo_id, REPOCARD_NAME, repo_type="dataset", revision=branch):
-            card = create_lerobot_dataset_card(
-                tags=tags, dataset_info=self.meta.info, license=license, **card_kwargs
-            )
-            card.push_to_hub(repo_id=self.repo_id, repo_type="dataset", revision=branch)
-
-        if tag_version:
-            with contextlib.suppress(RevisionNotFoundError):
-                hub_api.delete_tag(self.repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
-            hub_api.create_tag(self.repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset")
-
-    def pull_from_repo(
-        self,
-        allow_patterns: list[str] | str | None = None,
-        ignore_patterns: list[str] | str | None = None,
-    ) -> None:
-        snapshot_download(
-            self.repo_id,
-            repo_type="dataset",
-            revision=self.revision,
-            local_dir=self.root,
-            allow_patterns=allow_patterns,
-            ignore_patterns=ignore_patterns,
-        )
-
-    def download_episodes(self, download_videos: bool = True) -> None:
-        """Downloads the dataset from the given 'repo_id' at the provided version. If 'episodes' is given, this
-        will only download those episodes (selected by their episode_index). If 'episodes' is None, the whole
-        dataset will be downloaded. Thanks to the behavior of snapshot_download, if the files are already present
-        in 'local_dir', they won't be downloaded again.
-        """
-        # TODO(rcadene, aliberts): implement faster transfer
-        # https://huggingface.co/docs/huggingface_hub/en/guides/download#faster-downloads
-        files = None
-        ignore_patterns = None if download_videos else "videos/"
-        if self.episodes is not None:
-            files = self.get_episodes_file_paths()
-
-        self.pull_from_repo(allow_patterns=files, ignore_patterns=ignore_patterns)
-
-    def get_episodes_file_paths(self) -> list[Path]:
-        episodes = self.episodes if self.episodes is not None else list(range(self.meta.total_episodes))
-        fpaths = [str(self.meta.get_data_file_path(ep_idx)) for ep_idx in episodes]
-        if len(self.meta.video_keys) > 0:
-            video_files = [
-                str(self.meta.get_video_file_path(ep_idx, vid_key))
-                for vid_key in self.meta.video_keys
-                for ep_idx in episodes
-            ]
-            fpaths += video_files
-
-        return fpaths
-
-    def load_hf_dataset(self) -> datasets.Dataset:
-        """hf_dataset contains all the observations, states, actions, rewards, etc."""
-        if self.episodes is None:
-            path = str(self.root / "data")
-            hf_dataset = load_dataset("parquet", data_dir=path, split="train")
-        else:
-            files = [str(self.root / self.meta.get_data_file_path(ep_idx)) for ep_idx in self.episodes]
-            hf_dataset = load_dataset("parquet", data_files=files, split="train")
-
-        # TODO(aliberts): hf_dataset.set_format("torch")
-        hf_dataset.set_transform(hf_transform_to_torch)
-        return hf_dataset
-
-    def create_hf_dataset(self) -> datasets.Dataset:
-        features = get_hf_features_from_features(self.features)
-        ft_dict = {col: [] for col in features}
-        hf_dataset = datasets.Dataset.from_dict(ft_dict, features=features, split="train")
-
-        # TODO(aliberts): hf_dataset.set_format("torch")
-        hf_dataset.set_transform(hf_transform_to_torch)
-        return hf_dataset
-
-    @property
-    def fps(self) -> int:
-        """Frames per second used during data collection."""
-        return self.meta.fps
-
-    @property
-    def num_frames(self) -> int:
-        """Number of frames in selected episodes."""
-        return len(self.hf_dataset) if self.hf_dataset is not None else self.meta.total_frames
-
-    @property
-    def num_episodes(self) -> int:
-        """Number of episodes selected."""
-        return len(self.episodes) if self.episodes is not None else self.meta.total_episodes
-
-    @property
-    def features(self) -> dict[str, dict]:
-        return self.meta.features
-
-    @property
-    def hf_features(self) -> datasets.Features:
-        """Features of the hf_dataset."""
-        if self.hf_dataset is not None:
-            return self.hf_dataset.features
-        else:
-            return get_hf_features_from_features(self.features)
-
-    def _get_query_indices(self, idx: int, ep_idx: int) -> tuple[dict[str, list[int | bool]]]:
-        ep_start = self.episode_data_index["from"][ep_idx]
-        ep_end = self.episode_data_index["to"][ep_idx]
-        query_indices = {
-            key: [max(ep_start.item(), min(ep_end.item() - 1, idx + delta)) for delta in delta_idx]
-            for key, delta_idx in self.delta_indices.items()
-        }
-        padding = {  # Pad values outside of current episode range
-            f"{key}_is_pad": torch.BoolTensor(
-                [(idx + delta < ep_start.item()) | (idx + delta >= ep_end.item()) for delta in delta_idx]
-            )
-            for key, delta_idx in self.delta_indices.items()
-        }
-        return query_indices, padding
-
-    def _get_query_timestamps(
-        self,
-        current_ts: float,
-        query_indices: dict[str, list[int]] | None = None,
-    ) -> dict[str, list[float]]:
-        query_timestamps = {}
-        for key in self.meta.video_keys:
-            if query_indices is not None and key in query_indices:
-                timestamps = self.hf_dataset.select(query_indices[key])["timestamp"]
-                query_timestamps[key] = torch.stack(timestamps).tolist()
-            else:
-                query_timestamps[key] = [current_ts]
-
-        return query_timestamps
-
-    def _query_hf_dataset(self, query_indices: dict[str, list[int]]) -> dict:
-        return {
-            key: torch.stack(self.hf_dataset.select(q_idx)[key])
-            for key, q_idx in query_indices.items()
-            if key not in self.meta.video_keys
-        }
-
-    def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict[str, torch.Tensor]:
-        """Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function
-        in the main process (e.g. by using a second Dataloader with num_workers=0). It will result in a
-        Segmentation Fault. This probably happens because a memory reference to the video loader is created in
-        the main process and a subprocess fails to access it.
-        """
-        item = {}
-        for vid_key, query_ts in query_timestamps.items():
-            video_path = self.root / self.meta.get_video_file_path(ep_idx, vid_key)
-            frames = decode_video_frames(video_path, query_ts, self.tolerance_s, self.video_backend)
-            item[vid_key] = frames.squeeze(0)
-
-        return item
-
-    def _add_padding_keys(self, item: dict, padding: dict[str, list[bool]]) -> dict:
-        for key, val in padding.items():
-            item[key] = torch.BoolTensor(val)
-        return item
-
-    def __len__(self):
-        return self.num_frames
-
-    def __getitem__(self, idx) -> dict:
-        item = self.hf_dataset[idx]
-        ep_idx = item["episode_index"].item()
-
-        query_indices = None
-        if self.delta_indices is not None:
-            query_indices, padding = self._get_query_indices(idx, ep_idx)
-            query_result = self._query_hf_dataset(query_indices)
-            item = {**item, **padding}
-            for key, val in query_result.items():
-                item[key] = val
-
-        if len(self.meta.video_keys) > 0:
-            current_ts = item["timestamp"].item()
-            query_timestamps = self._get_query_timestamps(current_ts, query_indices)
-            video_frames = self._query_videos(query_timestamps, ep_idx)
-            item = {**video_frames, **item}
-
-        if self.image_transforms is not None:
-            image_keys = self.meta.camera_keys
-            for cam in image_keys:
-                item[cam] = self.image_transforms(item[cam])
-
-        # Add task as a string
-        task_idx = item["task_index"].item()
-        item["task"] = self.meta.tasks[task_idx]
-
-        return item
-
-    def __repr__(self):
-        feature_keys = list(self.features)
-        return (
-            f"{self.__class__.__name__}({{\n"
-            f"    Repository ID: '{self.repo_id}',\n"
-            f"    Number of selected episodes: '{self.num_episodes}',\n"
-            f"    Number of selected samples: '{self.num_frames}',\n"
-            f"    Features: '{feature_keys}',\n"
-            "})',\n"
-        )
-
-    def create_episode_buffer(self, episode_index: int | None = None) -> dict:
-        current_ep_idx = self.meta.total_episodes if episode_index is None else episode_index
-        ep_buffer = {}
-        # size and task are special cases that are not in self.features
-        ep_buffer["size"] = 0
-        ep_buffer["task"] = []
-        for key in self.features:
-            ep_buffer[key] = current_ep_idx if key == "episode_index" else []
-        return ep_buffer
-
-    def _get_image_file_path(self, episode_index: int, image_key: str, frame_index: int) -> Path:
-        fpath = DEFAULT_IMAGE_PATH.format(
-            image_key=image_key, episode_index=episode_index, frame_index=frame_index
-        )
-        return self.root / fpath
-
-    def _save_image(self, image: torch.Tensor | np.ndarray | PIL.Image.Image, fpath: Path) -> None:
-        if self.image_writer is None:
-            if isinstance(image, torch.Tensor):
-                image = image.cpu().numpy()
-            write_image(image, fpath)
-        else:
-            self.image_writer.save_image(image=image, fpath=fpath)
-
-    def add_frame(self, frame: dict, task: str, timestamp: float | None = None) -> None:
-        """
-        This function only adds the frame to the episode_buffer. Apart from images — which are written in a
-        temporary directory — nothing is written to disk. To save those frames, the 'save_episode()' method
-        then needs to be called.
-        """
-        # Convert torch to numpy if needed
-        for name in frame:
-            if isinstance(frame[name], torch.Tensor):
-                frame[name] = frame[name].numpy()
-
-        validate_frame(frame, self.features)
-
-        if self.episode_buffer is None:
-            self.episode_buffer = self.create_episode_buffer()
-
-        # Automatically add frame_index and timestamp to episode buffer
-        frame_index = self.episode_buffer["size"]
-        if timestamp is None:
-            timestamp = frame_index / self.fps
-        self.episode_buffer["frame_index"].append(frame_index)
-        self.episode_buffer["timestamp"].append(timestamp)
-        self.episode_buffer["task"].append(task)
-
-        # Add frame features to episode_buffer
-        for key in frame:
-            
-            if key not in self.features:
-                raise ValueError(
-                    f"An element of the frame is not in the features. '{key}' not in '{self.features.keys()}'."
-                )
-
-            if self.features[key]["dtype"] in ["image", "video"]:
-                img_path = self._get_image_file_path(
-                    episode_index=self.episode_buffer["episode_index"], image_key=key, frame_index=frame_index
-                )
-                if frame_index == 0:
-                    img_path.parent.mkdir(parents=True, exist_ok=True)
-                self._save_image(frame[key], img_path)
-                self.episode_buffer[key].append(str(img_path))
-            else:
-                self.episode_buffer[key].append(frame[key])
-
-        self.episode_buffer["size"] += 1
-
-    def save_episode(self, episode_data: dict | None = None) -> None:
-        """
-        This will save to disk the current episode in self.episode_buffer.
-
-        Args:
-            episode_data (dict | None, optional): Dict containing the episode data to save. If None, this will
-                save the current episode in self.episode_buffer, which is filled with 'add_frame'. Defaults to
-                None.
-        """
-        if not episode_data:
-            episode_buffer = self.episode_buffer
-
-        validate_episode_buffer(episode_buffer, self.meta.total_episodes, self.features)
-
-        # size and task are special cases that won't be added to hf_dataset
-        episode_length = episode_buffer.pop("size")
-        tasks = episode_buffer.pop("task")
-        episode_tasks = list(set(tasks))
-        episode_index = episode_buffer["episode_index"]
-
-        episode_buffer["index"] = np.arange(self.meta.total_frames, self.meta.total_frames + episode_length)
-        episode_buffer["episode_index"] = np.full((episode_length,), episode_index)
-
-        # Add new tasks to the tasks dictionary
-        for task in episode_tasks:
-            task_index = self.meta.get_task_index(task)
-            if task_index is None:
-                self.meta.add_task(task)
-
-        # Given tasks in natural language, find their corresponding task indices
-        episode_buffer["task_index"] = np.array([self.meta.get_task_index(task) for task in tasks])
-
-        for key, ft in self.features.items():
-            # index, episode_index, task_index are already processed above, and image and video
-            # are processed separately by storing image path and frame info as meta data
-            if key in ["index", "episode_index", "task_index"] or ft["dtype"] in ["image", "video"]:
-                continue
-            episode_buffer[key] = np.stack(episode_buffer[key])
-
-        self._wait_image_writer()
-        self._save_episode_table(episode_buffer, episode_index)
-        ep_stats = compute_episode_stats(episode_buffer, self.features)
-
-        if len(self.meta.video_keys) > 0:
-            video_paths = self.encode_episode_videos(episode_index)
-            for key in self.meta.video_keys:
-                episode_buffer[key] = video_paths[key]
-
-        # `meta.save_episode` be executed after encoding the videos
-        self.meta.save_episode(episode_index, episode_length, episode_tasks, ep_stats)
-
-        ep_data_index = get_episode_data_index(self.meta.episodes, [episode_index])
-        ep_data_index_np = {k: t.numpy() for k, t in ep_data_index.items()}
-        check_timestamps_sync(
-            episode_buffer["timestamp"],
-            episode_buffer["episode_index"],
-            ep_data_index_np,
-            self.fps,
-            self.tolerance_s,
-        )
-
-        video_files = list(self.root.rglob("*.mp4"))
-        assert len(video_files) == self.num_episodes * len(self.meta.video_keys)
-
-        parquet_files = list(self.root.rglob("*.parquet"))
-        assert len(parquet_files) == self.num_episodes
-
-        # delete images
-        img_dir = self.root / "images"
-        if img_dir.is_dir():
-            shutil.rmtree(self.root / "images")
-
-        if not episode_data:  # Reset the buffer
-            self.episode_buffer = self.create_episode_buffer()
-
-    def _save_episode_table(self, episode_buffer: dict, episode_index: int) -> None:
-        episode_dict = {key: episode_buffer[key] for key in self.hf_features}
-        ep_dataset = datasets.Dataset.from_dict(episode_dict, features=self.hf_features, split="train")
-        ep_dataset = embed_images(ep_dataset)
-        self.hf_dataset = concatenate_datasets([self.hf_dataset, ep_dataset])
-        self.hf_dataset.set_transform(hf_transform_to_torch)
-        ep_data_path = self.root / self.meta.get_data_file_path(ep_index=episode_index)
-        ep_data_path.parent.mkdir(parents=True, exist_ok=True)
-        ep_dataset.to_parquet(ep_data_path)
-
-    def clear_episode_buffer(self) -> None:
-        episode_index = self.episode_buffer["episode_index"]
-        if self.image_writer is not None:
-            for cam_key in self.meta.camera_keys:
-                img_dir = self._get_image_file_path(
-                    episode_index=episode_index, image_key=cam_key, frame_index=0
-                ).parent
-                if img_dir.is_dir():
-                    shutil.rmtree(img_dir)
-
-        # Reset the buffer
-        self.episode_buffer = self.create_episode_buffer()
-
-    def start_image_writer(self, num_processes: int = 0, num_threads: int = 4) -> None:
-        if isinstance(self.image_writer, AsyncImageWriter):
-            logging.warning(
-                "You are starting a new AsyncImageWriter that is replacing an already existing one in the dataset."
-            )
-
-        self.image_writer = AsyncImageWriter(
-            num_processes=num_processes,
-            num_threads=num_threads,
-        )
-
-    def stop_image_writer(self) -> None:
-        """
-        Whenever wrapping this dataset inside a parallelized DataLoader, this needs to be called first to
-        remove the image_writer in order for the LeRobotDataset object to be picklable and parallelized.
-        """
-        if self.image_writer is not None:
-            self.image_writer.stop()
-            self.image_writer = None
-
-    def _wait_image_writer(self) -> None:
-        """Wait for asynchronous image writer to finish."""
-        if self.image_writer is not None:
-            self.image_writer.wait_until_done()
-
-    def encode_videos(self) -> None:
-        """
-        Use ffmpeg to convert frames stored as png into mp4 videos.
-        Note: `encode_video_frames` is a blocking call. Making it asynchronous shouldn't speedup encoding,
-        since video encoding with ffmpeg is already using multithreading.
-        """
-        for ep_idx in range(self.meta.total_episodes):
-            self.encode_episode_videos(ep_idx)
-
-    def encode_episode_videos(self, episode_index: int) -> dict:
-        """
-        Use ffmpeg to convert frames stored as png into mp4 videos.
-        Note: `encode_video_frames` is a blocking call. Making it asynchronous shouldn't speedup encoding,
-        since video encoding with ffmpeg is already using multithreading.
-        """
-        video_paths = {}
-        for key in self.meta.video_keys:
-            video_path = self.root / self.meta.get_video_file_path(episode_index, key)
-            video_paths[key] = str(video_path)
-            if video_path.is_file():
-                # Skip if video is already encoded. Could be the case when resuming data recording.
-                continue
-            img_dir = self._get_image_file_path(
-                episode_index=episode_index, image_key=key, frame_index=0
-            ).parent
-            encode_video_frames(img_dir, video_path, self.fps, overwrite=True)
-
-        return video_paths
-
-    @classmethod
-    def create(
-        cls,
-        repo_id: str,
-        fps: int,
-        features: dict,
-        root: str | Path | None = None,
-        robot_type: str | None = None,
-        use_videos: bool = True,
-        tolerance_s: float = 1e-4,
-        image_writer_processes: int = 0,
-        image_writer_threads: int = 0,
-        video_backend: str | None = None,
-    ) -> "LeRobotDataset":
-        """Create a LeRobot Dataset from scratch in order to record data."""
-        obj = cls.__new__(cls)
-        obj.meta = LeRobotDatasetMetadata.create(
-            repo_id=repo_id,
-            fps=fps,
-            robot_type=robot_type,
-            features=features,
-            root=root,
-            use_videos=use_videos,
-        )
-        obj.repo_id = obj.meta.repo_id
-        obj.root = obj.meta.root
-        obj.local_files_only = obj.meta.local_files_only
-        obj.revision = None
-        obj.tolerance_s = tolerance_s
-        obj.image_writer = None
-
-        if image_writer_processes or image_writer_threads:
-            obj.start_image_writer(image_writer_processes, image_writer_threads)
-
-        # TODO(aliberts, rcadene, alexander-soare): Merge this with OnlineBuffer/DataBuffer
-        obj.episode_buffer = obj.create_episode_buffer()
-
-        obj.episodes = None
-        obj.hf_dataset = obj.create_hf_dataset()
-        obj.image_transforms = None
-        obj.delta_timestamps = None
-        obj.delta_indices = None
-        obj.episode_data_index = None
-        obj.video_backend = video_backend if video_backend is not None else get_safe_default_codec()
-        return obj
-
-def reshape_features_to_max_dim(features: dict, reshape_dim: int = -1, keys_to_max_dim: dict = {}) -> dict:
-    """Reshape features to have a maximum dimension of `max_dim`."""
-    reshaped_features = {}
-    for key in features:
-        if key in keys_to_max_dim and keys_to_max_dim[key] is not None:
-            reshaped_features[key] = features[key]
-            shape = list(features[key]["shape"])
-            if any([k in key for k in [OBS_IMAGE, OBS_IMAGE_2, OBS_IMAGE_3]]):  # Assume square images
-                shape[-3] = keys_to_max_dim[key]
-                shape[-2] = keys_to_max_dim[key]
-            else:
-                shape[reshape_dim] = keys_to_max_dim[key]
-            reshaped_features[key]["shape"] = tuple(shape)
-        else:
-            reshaped_features[key] = features[key]
-    return reshaped_features
-
-def str_to_torch_dtype(dtype_str):
-    """Convert a dtype string to a torch dtype."""
-    mapping = {
-        "float32": torch.float32,
-        "int64": torch.int64,
-        "int16": torch.int16,
-        "bool": torch.bool,
-        "video": torch.float32,  # Assuming video is stored as uint8 images
-    }
-    return mapping.get(dtype_str, torch.float32)  
-def create_padded_features(item: dict, features: dict = {}):
-    for key, ft in features.items():
-        if any([k in key for k in ["cam", "effort", "absolute"]]): # FIXME(mshukor): temporary hack
-            continue
-        shape = ft["shape"]
-        if len(shape) == 3:  # images to torch format (C, H, W)
-            shape = (shape[2], shape[0], shape[1])
-        if len(shape) == 1 and shape[0] == 1:  # ft with shape are actually tensor(ele)
-            shape = []
-        if key not in item:
-            dtype = str_to_torch_dtype(ft["dtype"])
-            item[key] = torch.zeros(shape, dtype=dtype)
-            item[f"{key}_padding_mask"] = torch.tensor(0, dtype=torch.int64)
-            if "image" in key: # FIXME(mshukor): support other observations
-                item[f"{key}_is_pad"] = torch.BoolTensor([False])
-        else:
-            item[f"{key}_padding_mask"] = torch.tensor(1, dtype=torch.int64)
-    return item
-
-class MultiLeRobotDataset(torch.utils.data.Dataset):
-    """A dataset consisting of multiple underlying `LeRobotDataset`s.
-
-    The underlying `LeRobotDataset`s are effectively concatenated, and this class adopts much of the API
-    structure of `LeRobotDataset`.
-    """
-
-    def __init__(
-        self,
-        repo_ids: list[str],
-        root: str | Path | None = None,
-        episodes: dict | None = None,
-        image_transforms: Callable | None = None,
-        delta_timestamps: dict[list[float]] | None = None,
-        tolerances_s: dict | None = None,
-        download_videos: bool = True,
-        local_files_only: bool = False,
-        video_backend: str | None = None,
-        sampling_weights: list[float] | None = None,
-        feature_keys_mapping: dict[str, dict[str, str]] | None = None,
-        max_action_dim: int = None,
-        max_state_dim: int = None,
-        max_num_images: int = None,
-        max_image_dim: int = None,
-        train_on_all_features: bool = False,
-        training_features: list | None = None,
-        discard_first_n_frames: int = 0,
-        min_fps: int = 1,
-        max_fps: int = 100,
-        discard_first_idle_frames: bool = False,
-        motion_threshold: float = 0.05,
-        motion_window_size: int = 10,
-        motion_buffer: int = 3,
-    ):
-        super().__init__()
-        self.repo_ids = repo_ids
-        self.root = Path(root) if root else HF_LEROBOT_HOME
-        self.tolerances_s = tolerances_s if tolerances_s else dict.fromkeys(repo_ids, 0.0001)
-        # Construct the underlying datasets passing everything but `transform` and `delta_timestamps` which
-        # are handled by this class.
-        self._datasets = [
-            LeRobotDataset(
-                repo_id,
-                root=self.root / repo_id,
-                episodes=episodes[repo_id] if episodes else None,
-                image_transforms=image_transforms,
-                delta_timestamps=delta_timestamps,
-                tolerance_s=self.tolerances_s[repo_id],
-                download_videos=download_videos,
-                video_backend=video_backend,
-            )
-            for repo_id in repo_ids
-        ]
-
-        # Disable any data keys that are not common across all of the datasets. Note: we may relax this
-        # restriction in future iterations of this class. For now, this is necessary at least for being able
-        # to use PyTorch's default DataLoader collate function.
-        self.disabled_features = set()
-        intersection_features = set(self._datasets[0].features)
-        for ds in self._datasets:
-            intersection_features.intersection_update(ds.features)
-        if len(intersection_features) == 0:
-            raise RuntimeError(
-                "Multiple datasets were provided but they had no keys common to all of them. "
-                "The multi-dataset functionality currently only keeps common keys."
-            )
-        for repo_id, ds in zip(self.repo_ids, self._datasets, strict=True):
-            extra_keys = set(ds.features).difference(intersection_features)
-            logging.warning(
-                f"keys {extra_keys} of {repo_id} were disabled as they are not contained in all the "
-                "other datasets."
-            )
-            self.disabled_features.update(extra_keys)
-
-        self.image_transforms = image_transforms
-        self.delta_timestamps = self.delta_timestamps = delta_timestamps.get(
-            repo_id, None
-        )
-        # TODO(rcadene, aliberts): We should not perform this aggregation for datasets
-        # with multiple robots of different ranges. Instead we should have one normalization
-        # per robot.
-        for ds in _datasets:
-            ds.meta.info["robot_type"] = ROBOT_TYPE_KEYS_MAPPING.get(ds.repo_id, ds.meta.info["robot_type"])
-            ds.robot_type = ds.meta.info["robot_type"]
-        #self.stats = aggregate_stats([dataset.meta.stats for dataset in self._datasets])
-        _datasets = keep_datasets_with_valid_fps(_datasets, min_fps=min_fps, max_fps=max_fps)
-        self._datasets, datasets_maks = keep_datasets_with_the_same_features_per_robot_type(_datasets)
-        self.sampling_weights = [self.sampling_weights[i] for i in range(len(_datasets)) if datasets_maks[i]]
-        self.repo_ids = [repo_ids[i] for i in range(len(_datasets)) if datasets_maks[i]]
-        self.datasets_repo_ids = [datasets_repo_ids[i] for i in range(len(_datasets)) if datasets_maks[i]]
-        # Compute cumulative sizes for fast indexing
-        self.cumulative_sizes = np.array(
-            [0] + list(torch.cumsum(torch.tensor([len(d) for d in self._datasets]), dim=0))
-        )
-        self.sampling_weights = np.array(self.sampling_weights, dtype=np.float32)
-        self.stats = aggregate_stats_per_robot_type(self._datasets)
-        self.meta = copy.deepcopy(self._datasets[0].meta)  # FIXME(mshukor): aggregate meta from all datasets
-        self.meta.info = {
-            repo_id: ds.meta.info for repo_id, ds in zip(self.repo_ids, self._datasets, strict=False)
-        }
-        # self.meta.info["features"] = self._datasets[0].meta.info["features"] # Assume all datasets have the same features
-        # FIXME(mshukor): pad based on types in case we have more than one state?
-        self.keys_to_max_dim = {
-            ACTION: max_action_dim,
-            OBS_ENV: max_state_dim,
-            OBS_ROBOT: max_state_dim,
-            OBS_IMAGE: max_image_dim,
-            OBS_IMAGE_2: max_image_dim,
-            OBS_IMAGE_3: max_image_dim,
-        }
-        # self.meta.info["features"] = reshape_features_to_max_dim(self._datasets[0].meta.info["features"], reshape_dim=-1, keys_to_max_dim=self.keys_to_max_dim)
-        self.meta.info["features"] = reshape_features_to_max_dim(
-            union_features, reshape_dim=-1, keys_to_max_dim=self.keys_to_max_dim
-        )
-        # reshape stats
-        for robot_type_, stats_ in self.stats.items():
-            for feat_key, feat_stats in stats_.items():
-                if feat_key in [ACTION, OBS_ENV, OBS_ROBOT]:
-                    for k, v in feat_stats.items():
-                        if k in ["min", "mean"]:
-                            pad_value = 0
-                        elif k in ["max", "std"]:
-                            pad_value = 1
-                        else:
-                            continue
-                        self.stats[robot_type_][feat_key][k] = pad_tensor(v, max_size=self.keys_to_max_dim.get(feat_key, -1), pad_dim=-1, pad_value=pad_value)
-
-        self.meta.stats = self.stats
-        # self.meta.info["features"] = aggregate_features(self._datasets)
-        self.meta.tasks = {
-            repo_id: ds.meta.tasks for repo_id, ds in zip(self.repo_ids, self._datasets, strict=False)
-        }
-        self.meta.episodes = {
-            repo_id: ds.meta.episodes for repo_id, ds in zip(self.repo_ids, self._datasets, strict=False)
-        }
-        self.robot_types = [ds.meta.info["robot_type"] for ds in self._datasets]
-
-    @property
-    def repo_id_to_index(self):
-        """Return a mapping from dataset repo_id to a dataset index automatically created by this class.
-
-        This index is incorporated as a data key in the dictionary returned by `__getitem__`.
-        """
-        return {repo_id: i for i, repo_id in enumerate(self.repo_ids)}
-
-    @property
-    def repo_index_to_id(self):
-        """Return the inverse mapping if repo_id_to_index."""
-        return {v: k for k, v in self.repo_id_to_index}
-
-    @property
-    def fps(self) -> int:
-        """Frames per second used during data collection.
-
-        NOTE: Fow now, this relies on a check in __init__ to make sure all sub-datasets have the same info.
-        """
-        return self._datasets[0].meta.info["fps"]
-
-    @property
-    def video(self) -> bool:
-        """Returns True if this dataset loads video frames from mp4 files.
-
-        Returns False if it only loads images from png files.
-
-        NOTE: Fow now, this relies on a check in __init__ to make sure all sub-datasets have the same info.
-        """
-        return self._datasets[0].meta.info.get("video", False)
-
-    @property
-    def features(self) -> datasets.Features:
-        features = {}
-        for dataset in self._datasets:
-            features.update({k: v for k, v in dataset.hf_features.items() if k not in self.disabled_features})
-        return features
-
-    @property
-    def camera_keys(self) -> list[str]:
-        """Keys to access image and video stream from cameras."""
-        keys = []
-        for key, feats in self.features.items():
-            if isinstance(feats, (datasets.Image, VideoFrame)):
-                keys.append(key)
-        return keys
-
-    @property
-    def video_frame_keys(self) -> list[str]:
-        """Keys to access video frames that requires to be decoded into images.
-
-        Note: It is empty if the dataset contains images only,
-        or equal to `self.cameras` if the dataset contains videos only,
-        or can even be a subset of `self.cameras` in a case of a mixed image/video dataset.
-        """
-        video_frame_keys = []
-        for key, feats in self.features.items():
-            if isinstance(feats, VideoFrame):
-                video_frame_keys.append(key)
-        return video_frame_keys
-
-    @property
-    def num_frames(self) -> int:
-        """Number of samples/frames."""
-        return sum(d.num_frames for d in self._datasets)
-
-    @property
-    def num_episodes(self) -> int:
-        """Number of episodes."""
-        return sum(d.num_episodes for d in self._datasets)
-
-    @property
-    def tolerance_s(self) -> float:
-        """Tolerance in seconds used to discard loaded frames when their timestamps
-        are not close enough from the requested frames. It is only used when `delta_timestamps`
-        is provided or when loading video frames from mp4 files.
-        """
-        # 1e-4 to account for possible numerical error
-        return 1 / self.fps - 1e-4
-
-    def __len__(self):
-        return self.num_frames
-
-    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
-        if idx >= len(self):
-            raise IndexError(f"Index {idx} out of bounds.")
-        # Determine which dataset to get an item from based on the index.
-        dataset_idx = np.searchsorted(self.cumulative_sizes, idx, side="right").item() - 1
-        local_idx = (idx - self.cumulative_sizes[dataset_idx]).item()
-
-
-        item = self._datasets[dataset_idx][local_idx]
-        item["dataset_index"] = torch.tensor(dataset_idx)
-        item = create_padded_features(item, self.meta.info["features"])
-        for data_key in self.disabled_features:  # FIXME(mshukor): not in getitem?
-            if data_key in item:
-                del item[data_key]
-
-        return item
-
-    def __repr__(self):
-        return (
-            f"{self.__class__.__name__}(\n"
-            f"  Repository IDs: '{self.repo_ids}',\n"
-            f"  Number of Samples: {self.num_frames},\n"
-            f"  Number of Episodes: {self.num_episodes},\n"
-            f"  Type: {'video (.mp4)' if self.video else 'image (.png)'},\n"
-            f"  Recorded Frames per Second: {self.fps},\n"
-            f"  Camera Keys: {self.camera_keys},\n"
-            f"  Video Frame Keys: {self.video_frame_keys if self.video else 'N/A'},\n"
-            f"  Transformations: {self.image_transforms},\n"
-            f")"
-        )
diff --git a/lerobot/common/model/kinematics.py b/lerobot/common/model/kinematics.py
deleted file mode 100644
index 367b609e1..000000000
--- a/lerobot/common/model/kinematics.py
+++ /dev/null
@@ -1,483 +0,0 @@
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-import numpy as np
-from numpy.typing import NDArray
-from scipy.spatial.transform import Rotation
-
-
-def skew_symmetric(w: NDArray[np.float32]) -> NDArray[np.float32]:
-    """Creates the skew-symmetric matrix from a 3D vector."""
-    return np.array([[0, -w[2], w[1]], [w[2], 0, -w[0]], [-w[1], w[0], 0]])
-
-
-def rodrigues_rotation(w: NDArray[np.float32], theta: float) -> NDArray[np.float32]:
-    """Computes the rotation matrix using Rodrigues' formula."""
-    w_hat = skew_symmetric(w)
-    return np.eye(3) + np.sin(theta) * w_hat + (1 - np.cos(theta)) * w_hat @ w_hat
-
-
-def screw_axis_to_transform(s: NDArray[np.float32], theta: float) -> NDArray[np.float32]:
-    """Converts a screw axis to a 4x4 transformation matrix."""
-    screw_axis_rot = s[:3]
-    screw_axis_trans = s[3:]
-
-    # Pure translation
-    if np.allclose(screw_axis_rot, 0) and np.linalg.norm(screw_axis_trans) == 1:
-        transform = np.eye(4)
-        transform[:3, 3] = screw_axis_trans * theta
-
-    # Rotation (and potentially translation)
-    elif np.linalg.norm(screw_axis_rot) == 1:
-        w_hat = skew_symmetric(screw_axis_rot)
-        rot_mat = np.eye(3) + np.sin(theta) * w_hat + (1 - np.cos(theta)) * w_hat @ w_hat
-        t = (
-            np.eye(3) * theta + (1 - np.cos(theta)) * w_hat + (theta - np.sin(theta)) * w_hat @ w_hat
-        ) @ screw_axis_trans
-        transform = np.eye(4)
-        transform[:3, :3] = rot_mat
-        transform[:3, 3] = t
-    else:
-        raise ValueError("Invalid screw axis parameters")
-    return transform
-
-
-def pose_difference_se3(pose1: NDArray[np.float32], pose2: NDArray[np.float32]) -> NDArray[np.float32]:
-    """
-    Calculates the SE(3) difference between two 4x4 homogeneous transformation matrices.
-    SE(3) (Special Euclidean Group) represents rigid body transformations in 3D space,
-    combining rotation (SO(3)) and translation.
-
-    Each 4x4 matrix has the following structure:
-    [R11 R12 R13 tx]
-    [R21 R22 R23 ty]
-    [R31 R32 R33 tz]
-    [ 0   0   0   1]
-
-    where R is the 3x3 rotation matrix and [tx,ty,tz] is the translation vector.
-
-    Args:
-        pose1: A 4x4 numpy array representing the first pose.
-        pose2: A 4x4 numpy array representing the second pose.
-
-    Returns:
-        A 6D numpy array concatenating translation and rotation differences.
-        First 3 elements are the translational difference (position).
-        Last 3 elements are the rotational difference in axis-angle representation.
-    """
-    rot1 = pose1[:3, :3]
-    rot2 = pose2[:3, :3]
-
-    translation_diff = pose1[:3, 3] - pose2[:3, 3]
-
-    # Calculate rotational difference using scipy's Rotation library
-    rot_diff = Rotation.from_matrix(rot1 @ rot2.T)
-    rotation_diff = rot_diff.as_rotvec()  # Axis-angle representation
-
-    return np.concatenate([translation_diff, rotation_diff])
-
-
-def se3_error(target_pose: NDArray[np.float32], current_pose: NDArray[np.float32]) -> NDArray[np.float32]:
-    pos_error = target_pose[:3, 3] - current_pose[:3, 3]
-
-    rot_target = target_pose[:3, :3]
-    rot_current = current_pose[:3, :3]
-    rot_error_mat = rot_target @ rot_current.T
-    rot_error = Rotation.from_matrix(rot_error_mat).as_rotvec()
-
-    return np.concatenate([pos_error, rot_error])
-
-
-class RobotKinematics:
-    """Robot kinematics class supporting multiple robot models."""
-
-    # Robot measurements dictionary
-    ROBOT_MEASUREMENTS = {
-        "koch": {
-            "gripper": [0.239, -0.001, 0.024],
-            "wrist": [0.209, 0, 0.024],
-            "forearm": [0.108, 0, 0.02],
-            "humerus": [0, 0, 0.036],
-            "shoulder": [0, 0, 0],
-            "base": [0, 0, 0.02],
-        },
-        "moss": {
-            "gripper": [0.246, 0.013, 0.111],
-            "wrist": [0.245, 0.002, 0.064],
-            "forearm": [0.122, 0, 0.064],
-            "humerus": [0.001, 0.001, 0.063],
-            "shoulder": [0, 0, 0],
-            "base": [0, 0, 0.02],
-        },
-        "so_old_calibration": {
-            "gripper": [0.320, 0, 0.050],
-            "wrist": [0.278, 0, 0.050],
-            "forearm": [0.143, 0, 0.044],
-            "humerus": [0.031, 0, 0.072],
-            "shoulder": [0, 0, 0],
-            "base": [0, 0, 0.02],
-        },
-        "so_new_calibration": {
-            "gripper": [0.33, 0.0, 0.285],
-            "wrist": [0.30, 0.0, 0.267],
-            "forearm": [0.25, 0.0, 0.266],
-            "humerus": [0.06, 0.0, 0.264],
-            "shoulder": [0.0, 0.0, 0.238],
-            "base": [0.0, 0.0, 0.12],
-        },
-    }
-
-    def __init__(self, robot_type: str = "so100"):
-        """Initialize kinematics for the specified robot type.
-
-        Args:
-            robot_type: String specifying the robot model ("koch", "so100", or "moss")
-        """
-        if robot_type not in self.ROBOT_MEASUREMENTS:
-            raise ValueError(
-                f"Unknown robot type: {robot_type}. Available types: {list(self.ROBOT_MEASUREMENTS.keys())}"
-            )
-
-        self.robot_type = robot_type
-        self.measurements = self.ROBOT_MEASUREMENTS[robot_type]
-
-        # Initialize all transformation matrices and screw axes
-        self._setup_transforms()
-
-    def _create_translation_matrix(
-        self, x: float = 0.0, y: float = 0.0, z: float = 0.0
-    ) -> NDArray[np.float32]:
-        """Create a 4x4 translation matrix."""
-        return np.array([[1, 0, 0, x], [0, 1, 0, y], [0, 0, 1, z], [0, 0, 0, 1]])
-
-    def _setup_transforms(self):
-        """Setup all transformation matrices and screw axes for the robot."""
-        # Set up rotation matrices (constant across robot types)
-
-        # Gripper orientation
-        self.gripper_X0 = np.array(
-            [
-                [1, 0, 0, 0],
-                [0, 0, 1, 0],
-                [0, -1, 0, 0],
-                [0, 0, 0, 1],
-            ],
-            dtype=np.float32,
-        )
-
-        # Wrist orientation
-        self.wrist_X0 = np.array(
-            [
-                [0, -1, 0, 0],
-                [1, 0, 0, 0],
-                [0, 0, 1, 0],
-                [0, 0, 0, 1],
-            ],
-            dtype=np.float32,
-        )
-
-        # Base orientation
-        self.base_X0 = np.array(
-            [
-                [0, 0, 1, 0],
-                [1, 0, 0, 0],
-                [0, 1, 0, 0],
-                [0, 0, 0, 1],
-            ],
-            dtype=np.float32,
-        )
-
-        # Gripper
-        # Screw axis of gripper frame wrt base frame
-        self.S_BG = np.array(
-            [
-                1,
-                0,
-                0,
-                0,
-                self.measurements["gripper"][2],
-                -self.measurements["gripper"][1],
-            ],
-            dtype=np.float32,
-        )
-
-        # Gripper origin to centroid transform
-        self.X_GoGc = self._create_translation_matrix(x=0.07)
-
-        # Gripper origin to tip transform
-        self.X_GoGt = self._create_translation_matrix(x=0.12)
-
-        # 0-position gripper frame pose wrt base
-        self.X_BoGo = self._create_translation_matrix(
-            x=self.measurements["gripper"][0],
-            y=self.measurements["gripper"][1],
-            z=self.measurements["gripper"][2],
-        )
-
-        # Wrist
-        # Screw axis of wrist frame wrt base frame
-        self.S_BR = np.array(
-            [0, 1, 0, -self.measurements["wrist"][2], 0, self.measurements["wrist"][0]], dtype=np.float32
-        )
-
-        # 0-position origin to centroid transform
-        self.X_RoRc = self._create_translation_matrix(x=0.0035, y=-0.002)
-
-        # 0-position wrist frame pose wrt base
-        self.X_BR = self._create_translation_matrix(
-            x=self.measurements["wrist"][0],
-            y=self.measurements["wrist"][1],
-            z=self.measurements["wrist"][2],
-        )
-
-        # Forearm
-        # Screw axis of forearm frame wrt base frame
-        self.S_BF = np.array(
-            [
-                0,
-                1,
-                0,
-                -self.measurements["forearm"][2],
-                0,
-                self.measurements["forearm"][0],
-            ],
-            dtype=np.float32,
-        )
-
-        # Forearm origin + centroid transform
-        self.X_ForearmFc = self._create_translation_matrix(x=0.036)
-
-        # 0-position forearm frame pose wrt base
-        self.X_BF = self._create_translation_matrix(
-            x=self.measurements["forearm"][0],
-            y=self.measurements["forearm"][1],
-            z=self.measurements["forearm"][2],
-        )
-
-        # Humerus
-        # Screw axis of humerus frame wrt base frame
-        self.S_BH = np.array(
-            [
-                0,
-                -1,
-                0,
-                self.measurements["humerus"][2],
-                0,
-                -self.measurements["humerus"][0],
-            ],
-            dtype=np.float32,
-        )
-
-        # Humerus origin to centroid transform
-        self.X_HoHc = self._create_translation_matrix(x=0.0475)
-
-        # 0-position humerus frame pose wrt base
-        self.X_BH = self._create_translation_matrix(
-            x=self.measurements["humerus"][0],
-            y=self.measurements["humerus"][1],
-            z=self.measurements["humerus"][2],
-        )
-
-        # Shoulder
-        # Screw axis of shoulder frame wrt Base frame
-        self.S_BS = np.array([0, 0, -1, 0, 0, 0], dtype=np.float32)
-
-        # Shoulder origin to centroid transform
-        self.X_SoSc = self._create_translation_matrix(x=-0.017, z=0.0235)
-
-        # 0-position shoulder frame pose wrt base
-        self.X_BS = self._create_translation_matrix(
-            x=self.measurements["shoulder"][0],
-            y=self.measurements["shoulder"][1],
-            z=self.measurements["shoulder"][2],
-        )
-
-        # Base
-        # Base origin to centroid transform
-        self.X_BoBc = self._create_translation_matrix(y=0.015)
-
-        # World to base transform
-        self.X_WoBo = self._create_translation_matrix(
-            x=self.measurements["base"][0],
-            y=self.measurements["base"][1],
-            z=self.measurements["base"][2],
-        )
-
-        # Pre-compute gripper post-multiplication matrix
-        self._fk_gripper_post = self.X_GoGc @ self.X_BoGo @ self.gripper_X0
-
-    def forward_kinematics(
-        self,
-        robot_pos_deg: NDArray[np.float32],
-        frame: str = "gripper_tip",
-    ) -> NDArray[np.float32]:
-        """Generic forward kinematics.
-
-        Args:
-            robot_pos_deg: Joint positions in degrees. Can be ``None`` when
-                computing the *base* frame as it does not depend on joint
-                angles.
-            frame: Target frame. One of
-                ``{"base", "shoulder", "humerus", "forearm", "wrist", "gripper", "gripper_tip"}``.
-
-        Returns
-        -------
-        NDArray[np.float32]
-            4×4 homogeneous transformation matrix of the requested frame
-            expressed in the world coordinate system.
-        """
-        frame = frame.lower()
-        if frame not in {
-            "base",
-            "shoulder",
-            "humerus",
-            "forearm",
-            "wrist",
-            "gripper",
-            "gripper_tip",
-        }:
-            raise ValueError(
-                f"Unknown frame '{frame}'. Valid options are base, shoulder, humerus, forearm, wrist, gripper, gripper_tip."
-            )
-
-        # Base frame does not rely on joint angles.
-        if frame == "base":
-            return self.X_WoBo @ self.X_BoBc @ self.base_X0
-
-        robot_pos_rad = robot_pos_deg / 180 * np.pi
-
-        # Extract joint angles (note the sign convention for shoulder lift).
-        theta_shoulder_pan = robot_pos_rad[0]
-        theta_shoulder_lift = -robot_pos_rad[1]
-        theta_elbow_flex = robot_pos_rad[2]
-        theta_wrist_flex = robot_pos_rad[3]
-        theta_wrist_roll = robot_pos_rad[4]
-
-        # Start with the world-to-base transform; incrementally add successive links.
-        transformation_matrix = self.X_WoBo @ screw_axis_to_transform(self.S_BS, theta_shoulder_pan)
-        if frame == "shoulder":
-            return transformation_matrix @ self.X_SoSc @ self.X_BS
-
-        transformation_matrix = transformation_matrix @ screw_axis_to_transform(
-            self.S_BH, theta_shoulder_lift
-        )
-        if frame == "humerus":
-            return transformation_matrix @ self.X_HoHc @ self.X_BH
-
-        transformation_matrix = transformation_matrix @ screw_axis_to_transform(self.S_BF, theta_elbow_flex)
-        if frame == "forearm":
-            return transformation_matrix @ self.X_ForearmFc @ self.X_BF
-
-        transformation_matrix = transformation_matrix @ screw_axis_to_transform(self.S_BR, theta_wrist_flex)
-        if frame == "wrist":
-            return transformation_matrix @ self.X_RoRc @ self.X_BR @ self.wrist_X0
-
-        transformation_matrix = transformation_matrix @ screw_axis_to_transform(self.S_BG, theta_wrist_roll)
-        if frame == "gripper":
-            return transformation_matrix @ self._fk_gripper_post
-        else:  # frame == "gripper_tip"
-            return transformation_matrix @ self.X_GoGt @ self.X_BoGo @ self.gripper_X0
-
-    def compute_jacobian(
-        self, robot_pos_deg: NDArray[np.float32], frame: str = "gripper_tip"
-    ) -> NDArray[np.float32]:
-        """Finite differences to compute the Jacobian.
-        J(i, j) represents how the ith component of the end-effector's velocity changes wrt a small change
-        in the jth joint's velocity.
-
-        Args:
-            robot_pos_deg: Current joint positions in degrees
-            fk_func: Forward kinematics function to use (defaults to fk_gripper)
-        """
-
-        eps = 1e-8
-        jac = np.zeros(shape=(6, 5))
-        delta = np.zeros(len(robot_pos_deg[:-1]), dtype=np.float64)
-        for el_ix in range(len(robot_pos_deg[:-1])):
-            delta *= 0
-            delta[el_ix] = eps / 2
-            sdot = (
-                pose_difference_se3(
-                    self.forward_kinematics(robot_pos_deg[:-1] + delta, frame),
-                    self.forward_kinematics(robot_pos_deg[:-1] - delta, frame),
-                )
-                / eps
-            )
-            jac[:, el_ix] = sdot
-        return jac
-
-    def compute_positional_jacobian(
-        self, robot_pos_deg: NDArray[np.float32], frame: str = "gripper_tip"
-    ) -> NDArray[np.float32]:
-        """Finite differences to compute the positional Jacobian.
-        J(i, j) represents how the ith component of the end-effector's position changes wrt a small change
-        in the jth joint's velocity.
-
-        Args:
-            robot_pos_deg: Current joint positions in degrees
-            fk_func: Forward kinematics function to use (defaults to fk_gripper)
-        """
-        eps = 1e-8
-        jac = np.zeros(shape=(3, 5))
-        delta = np.zeros(len(robot_pos_deg[:-1]), dtype=np.float64)
-        for el_ix in range(len(robot_pos_deg[:-1])):
-            delta *= 0
-            delta[el_ix] = eps / 2
-            sdot = (
-                self.forward_kinematics(robot_pos_deg[:-1] + delta, frame)[:3, 3]
-                - self.forward_kinematics(robot_pos_deg[:-1] - delta, frame)[:3, 3]
-            ) / eps
-            jac[:, el_ix] = sdot
-        return jac
-
-    def ik(
-        self,
-        current_joint_pos: NDArray[np.float32],
-        desired_ee_pose: NDArray[np.float32],
-        position_only: bool = True,
-        frame: str = "gripper_tip",
-        max_iterations: int = 5,
-        learning_rate: float = 1,
-    ) -> NDArray[np.float32]:
-        """Inverse kinematics using gradient descent.
-
-        Args:
-            current_joint_state: Initial joint positions in degrees
-            desired_ee_pose: Target end-effector pose as a 4x4 transformation matrix
-            position_only: If True, only match end-effector position, not orientation
-            frame: Target frame. One of
-                ``{"base", "shoulder", "humerus", "forearm", "wrist", "gripper", "gripper_tip"}``.
-            max_iterations: Maximum number of iterations to run
-            learning_rate: Learning rate for gradient descent
-
-        Returns:
-            Joint positions in degrees that achieve the desired end-effector pose
-        """
-        # Do gradient descent.
-        current_joint_state = current_joint_pos.copy()
-        for _ in range(max_iterations):
-            current_ee_pose = self.forward_kinematics(current_joint_state, frame)
-            if not position_only:
-                error = se3_error(desired_ee_pose, current_ee_pose)
-                jac = self.compute_jacobian(current_joint_state, frame)
-            else:
-                error = desired_ee_pose[:3, 3] - current_ee_pose[:3, 3]
-                jac = self.compute_positional_jacobian(current_joint_state, frame)
-            delta_angles = np.linalg.pinv(jac) @ error
-            current_joint_state[:-1] += learning_rate * delta_angles
-
-            if np.linalg.norm(error) < 5e-3:
-                return current_joint_state
-        return current_joint_state
diff --git a/lerobot/common/transport/services_pb2.py b/lerobot/common/transport/services_pb2.py
deleted file mode 100644
index 727beb60d..000000000
--- a/lerobot/common/transport/services_pb2.py
+++ /dev/null
@@ -1,45 +0,0 @@
-# Generated by the protocol buffer compiler.  DO NOT EDIT!
-# NO CHECKED-IN PROTOBUF GENCODE
-# source: lerobot/common/transport/services.proto
-# Protobuf Python Version: 5.29.0
-"""Generated protocol buffer code."""
-from google.protobuf import descriptor as _descriptor
-from google.protobuf import descriptor_pool as _descriptor_pool
-from google.protobuf import runtime_version as _runtime_version
-from google.protobuf import symbol_database as _symbol_database
-from google.protobuf.internal import builder as _builder
-_runtime_version.ValidateProtobufRuntimeVersion(
-    _runtime_version.Domain.PUBLIC,
-    5,
-    29,
-    0,
-    '',
-    'lerobot/common/transport/services.proto'
-)
-# @@protoc_insertion_point(imports)
-
-_sym_db = _symbol_database.Default()
-
-
-
-
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\'lerobot/common/transport/services.proto\x12\ttransport\"L\n\nTransition\x12\x30\n\x0etransfer_state\x18\x01 \x01(\x0e\x32\x18.transport.TransferState\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"L\n\nParameters\x12\x30\n\x0etransfer_state\x18\x01 \x01(\x0e\x32\x18.transport.TransferState\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"T\n\x12InteractionMessage\x12\x30\n\x0etransfer_state\x18\x01 \x01(\x0e\x32\x18.transport.TransferState\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"\x07\n\x05\x45mpty*`\n\rTransferState\x12\x14\n\x10TRANSFER_UNKNOWN\x10\x00\x12\x12\n\x0eTRANSFER_BEGIN\x10\x01\x12\x13\n\x0fTRANSFER_MIDDLE\x10\x02\x12\x10\n\x0cTRANSFER_END\x10\x03\x32\x81\x02\n\x0eLearnerService\x12=\n\x10StreamParameters\x12\x10.transport.Empty\x1a\x15.transport.Parameters0\x01\x12<\n\x0fSendTransitions\x12\x15.transport.Transition\x1a\x10.transport.Empty(\x01\x12\x45\n\x10SendInteractions\x12\x1d.transport.InteractionMessage\x1a\x10.transport.Empty(\x01\x12+\n\x05Ready\x12\x10.transport.Empty\x1a\x10.transport.Emptyb\x06proto3')
-
-_globals = globals()
-_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
-_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'lerobot.common.transport.services_pb2', _globals)
-if not _descriptor._USE_C_DESCRIPTORS:
-  DESCRIPTOR._loaded_options = None
-  _globals['_TRANSFERSTATE']._serialized_start=305
-  _globals['_TRANSFERSTATE']._serialized_end=401
-  _globals['_TRANSITION']._serialized_start=54
-  _globals['_TRANSITION']._serialized_end=130
-  _globals['_PARAMETERS']._serialized_start=132
-  _globals['_PARAMETERS']._serialized_end=208
-  _globals['_INTERACTIONMESSAGE']._serialized_start=210
-  _globals['_INTERACTIONMESSAGE']._serialized_end=294
-  _globals['_EMPTY']._serialized_start=296
-  _globals['_EMPTY']._serialized_end=303
-  _globals['_LEARNERSERVICE']._serialized_start=404
-  _globals['_LEARNERSERVICE']._serialized_end=661
-# @@protoc_insertion_point(module_scope)
diff --git a/lerobot/scripts/push_pretrained.py b/lerobot/scripts/push_pretrained.py
deleted file mode 100644
index e3c683f96..000000000
--- a/lerobot/scripts/push_pretrained.py
+++ /dev/null
@@ -1,71 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""
-Once you have trained a policy with our training script (lerobot/scripts/train.py), use this script to push it
-to the hub.
-
-Example:
-
-```bash
-python lerobot/scripts/push_pretrained.py \
-    --pretrained_path=outputs/train/act_aloha_sim_transfer_cube_human/checkpoints/last/pretrained_model \
-    --repo_id=lerobot/act_aloha_sim_transfer_cube_human
-```
-"""
-
-from dataclasses import dataclass
-from pathlib import Path
-
-import draccus
-from huggingface_hub import HfApi
-
-
-@dataclass
-class PushPreTrainedConfig:
-    pretrained_path: Path
-    repo_id: str
-    branch: str | None = None
-    private: bool = False
-    exist_ok: bool = False
-
-
-@draccus.wrap()
-def main(cfg: PushPreTrainedConfig):
-    hub_api = HfApi()
-    hub_api.create_repo(
-        repo_id=cfg.repo_id,
-        private=cfg.private,
-        repo_type="model",
-        exist_ok=cfg.exist_ok,
-    )
-    if cfg.branch:
-        hub_api.create_branch(
-            repo_id=cfg.repo_id,
-            branch=cfg.branch,
-            repo_type="model",
-            exist_ok=cfg.exist_ok,
-        )
-
-    hub_api.upload_folder(
-        repo_id=cfg.repo_id,
-        folder_path=cfg.pretrained_path,
-        repo_type="model",
-        revision=cfg.branch,
-    )
-
-
-if __name__ == "__main__":
-    main()
diff --git a/pyproject.toml b/pyproject.toml
index 5bff0fca6..9fc84d903 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -68,7 +68,6 @@ dependencies = [
     "pyserial>=3.5",
     "pyzmq>=26.2.1",
     "rerun-sdk>=0.21.0",
-    "scipy>=1.14.0",
     "termcolor>=2.4.0",
     "torch>=2.2.1",
     "torchcodec>=0.2.1; sys_platform != 'win32' and (sys_platform != 'linux' or (platform_machine != 'aarch64' and platform_machine != 'arm64' and platform_machine != 'armv7l')) and (sys_platform != 'darwin' or platform_machine != 'x86_64')",
@@ -87,6 +86,7 @@ dora = [
 dynamixel = ["dynamixel-sdk>=3.7.31"]
 feetech = ["feetech-servo-sdk>=1.0.0"]
 gamepad = ["pygame>=2.5.1", "hidapi>=0.14.0"]
+kinematics = ["placo>=0.9.6"]
 intelrealsense = [
     "pyrealsense2>=2.55.1.6486 ; sys_platform != 'darwin'",
     "pyrealsense2-macosx>=2.54 ; sys_platform == 'darwin'",
@@ -100,13 +100,16 @@ stretch = [
     "pyrealsense2>=2.55.1.6486 ; sys_platform != 'darwin'"
 ]
 test = ["pytest>=8.1.0", "pytest-timeout>=2.4.0", "pytest-cov>=5.0.0", "pyserial>=3.5", "mock-serial>=0.0.1 ; sys_platform != 'win32'"]
-hilserl = ["transformers>=4.50.3", "gym-hil>=0.1.8", "protobuf>=5.29.3", "grpcio==1.71.0"]
+hilserl = ["transformers>=4.50.3", "gym-hil>=0.1.9", "protobuf>=5.29.3", "grpcio==1.71.0", "placo>=0.9.6"]
 umi = ["imagecodecs>=2024.1.1"]
 video_benchmark = ["scikit-image>=0.23.2", "pandas>=2.2.2"]
 xarm = ["gym-xarm>=0.1.1 ; python_version < '4.0'"]
 
 [tool.poetry]
 requires-poetry = ">=2.1"
+packages = [
+    { include = "lerobot", from = "src" }
+]
 
 [tool.ruff]
 line-length = 110
@@ -123,10 +126,10 @@ select = ["E4", "E7", "E9", "F", "I", "N", "B", "C4", "SIM"]
 exclude_dirs = [
     "tests",
     "benchmarks",
-    "lerobot/common/datasets/push_dataset_to_hub",
-    "lerobot/common/datasets/v2/convert_dataset_v1_to_v2",
-    "lerobot/common/policies/pi0/conversion_scripts",
-    "lerobot/scripts/push_dataset_to_hub.py",
+    "src/lerobot/datasets/push_dataset_to_hub",
+    "src/lerobot/datasets/v2/convert_dataset_v1_to_v2",
+    "src/lerobot/policies/pi0/conversion_scripts",
+    "src/lerobot/scripts/push_dataset_to_hub.py",
 ]
 skips = ["B101", "B311", "B404", "B603"]
 
diff --git a/lerobot/__init__.py b/src/lerobot/__init__.py
similarity index 96%
rename from lerobot/__init__.py
rename to src/lerobot/__init__.py
index 11114da0a..38d4e8644 100644
--- a/lerobot/__init__.py
+++ b/src/lerobot/__init__.py
@@ -167,10 +167,10 @@ available_datasets = sorted(
     set(itertools.chain(*available_datasets_per_env.values(), available_real_world_datasets))
 )
 
-# lists all available policies from `lerobot/common/policies`
+# lists all available policies from `lerobot/policies`
 available_policies = ["act", "diffusion", "tdmpc", "vqbet"]
 
-# lists all available robots from `lerobot/common/robot_devices/robots`
+# lists all available robots from `lerobot/robot_devices/robots`
 available_robots = [
     "koch",
     "koch_bimanual",
@@ -179,13 +179,13 @@ available_robots = [
     "so101",
 ]
 
-# lists all available cameras from `lerobot/common/robot_devices/cameras`
+# lists all available cameras from `lerobot/robot_devices/cameras`
 available_cameras = [
     "opencv",
     "intelrealsense",
 ]
 
-# lists all available motors from `lerobot/common/robot_devices/motors`
+# lists all available motors from `lerobot/robot_devices/motors`
 available_motors = [
     "dynamixel",
     "feetech",
diff --git a/lerobot/__version__.py b/src/lerobot/__version__.py
similarity index 100%
rename from lerobot/__version__.py
rename to src/lerobot/__version__.py
diff --git a/lerobot/calibrate.py b/src/lerobot/calibrate.py
similarity index 84%
rename from lerobot/calibrate.py
rename to src/lerobot/calibrate.py
index 6780577ff..37a9d5bdf 100644
--- a/lerobot/calibrate.py
+++ b/src/lerobot/calibrate.py
@@ -31,9 +31,9 @@ from pprint import pformat
 
 import draccus
 
-from lerobot.common.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
-from lerobot.common.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
-from lerobot.common.robots import (  # noqa: F401
+from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
+from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
+from lerobot.robots import (  # noqa: F401
     Robot,
     RobotConfig,
     koch_follower,
@@ -42,7 +42,7 @@ from lerobot.common.robots import (  # noqa: F401
     so100_follower,
     so101_follower,
 )
-from lerobot.common.teleoperators import (  # noqa: F401
+from lerobot.teleoperators import (  # noqa: F401
     Teleoperator,
     TeleoperatorConfig,
     koch_leader,
@@ -50,7 +50,7 @@ from lerobot.common.teleoperators import (  # noqa: F401
     so100_leader,
     so101_leader,
 )
-from lerobot.common.utils.utils import init_logging
+from lerobot.utils.utils import init_logging
 
 
 @dataclass
diff --git a/lerobot/common/cameras/__init__.py b/src/lerobot/cameras/__init__.py
similarity index 100%
rename from lerobot/common/cameras/__init__.py
rename to src/lerobot/cameras/__init__.py
diff --git a/lerobot/common/cameras/camera.py b/src/lerobot/cameras/camera.py
similarity index 100%
rename from lerobot/common/cameras/camera.py
rename to src/lerobot/cameras/camera.py
diff --git a/lerobot/common/cameras/configs.py b/src/lerobot/cameras/configs.py
similarity index 100%
rename from lerobot/common/cameras/configs.py
rename to src/lerobot/cameras/configs.py
diff --git a/lerobot/common/cameras/opencv/__init__.py b/src/lerobot/cameras/opencv/__init__.py
similarity index 100%
rename from lerobot/common/cameras/opencv/__init__.py
rename to src/lerobot/cameras/opencv/__init__.py
diff --git a/lerobot/common/cameras/opencv/camera_opencv.py b/src/lerobot/cameras/opencv/camera_opencv.py
similarity index 98%
rename from lerobot/common/cameras/opencv/camera_opencv.py
rename to src/lerobot/cameras/opencv/camera_opencv.py
index 3e9370fc4..fd99922a4 100644
--- a/lerobot/common/cameras/opencv/camera_opencv.py
+++ b/src/lerobot/cameras/opencv/camera_opencv.py
@@ -27,7 +27,7 @@ from typing import Any, Dict, List
 import cv2
 import numpy as np
 
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 from ..camera import Camera
 from ..utils import get_cv2_backend, get_cv2_rotation
@@ -64,8 +64,8 @@ class OpenCVCamera(Camera):
 
     Example:
         ```python
-        from lerobot.common.cameras.opencv import OpenCVCamera
-        from lerobot.common.cameras.configuration_opencv import OpenCVCameraConfig, ColorMode, Cv2Rotation
+        from lerobot.cameras.opencv import OpenCVCamera
+        from lerobot.cameras.configuration_opencv import OpenCVCameraConfig, ColorMode, Cv2Rotation
 
         # Basic usage with camera index 0
         config = OpenCVCameraConfig(index_or_path=0)
diff --git a/lerobot/common/cameras/opencv/configuration_opencv.py b/src/lerobot/cameras/opencv/configuration_opencv.py
similarity index 100%
rename from lerobot/common/cameras/opencv/configuration_opencv.py
rename to src/lerobot/cameras/opencv/configuration_opencv.py
diff --git a/lerobot/common/cameras/realsense/__init__.py b/src/lerobot/cameras/realsense/__init__.py
similarity index 100%
rename from lerobot/common/cameras/realsense/__init__.py
rename to src/lerobot/cameras/realsense/__init__.py
diff --git a/lerobot/common/cameras/realsense/camera_realsense.py b/src/lerobot/cameras/realsense/camera_realsense.py
similarity index 98%
rename from lerobot/common/cameras/realsense/camera_realsense.py
rename to src/lerobot/cameras/realsense/camera_realsense.py
index 2bcbee75c..96531b694 100644
--- a/lerobot/common/cameras/realsense/camera_realsense.py
+++ b/src/lerobot/cameras/realsense/camera_realsense.py
@@ -29,7 +29,7 @@ try:
 except Exception as e:
     logging.info(f"Could not import realsense: {e}")
 
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 from ..camera import Camera
 from ..configs import ColorMode
@@ -63,8 +63,8 @@ class RealSenseCamera(Camera):
 
     Example:
         ```python
-        from lerobot.common.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
-        from lerobot.common.cameras import ColorMode, Cv2Rotation
+        from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
+        from lerobot.cameras import ColorMode, Cv2Rotation
 
         # Basic usage with serial number
         config = RealSenseCameraConfig(serial_number_or_name="0123456789") # Replace with actual SN
diff --git a/lerobot/common/cameras/realsense/configuration_realsense.py b/src/lerobot/cameras/realsense/configuration_realsense.py
similarity index 100%
rename from lerobot/common/cameras/realsense/configuration_realsense.py
rename to src/lerobot/cameras/realsense/configuration_realsense.py
diff --git a/lerobot/common/cameras/utils.py b/src/lerobot/cameras/utils.py
similarity index 100%
rename from lerobot/common/cameras/utils.py
rename to src/lerobot/cameras/utils.py
diff --git a/lerobot/configs/datasets.py b/src/lerobot/configs/datasets.py
similarity index 100%
rename from lerobot/configs/datasets.py
rename to src/lerobot/configs/datasets.py
diff --git a/lerobot/configs/default.py b/src/lerobot/configs/default.py
similarity index 94%
rename from lerobot/configs/default.py
rename to src/lerobot/configs/default.py
index ce72466a8..53cfe58e7 100644
--- a/lerobot/configs/default.py
+++ b/src/lerobot/configs/default.py
@@ -16,11 +16,11 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.common import (
+from lerobot import (
     policies,  # noqa: F401
 )
-from lerobot.common.datasets.transforms import ImageTransformsConfig
-from lerobot.common.datasets.video_utils import get_safe_default_codec
+from lerobot.datasets.transforms import ImageTransformsConfig
+from lerobot.datasets.video_utils import get_safe_default_codec
 
 
 @dataclass
diff --git a/lerobot/configs/eval.py b/src/lerobot/configs/eval.py
similarity index 97%
rename from lerobot/configs/eval.py
rename to src/lerobot/configs/eval.py
index 16b352913..cfe48cf87 100644
--- a/lerobot/configs/eval.py
+++ b/src/lerobot/configs/eval.py
@@ -17,7 +17,7 @@ import logging
 from dataclasses import dataclass, field
 from pathlib import Path
 
-from lerobot.common import envs, policies  # noqa: F401
+from lerobot import envs, policies  # noqa: F401
 from lerobot.configs import parser
 from lerobot.configs.default import EvalConfig
 from lerobot.configs.policies import PreTrainedConfig
diff --git a/lerobot/configs/parser.py b/src/lerobot/configs/parser.py
similarity index 99%
rename from lerobot/configs/parser.py
rename to src/lerobot/configs/parser.py
index f69b5a7fa..1da7ad83f 100644
--- a/lerobot/configs/parser.py
+++ b/src/lerobot/configs/parser.py
@@ -22,7 +22,7 @@ from typing import Sequence
 
 import draccus
 
-from lerobot.common.utils.utils import has_method
+from lerobot.utils.utils import has_method
 
 PATH_KEY = "path"
 PLUGIN_DISCOVERY_SUFFIX = "discover_packages_path"
diff --git a/lerobot/configs/policies.py b/src/lerobot/configs/policies.py
similarity index 92%
rename from lerobot/configs/policies.py
rename to src/lerobot/configs/policies.py
index 1302db1fa..36e6ea2e5 100644
--- a/lerobot/configs/policies.py
+++ b/src/lerobot/configs/policies.py
@@ -23,11 +23,11 @@ from huggingface_hub import hf_hub_download
 from huggingface_hub.constants import CONFIG_NAME
 from huggingface_hub.errors import HfHubHTTPError
 
-from lerobot.common.optim.optimizers import OptimizerConfig
-from lerobot.common.optim.schedulers import LRSchedulerConfig
-from lerobot.common.utils.hub import HubMixin
-from lerobot.common.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim.optimizers import OptimizerConfig
+from lerobot.optim.schedulers import LRSchedulerConfig
+from lerobot.utils.hub import HubMixin
+from lerobot.utils.utils import auto_select_torch_device, is_amp_available, is_torch_device_available
 
 # Generic variable that is either PreTrainedConfig or a subclass thereof
 T = TypeVar("T", bound="PreTrainedConfig")
@@ -60,6 +60,16 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
     # automatic gradient scaling is used.
     use_amp: bool = False
 
+    push_to_hub: bool = True
+    repo_id: str | None = None
+
+    # Upload on private repository on the Hugging Face hub.
+    private: bool | None = None
+    # Add tags to your policy on the hub.
+    tags: list[str] | None = None
+    # Add tags to your policy on the hub.
+    license: str | None = None
+
     def __post_init__(self):
         self.pretrained_path = None
         if not self.device or not is_torch_device_available(self.device):
diff --git a/lerobot/configs/train.py b/src/lerobot/configs/train.py
similarity index 95%
rename from lerobot/configs/train.py
rename to src/lerobot/configs/train.py
index 96a460bdf..c088a5fa1 100644
--- a/lerobot/configs/train.py
+++ b/src/lerobot/configs/train.py
@@ -21,13 +21,13 @@ import draccus
 from huggingface_hub import hf_hub_download
 from huggingface_hub.errors import HfHubHTTPError
 
-from lerobot.common import envs
-from lerobot.common.optim import OptimizerConfig
-from lerobot.common.optim.schedulers import LRSchedulerConfig
-from lerobot.common.utils.hub import HubMixin
+from lerobot import envs
 from lerobot.configs import parser
 from lerobot.configs.default import DatasetConfig, EvalConfig, WandBConfig
 from lerobot.configs.policies import PreTrainedConfig
+from lerobot.optim import OptimizerConfig
+from lerobot.optim.schedulers import LRSchedulerConfig
+from lerobot.utils.hub import HubMixin
 
 TRAIN_CONFIG_NAME = "train_config.json"
 
@@ -116,6 +116,11 @@ class TrainPipelineConfig(HubMixin):
             self.optimizer = self.policy.get_optimizer_preset()
             self.scheduler = self.policy.get_scheduler_preset()
 
+        if self.policy.push_to_hub and not self.policy.repo_id:
+            raise ValueError(
+                "'policy.repo_id' argument missing. Please specify it to push the model to the hub."
+            )
+
     @classmethod
     def __get_path_fields__(cls) -> list[str]:
         """This enables the parser to load config from the policy using `--policy.path=local/dir`"""
diff --git a/lerobot/configs/types.py b/src/lerobot/configs/types.py
similarity index 100%
rename from lerobot/configs/types.py
rename to src/lerobot/configs/types.py
diff --git a/lerobot/common/constants.py b/src/lerobot/constants.py
similarity index 98%
rename from lerobot/common/constants.py
rename to src/lerobot/constants.py
index 83d736392..e3720bdf0 100644
--- a/lerobot/common/constants.py
+++ b/src/lerobot/constants.py
@@ -29,7 +29,7 @@ REWARD = "next.reward"
 
 ROBOTS = "robots"
 TASK = "task"
-ROBOT = "robot_type"
+ROBOT_TYPE = "robot_type"
 TELEOPERATORS = "teleoperators"
 
 # files & directories
diff --git a/lerobot/common/datasets/backward_compatibility.py b/src/lerobot/datasets/backward_compatibility.py
similarity index 95%
rename from lerobot/common/datasets/backward_compatibility.py
rename to src/lerobot/datasets/backward_compatibility.py
index cf8e31c4f..fae485058 100644
--- a/lerobot/common/datasets/backward_compatibility.py
+++ b/src/lerobot/datasets/backward_compatibility.py
@@ -20,7 +20,7 @@ The dataset you requested ({repo_id}) is in {version} format.
 We introduced a new format since v2.0 which is not backward compatible with v1.x.
 Please, use our conversion script. Modify the following command with your own task description:
 ```
-python lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py \\
+python -m lerobot.datasets.v2.convert_dataset_v1_to_v2 \\
     --repo-id {repo_id} \\
     --single-task "TASK DESCRIPTION."  # <---- /!\\ Replace TASK DESCRIPTION /!\\
 ```
@@ -40,7 +40,7 @@ The dataset you requested ({repo_id}) is in {version} format.
 While current version of LeRobot is backward-compatible with it, the version of your dataset still uses global
 stats instead of per-episode stats. Update your dataset stats to the new format using this command:
 ```
-python lerobot/common/datasets/v21/convert_dataset_v20_to_v21.py --repo-id={repo_id}
+python -m lerobot.datasets.v21.convert_dataset_v20_to_v21 --repo-id={repo_id}
 ```
 
 If you encounter a problem, contact LeRobot maintainers on [Discord](https://discord.com/invite/s3KuuzsPFb)
diff --git a/lerobot/common/datasets/card_template.md b/src/lerobot/datasets/card_template.md
similarity index 100%
rename from lerobot/common/datasets/card_template.md
rename to src/lerobot/datasets/card_template.md
diff --git a/lerobot/common/datasets/collators.py b/src/lerobot/datasets/collators.py
similarity index 100%
rename from lerobot/common/datasets/collators.py
rename to src/lerobot/datasets/collators.py
diff --git a/lerobot/common/datasets/compute_stats.py b/src/lerobot/datasets/compute_stats.py
similarity index 99%
rename from lerobot/common/datasets/compute_stats.py
rename to src/lerobot/datasets/compute_stats.py
index 1149ec83e..bfe7b18b4 100644
--- a/lerobot/common/datasets/compute_stats.py
+++ b/src/lerobot/datasets/compute_stats.py
@@ -15,7 +15,7 @@
 # limitations under the License.
 import numpy as np
 
-from lerobot.common.datasets.utils import load_image_as_numpy
+from lerobot.datasets.utils import load_image_as_numpy
 
 
 def estimate_num_samples(
diff --git a/lerobot/common/datasets/factory.py b/src/lerobot/datasets/factory.py
similarity index 97%
rename from lerobot/common/datasets/factory.py
rename to src/lerobot/datasets/factory.py
index 88d3f767f..e06650bc9 100644
--- a/lerobot/common/datasets/factory.py
+++ b/src/lerobot/datasets/factory.py
@@ -18,14 +18,14 @@ from pprint import pformat
 
 import torch
 
-from lerobot.common.datasets.lerobot_dataset import (
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.train import TrainPipelineConfig
+from lerobot.datasets.lerobot_dataset import (
     LeRobotDataset,
     LeRobotDatasetMetadata,
     MultiLeRobotDataset,
 )
-from lerobot.common.datasets.transforms import ImageTransforms
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.configs.train import TrainPipelineConfig
+from lerobot.datasets.transforms import ImageTransforms
 
 IMAGENET_STATS = {
     "mean": [[[0.485]], [[0.456]], [[0.406]]],  # (c,1,1)
diff --git a/lerobot/common/datasets/image_writer.py b/src/lerobot/datasets/image_writer.py
similarity index 100%
rename from lerobot/common/datasets/image_writer.py
rename to src/lerobot/datasets/image_writer.py
diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/lerobot/common/datasets/online_buffer.py b/src/lerobot/datasets/online_buffer.py
similarity index 99%
rename from lerobot/common/datasets/online_buffer.py
rename to src/lerobot/datasets/online_buffer.py
index d907e4687..79f48f49d 100644
--- a/lerobot/common/datasets/online_buffer.py
+++ b/src/lerobot/datasets/online_buffer.py
@@ -28,7 +28,7 @@ from typing import Any
 import numpy as np
 import torch
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
 
 
 def _make_memmap_safe(**kwargs) -> np.memmap:
diff --git a/lerobot/common/datasets/push_dataset_to_hub/utils.py b/src/lerobot/datasets/push_dataset_to_hub/utils.py
similarity index 98%
rename from lerobot/common/datasets/push_dataset_to_hub/utils.py
rename to src/lerobot/datasets/push_dataset_to_hub/utils.py
index ebcf87f77..6aca7b03b 100644
--- a/lerobot/common/datasets/push_dataset_to_hub/utils.py
+++ b/src/lerobot/datasets/push_dataset_to_hub/utils.py
@@ -23,7 +23,7 @@ import numpy
 import PIL
 import torch
 
-from lerobot.common.datasets.video_utils import encode_video_frames
+from lerobot.datasets.video_utils import encode_video_frames
 
 
 def concatenate_episodes(ep_dicts):
diff --git a/lerobot/common/datasets/sampler.py b/src/lerobot/datasets/sampler.py
similarity index 100%
rename from lerobot/common/datasets/sampler.py
rename to src/lerobot/datasets/sampler.py
diff --git a/lerobot/common/datasets/transforms.py b/src/lerobot/datasets/transforms.py
similarity index 100%
rename from lerobot/common/datasets/transforms.py
rename to src/lerobot/datasets/transforms.py
diff --git a/lerobot/common/datasets/utils.py b/src/lerobot/datasets/utils.py
similarity index 98%
rename from lerobot/common/datasets/utils.py
rename to src/lerobot/datasets/utils.py
index 927a784b0..53a91ced0 100644
--- a/lerobot/common/datasets/utils.py
+++ b/src/lerobot/datasets/utils.py
@@ -35,14 +35,14 @@ from huggingface_hub.errors import RevisionNotFoundError
 from PIL import Image as PILImage
 from torchvision import transforms
 
-from lerobot.common.datasets.backward_compatibility import (
+from lerobot.configs.types import DictLike, FeatureType, PolicyFeature
+from lerobot.datasets.backward_compatibility import (
     V21_MESSAGE,
     BackwardCompatibilityError,
     ForwardCompatibilityError,
 )
-from lerobot.common.robots import Robot
-from lerobot.common.utils.utils import is_valid_numpy_dtype_string
-from lerobot.configs.types import DictLike, FeatureType, PolicyFeature
+from lerobot.robots import Robot
+from lerobot.utils.utils import is_valid_numpy_dtype_string
 
 DEFAULT_CHUNK_SIZE = 1000  # Max number of episodes per chunk
 
@@ -664,7 +664,7 @@ def create_lerobot_dataset_card(
     **kwargs,
 ) -> DatasetCard:
     """
-    Keyword arguments will be used to replace values in ./lerobot/common/datasets/card_template.md.
+    Keyword arguments will be used to replace values in src/lerobot/datasets/card_template.md.
     Note: If specified, license must be one of https://huggingface.co/docs/hub/repositories-licenses.
     """
     card_tags = ["LeRobot"]
@@ -687,7 +687,7 @@ def create_lerobot_dataset_card(
         ],
     )
 
-    card_template = (importlib.resources.files("lerobot.common.datasets") / "card_template.md").read_text()
+    card_template = (importlib.resources.files("lerobot.datasets") / "card_template.md").read_text()
 
     return DatasetCard.from_template(
         card_data=card_data,
diff --git a/lerobot/common/datasets/v2/batch_convert_dataset_v1_to_v2.py b/src/lerobot/datasets/v2/batch_convert_dataset_v1_to_v2.py
similarity index 99%
rename from lerobot/common/datasets/v2/batch_convert_dataset_v1_to_v2.py
rename to src/lerobot/datasets/v2/batch_convert_dataset_v1_to_v2.py
index 9b21cf7ca..fa99c725e 100644
--- a/lerobot/common/datasets/v2/batch_convert_dataset_v1_to_v2.py
+++ b/src/lerobot/datasets/v2/batch_convert_dataset_v1_to_v2.py
@@ -26,8 +26,8 @@ from pathlib import Path
 from textwrap import dedent
 
 from lerobot import available_datasets
-from lerobot.common.datasets.v2.convert_dataset_v1_to_v2 import convert_dataset
-from lerobot.common.robots.aloha.configuration_aloha import AlohaRobotConfig
+from lerobot.datasets.v2.convert_dataset_v1_to_v2 import convert_dataset
+from lerobot.robots.aloha.configuration_aloha import AlohaRobotConfig
 
 LOCAL_DIR = Path("data/")
 
diff --git a/lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py b/src/lerobot/datasets/v2/convert_dataset_v1_to_v2.py
similarity index 97%
rename from lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py
rename to src/lerobot/datasets/v2/convert_dataset_v1_to_v2.py
index 136a7a684..cddfc4c18 100644
--- a/lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py
+++ b/src/lerobot/datasets/v2/convert_dataset_v1_to_v2.py
@@ -38,7 +38,7 @@ If your dataset contains a single task, you can simply provide it directly via t
 Examples:
 
 ```bash
-python lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py \
+python -m lerobot.datasets.v2.convert_dataset_v1_to_v2 \
     --repo-id lerobot/aloha_sim_insertion_human_image \
     --single-task "Insert the peg into the socket." \
     --robot-config lerobot/configs/robot/aloha.yaml \
@@ -46,7 +46,7 @@ python lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py \
 ```
 
 ```bash
-python lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py \
+python -m lerobot.datasets.v2.convert_dataset_v1_to_v2 \
     --repo-id aliberts/koch_tutorial \
     --single-task "Pick the Lego block and drop it in the box on the right." \
     --robot-config lerobot/configs/robot/koch.yaml \
@@ -63,7 +63,7 @@ If your dataset is a multi-task dataset, you have two options to provide the tas
     Example:
 
     ```bash
-    python lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py \
+    python -m lerobot.datasets.v2.convert_dataset_v1_to_v2 \
         --repo-id lerobot/stanford_kuka_multimodal_dataset \
         --tasks-col "language_instruction" \
         --local-dir data
@@ -92,7 +92,7 @@ parquet file, and you must provide this column's name with the '--tasks-col' arg
 Example:
 
 ```bash
-python lerobot/common/datasets/v2/convert_dataset_v1_to_v2.py \
+python -m lerobot.datasets.v2.convert_dataset_v1_to_v2 \
     --repo-id lerobot/stanford_kuka_multimodal_dataset \
     --tasks-col "language_instruction" \
     --local-dir data
@@ -119,7 +119,7 @@ from huggingface_hub import HfApi
 from huggingface_hub.errors import EntryNotFoundError, HfHubHTTPError
 from safetensors.torch import load_file
 
-from lerobot.common.datasets.utils import (
+from lerobot.datasets.utils import (
     DEFAULT_CHUNK_SIZE,
     DEFAULT_PARQUET_PATH,
     DEFAULT_VIDEO_PATH,
@@ -136,12 +136,12 @@ from lerobot.common.datasets.utils import (
     write_json,
     write_jsonlines,
 )
-from lerobot.common.datasets.video_utils import (
+from lerobot.datasets.video_utils import (
     VideoFrame,  # noqa: F401
     get_image_pixel_channels,
     get_video_info,
 )
-from lerobot.common.robots import RobotConfig
+from lerobot.robots import RobotConfig
 
 V16 = "v1.6"
 V20 = "v2.0"
@@ -602,19 +602,19 @@ def make_robot_config(robot_type: str, **kwargs) -> RobotConfig:
         raise NotImplementedError  # TODO
 
     elif robot_type == "koch_follower":
-        from lerobot.common.robots.koch_follower import KochFollowerConfig
+        from lerobot.robots.koch_follower import KochFollowerConfig
 
         return KochFollowerConfig(**kwargs)
     elif robot_type == "so100_follower":
-        from lerobot.common.robots.so100_follower import SO100FollowerConfig
+        from lerobot.robots.so100_follower import SO100FollowerConfig
 
         return SO100FollowerConfig(**kwargs)
     elif robot_type == "stretch":
-        from lerobot.common.robots.stretch3 import Stretch3RobotConfig
+        from lerobot.robots.stretch3 import Stretch3RobotConfig
 
         return Stretch3RobotConfig(**kwargs)
     elif robot_type == "lekiwi":
-        from lerobot.common.robots.lekiwi import LeKiwiConfig
+        from lerobot.robots.lekiwi import LeKiwiConfig
 
         return LeKiwiConfig(**kwargs)
     else:
diff --git a/lerobot/common/datasets/v21/_remove_language_instruction.py b/src/lerobot/datasets/v21/_remove_language_instruction.py
similarity index 92%
rename from lerobot/common/datasets/v21/_remove_language_instruction.py
rename to src/lerobot/datasets/v21/_remove_language_instruction.py
index 643ddd3f2..1f1cb1855 100644
--- a/lerobot/common/datasets/v21/_remove_language_instruction.py
+++ b/src/lerobot/datasets/v21/_remove_language_instruction.py
@@ -20,9 +20,9 @@ from datasets import get_dataset_config_info
 from huggingface_hub import HfApi
 
 from lerobot import available_datasets
-from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata
-from lerobot.common.datasets.utils import INFO_PATH, write_info
-from lerobot.common.datasets.v21.convert_dataset_v20_to_v21 import V20, SuppressWarnings
+from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
+from lerobot.datasets.utils import INFO_PATH, write_info
+from lerobot.datasets.v21.convert_dataset_v20_to_v21 import V20, SuppressWarnings
 
 LOCAL_DIR = Path("data/")
 
diff --git a/lerobot/common/datasets/v21/batch_convert_dataset_v20_to_v21.py b/src/lerobot/datasets/v21/batch_convert_dataset_v20_to_v21.py
similarity index 95%
rename from lerobot/common/datasets/v21/batch_convert_dataset_v20_to_v21.py
rename to src/lerobot/datasets/v21/batch_convert_dataset_v20_to_v21.py
index cc9272a83..b4f1c36c4 100644
--- a/lerobot/common/datasets/v21/batch_convert_dataset_v20_to_v21.py
+++ b/src/lerobot/datasets/v21/batch_convert_dataset_v20_to_v21.py
@@ -24,7 +24,7 @@ from pathlib import Path
 from huggingface_hub import HfApi
 
 from lerobot import available_datasets
-from lerobot.common.datasets.v21.convert_dataset_v20_to_v21 import V21, convert_dataset
+from lerobot.datasets.v21.convert_dataset_v20_to_v21 import V21, convert_dataset
 
 LOCAL_DIR = Path("data/")
 
diff --git a/lerobot/common/datasets/v21/convert_dataset_v20_to_v21.py b/src/lerobot/datasets/v21/convert_dataset_v20_to_v21.py
similarity index 90%
rename from lerobot/common/datasets/v21/convert_dataset_v20_to_v21.py
rename to src/lerobot/datasets/v21/convert_dataset_v20_to_v21.py
index 176d16d0f..4ebc1086a 100644
--- a/lerobot/common/datasets/v21/convert_dataset_v20_to_v21.py
+++ b/src/lerobot/datasets/v21/convert_dataset_v20_to_v21.py
@@ -25,7 +25,7 @@ This script will help you convert any LeRobot dataset already pushed to the hub
 Usage:
 
 ```bash
-python lerobot/common/datasets/v21/convert_dataset_v20_to_v21.py \
+python -m lerobot.datasets.v21.convert_dataset_v20_to_v21 \
     --repo-id=aliberts/koch_tutorial
 ```
 
@@ -36,9 +36,9 @@ import logging
 
 from huggingface_hub import HfApi
 
-from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
-from lerobot.common.datasets.utils import EPISODES_STATS_PATH, STATS_PATH, load_stats, write_info
-from lerobot.common.datasets.v21.convert_stats import check_aggregate_stats, convert_stats
+from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset
+from lerobot.datasets.utils import EPISODES_STATS_PATH, STATS_PATH, load_stats, write_info
+from lerobot.datasets.v21.convert_stats import check_aggregate_stats, convert_stats
 
 V20 = "v2.0"
 V21 = "v2.1"
diff --git a/lerobot/common/datasets/v21/convert_stats.py b/src/lerobot/datasets/v21/convert_stats.py
similarity index 94%
rename from lerobot/common/datasets/v21/convert_stats.py
rename to src/lerobot/datasets/v21/convert_stats.py
index 4a20b4276..462781c15 100644
--- a/lerobot/common/datasets/v21/convert_stats.py
+++ b/src/lerobot/datasets/v21/convert_stats.py
@@ -17,9 +17,9 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 import numpy as np
 from tqdm import tqdm
 
-from lerobot.common.datasets.compute_stats import aggregate_stats, get_feature_stats, sample_indices
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.datasets.utils import write_episode_stats
+from lerobot.datasets.compute_stats import aggregate_stats, get_feature_stats, sample_indices
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import write_episode_stats
 
 
 def sample_episode_video_frames(dataset: LeRobotDataset, episode_index: int, ft_key: str) -> np.ndarray:
diff --git a/lerobot/common/datasets/video_utils.py b/src/lerobot/datasets/video_utils.py
similarity index 100%
rename from lerobot/common/datasets/video_utils.py
rename to src/lerobot/datasets/video_utils.py
diff --git a/lerobot/common/envs/__init__.py b/src/lerobot/envs/__init__.py
similarity index 100%
rename from lerobot/common/envs/__init__.py
rename to src/lerobot/envs/__init__.py
diff --git a/lerobot/common/envs/configs.py b/src/lerobot/envs/configs.py
similarity index 97%
rename from lerobot/common/envs/configs.py
rename to src/lerobot/envs/configs.py
index ea081e9fb..de969d618 100644
--- a/lerobot/common/envs/configs.py
+++ b/src/lerobot/envs/configs.py
@@ -18,10 +18,10 @@ from typing import Any, Optional
 
 import draccus
 
-from lerobot.common.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_IMAGES, OBS_STATE
-from lerobot.common.robots import RobotConfig
-from lerobot.common.teleoperators.config import TeleoperatorConfig
 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
 
 
 @dataclass
diff --git a/lerobot/common/envs/factory.py b/src/lerobot/envs/factory.py
similarity index 96%
rename from lerobot/common/envs/factory.py
rename to src/lerobot/envs/factory.py
index 4f5d59c69..dc6d96d61 100644
--- a/lerobot/common/envs/factory.py
+++ b/src/lerobot/envs/factory.py
@@ -17,7 +17,7 @@ import importlib
 
 import gymnasium as gym
 
-from lerobot.common.envs.configs import AlohaEnv, EnvConfig, HILEnvConfig, PushtEnv, XarmEnv
+from lerobot.envs.configs import AlohaEnv, EnvConfig, HILEnvConfig, PushtEnv, XarmEnv
 
 
 def make_env_config(env_type: str, **kwargs) -> EnvConfig:
diff --git a/lerobot/common/envs/utils.py b/src/lerobot/envs/utils.py
similarity index 97%
rename from lerobot/common/envs/utils.py
rename to src/lerobot/envs/utils.py
index 66d6e5f93..00676a011 100644
--- a/lerobot/common/envs/utils.py
+++ b/src/lerobot/envs/utils.py
@@ -22,9 +22,9 @@ import numpy as np
 import torch
 from torch import Tensor
 
-from lerobot.common.envs.configs import EnvConfig
-from lerobot.common.utils.utils import get_channel_first_image_shape
 from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.envs.configs import EnvConfig
+from lerobot.utils.utils import get_channel_first_image_shape
 
 
 def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Tensor]:
diff --git a/lerobot/common/errors.py b/src/lerobot/errors.py
similarity index 100%
rename from lerobot/common/errors.py
rename to src/lerobot/errors.py
diff --git a/lerobot/find_cameras.py b/src/lerobot/find_cameras.py
similarity index 96%
rename from lerobot/find_cameras.py
rename to src/lerobot/find_cameras.py
index 34f4865b1..aff2f8c19 100644
--- a/lerobot/find_cameras.py
+++ b/src/lerobot/find_cameras.py
@@ -37,11 +37,11 @@ from typing import Any, Dict, List
 import numpy as np
 from PIL import Image
 
-from lerobot.common.cameras.configs import ColorMode
-from lerobot.common.cameras.opencv.camera_opencv import OpenCVCamera
-from lerobot.common.cameras.opencv.configuration_opencv import OpenCVCameraConfig
-from lerobot.common.cameras.realsense.camera_realsense import RealSenseCamera
-from lerobot.common.cameras.realsense.configuration_realsense import RealSenseCameraConfig
+from lerobot.cameras.configs import ColorMode
+from lerobot.cameras.opencv.camera_opencv import OpenCVCamera
+from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
+from lerobot.cameras.realsense.camera_realsense import RealSenseCamera
+from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig
 
 logger = logging.getLogger(__name__)
 
diff --git a/lerobot/find_port.py b/src/lerobot/find_port.py
similarity index 100%
rename from lerobot/find_port.py
rename to src/lerobot/find_port.py
diff --git a/src/lerobot/model/kinematics.py b/src/lerobot/model/kinematics.py
new file mode 100644
index 000000000..f059b9790
--- /dev/null
+++ b/src/lerobot/model/kinematics.py
@@ -0,0 +1,128 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+
+class RobotKinematics:
+    """Robot kinematics using placo library for forward and inverse kinematics."""
+
+    def __init__(
+        self,
+        urdf_path: str,
+        target_frame_name: str = "gripper_frame_link",
+        joint_names: list[str] = None,
+    ):
+        """
+        Initialize placo-based kinematics solver.
+
+        Args:
+            urdf_path: Path to the robot URDF file
+            target_frame_name: Name of the end-effector frame in the URDF
+            joint_names: List of joint names to use for the kinematics solver
+        """
+        try:
+            import placo
+        except ImportError as e:
+            raise ImportError(
+                "placo is required for RobotKinematics. "
+                "Please install the optional dependencies of `kinematics` in the package."
+            ) from e
+
+        self.robot = placo.RobotWrapper(urdf_path)
+        self.solver = placo.KinematicsSolver(self.robot)
+        self.solver.mask_fbase(True)  # Fix the base
+
+        self.target_frame_name = target_frame_name
+
+        # Set joint names
+        self.joint_names = list(self.robot.joint_names()) if joint_names is None else joint_names
+
+        # Initialize frame task for IK
+        self.tip_frame = self.solver.add_frame_task(self.target_frame_name, np.eye(4))
+
+    def forward_kinematics(self, joint_pos_deg):
+        """
+        Compute forward kinematics for given joint configuration given the target frame name in the constructor.
+
+        Args:
+            joint_pos_deg: Joint positions in degrees (numpy array)
+
+        Returns:
+            4x4 transformation matrix of the end-effector pose
+        """
+
+        # Convert degrees to radians
+        joint_pos_rad = np.deg2rad(joint_pos_deg[: len(self.joint_names)])
+
+        # Update joint positions in placo robot
+        for i, joint_name in enumerate(self.joint_names):
+            self.robot.set_joint(joint_name, joint_pos_rad[i])
+
+        # Update kinematics
+        self.robot.update_kinematics()
+
+        # Get the transformation matrix
+        return self.robot.get_T_world_frame(self.target_frame_name)
+
+    def inverse_kinematics(
+        self, current_joint_pos, desired_ee_pose, position_weight=1.0, orientation_weight=0.01
+    ):
+        """
+        Compute inverse kinematics using placo solver.
+
+        Args:
+            current_joint_pos: Current joint positions in degrees (used as initial guess)
+            desired_ee_pose: Target end-effector pose as a 4x4 transformation matrix
+            position_weight: Weight for position constraint in IK
+            orientation_weight: Weight for orientation constraint in IK, set to 0.0 to only constrain position
+
+        Returns:
+            Joint positions in degrees that achieve the desired end-effector pose
+        """
+
+        # Convert current joint positions to radians for initial guess
+        current_joint_rad = np.deg2rad(current_joint_pos[: len(self.joint_names)])
+
+        # Set current joint positions as initial guess
+        for i, joint_name in enumerate(self.joint_names):
+            self.robot.set_joint(joint_name, current_joint_rad[i])
+
+        # Update the target pose for the frame task
+        self.tip_frame.T_world_frame = desired_ee_pose
+
+        # Configure the task based on position_only flag
+        self.tip_frame.configure(self.target_frame_name, "soft", position_weight, orientation_weight)
+
+        # Solve IK
+        self.solver.solve(True)
+        self.robot.update_kinematics()
+
+        # Extract joint positions
+        joint_pos_rad = []
+        for joint_name in self.joint_names:
+            joint = self.robot.get_joint(joint_name)
+            joint_pos_rad.append(joint)
+
+        # Convert back to degrees
+        joint_pos_deg = np.rad2deg(joint_pos_rad)
+
+        # Preserve gripper position if present in current_joint_pos
+        if len(current_joint_pos) > len(self.joint_names):
+            result = np.zeros_like(current_joint_pos)
+            result[: len(self.joint_names)] = joint_pos_deg
+            result[len(self.joint_names) :] = current_joint_pos[len(self.joint_names) :]
+            return result
+        else:
+            return joint_pos_deg
diff --git a/lerobot/common/motors/__init__.py b/src/lerobot/motors/__init__.py
similarity index 100%
rename from lerobot/common/motors/__init__.py
rename to src/lerobot/motors/__init__.py
diff --git a/lerobot/common/motors/dynamixel/__init__.py b/src/lerobot/motors/dynamixel/__init__.py
similarity index 100%
rename from lerobot/common/motors/dynamixel/__init__.py
rename to src/lerobot/motors/dynamixel/__init__.py
diff --git a/lerobot/common/motors/dynamixel/dynamixel.py b/src/lerobot/motors/dynamixel/dynamixel.py
similarity index 99%
rename from lerobot/common/motors/dynamixel/dynamixel.py
rename to src/lerobot/motors/dynamixel/dynamixel.py
index 9f0db901d..d4f41643c 100644
--- a/lerobot/common/motors/dynamixel/dynamixel.py
+++ b/src/lerobot/motors/dynamixel/dynamixel.py
@@ -22,7 +22,7 @@ import logging
 from copy import deepcopy
 from enum import Enum
 
-from lerobot.common.utils.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 (
diff --git a/lerobot/common/motors/dynamixel/tables.py b/src/lerobot/motors/dynamixel/tables.py
similarity index 100%
rename from lerobot/common/motors/dynamixel/tables.py
rename to src/lerobot/motors/dynamixel/tables.py
diff --git a/lerobot/common/motors/feetech/__init__.py b/src/lerobot/motors/feetech/__init__.py
similarity index 100%
rename from lerobot/common/motors/feetech/__init__.py
rename to src/lerobot/motors/feetech/__init__.py
diff --git a/lerobot/common/motors/feetech/feetech.py b/src/lerobot/motors/feetech/feetech.py
similarity index 99%
rename from lerobot/common/motors/feetech/feetech.py
rename to src/lerobot/motors/feetech/feetech.py
index 4937fdea7..7edf869a4 100644
--- a/lerobot/common/motors/feetech/feetech.py
+++ b/src/lerobot/motors/feetech/feetech.py
@@ -17,7 +17,7 @@ from copy import deepcopy
 from enum import Enum
 from pprint import pformat
 
-from lerobot.common.utils.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 (
diff --git a/lerobot/common/motors/feetech/tables.py b/src/lerobot/motors/feetech/tables.py
similarity index 100%
rename from lerobot/common/motors/feetech/tables.py
rename to src/lerobot/motors/feetech/tables.py
diff --git a/lerobot/common/motors/motors_bus.py b/src/lerobot/motors/motors_bus.py
similarity index 99%
rename from lerobot/common/motors/motors_bus.py
rename to src/lerobot/motors/motors_bus.py
index 7ac9e6813..7386bfb1c 100644
--- a/lerobot/common/motors/motors_bus.py
+++ b/src/lerobot/motors/motors_bus.py
@@ -32,8 +32,8 @@ import serial
 from deepdiff import DeepDiff
 from tqdm import tqdm
 
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
-from lerobot.common.utils.utils import enter_pressed, move_cursor_up
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.utils.utils import enter_pressed, move_cursor_up
 
 NameOrID: TypeAlias = str | int
 Value: TypeAlias = int | float
@@ -446,7 +446,7 @@ class MotorsBus(abc.ABC):
         except (FileNotFoundError, OSError, serial.SerialException) as e:
             raise ConnectionError(
                 f"\nCould not connect on port '{self.port}'. Make sure you are using the correct port."
-                "\nTry running `python lerobot/find_port.py`\n"
+                "\nTry running `python -m lerobot.find_port`\n"
             ) from e
 
     @abc.abstractmethod
diff --git a/lerobot/common/optim/__init__.py b/src/lerobot/optim/__init__.py
similarity index 100%
rename from lerobot/common/optim/__init__.py
rename to src/lerobot/optim/__init__.py
diff --git a/lerobot/common/optim/factory.py b/src/lerobot/optim/factory.py
similarity index 96%
rename from lerobot/common/optim/factory.py
rename to src/lerobot/optim/factory.py
index 10ff3df73..bab95d0ce 100644
--- a/lerobot/common/optim/factory.py
+++ b/src/lerobot/optim/factory.py
@@ -18,8 +18,8 @@
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 
-from lerobot.common.policies.pretrained import PreTrainedPolicy
 from lerobot.configs.train import TrainPipelineConfig
+from lerobot.policies.pretrained import PreTrainedPolicy
 
 
 def make_optimizer_and_scheduler(
diff --git a/lerobot/common/optim/optimizers.py b/src/lerobot/optim/optimizers.py
similarity index 97%
rename from lerobot/common/optim/optimizers.py
rename to src/lerobot/optim/optimizers.py
index 903434f59..ece4dc157 100644
--- a/lerobot/common/optim/optimizers.py
+++ b/src/lerobot/optim/optimizers.py
@@ -22,12 +22,12 @@ import draccus
 import torch
 from safetensors.torch import load_file, save_file
 
-from lerobot.common.constants import (
+from lerobot.constants import (
     OPTIMIZER_PARAM_GROUPS,
     OPTIMIZER_STATE,
 )
-from lerobot.common.datasets.utils import flatten_dict, unflatten_dict, write_json
-from lerobot.common.utils.io_utils import deserialize_json_into_object
+from lerobot.datasets.utils import flatten_dict, unflatten_dict, write_json
+from lerobot.utils.io_utils import deserialize_json_into_object
 
 
 @dataclass
diff --git a/lerobot/common/optim/schedulers.py b/src/lerobot/optim/schedulers.py
similarity index 96%
rename from lerobot/common/optim/schedulers.py
rename to src/lerobot/optim/schedulers.py
index 7e1583946..d08018175 100644
--- a/lerobot/common/optim/schedulers.py
+++ b/src/lerobot/optim/schedulers.py
@@ -22,9 +22,9 @@ import draccus
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LambdaLR, LRScheduler
 
-from lerobot.common.constants import SCHEDULER_STATE
-from lerobot.common.datasets.utils import write_json
-from lerobot.common.utils.io_utils import deserialize_json_into_object
+from lerobot.constants import SCHEDULER_STATE
+from lerobot.datasets.utils import write_json
+from lerobot.utils.io_utils import deserialize_json_into_object
 
 
 @dataclass
diff --git a/lerobot/common/policies/__init__.py b/src/lerobot/policies/__init__.py
similarity index 100%
rename from lerobot/common/policies/__init__.py
rename to src/lerobot/policies/__init__.py
diff --git a/lerobot/common/policies/act/configuration_act.py b/src/lerobot/policies/act/configuration_act.py
similarity index 99%
rename from lerobot/common/policies/act/configuration_act.py
rename to src/lerobot/policies/act/configuration_act.py
index 7a5819b74..6f6c1c4be 100644
--- a/lerobot/common/policies/act/configuration_act.py
+++ b/src/lerobot/policies/act/configuration_act.py
@@ -15,9 +15,9 @@
 # limitations under the License.
 from dataclasses import dataclass, field
 
-from lerobot.common.optim.optimizers import AdamWConfig
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
+from lerobot.optim.optimizers import AdamWConfig
 
 
 @PreTrainedConfig.register_subclass("act")
diff --git a/lerobot/common/policies/act/modeling_act.py b/src/lerobot/policies/act/modeling_act.py
similarity index 97%
rename from lerobot/common/policies/act/modeling_act.py
rename to src/lerobot/policies/act/modeling_act.py
index e7e74bf38..ed911e9be 100644
--- a/lerobot/common/policies/act/modeling_act.py
+++ b/src/lerobot/policies/act/modeling_act.py
@@ -33,9 +33,10 @@ from torch import Tensor, nn
 from torchvision.models._utils import IntermediateLayerGetter
 from torchvision.ops.misc import FrozenBatchNorm2d
 
-from lerobot.common.policies.act.configuration_act import ACTConfig
-from lerobot.common.policies.normalize import Normalize, Unnormalize
-from lerobot.common.policies.pretrained import PreTrainedPolicy
+from lerobot.constants import ACTION, OBS_IMAGES
+from lerobot.policies.act.configuration_act import ACTConfig
+from lerobot.policies.normalize import Normalize, Unnormalize
+from lerobot.policies.pretrained import PreTrainedPolicy
 
 
 class ACTPolicy(PreTrainedPolicy):
@@ -114,46 +115,49 @@ class ACTPolicy(PreTrainedPolicy):
         environment. It works by managing the actions in a queue and only calling `select_actions` when the
         queue is empty.
         """
-        self.eval()
+        self.eval()  # keeping the policy in eval mode as it could be set to train mode while queue is consumed
 
-        batch = self.normalize_inputs(batch)
-        if self.config.image_features:
-            batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
-            batch["observation.images"] = [batch[key] for key in self.config.image_features]
-
-        # If we are doing temporal ensembling, do online updates where we keep track of the number of actions
-        # we are ensembling over.
         if self.config.temporal_ensemble_coeff is not None:
-            actions = self.model(batch)[0]  # (batch_size, chunk_size, action_dim)
-            actions = self.unnormalize_outputs({"action": actions})["action"]
+            actions = self.predict_action_chunk(batch)
             action = self.temporal_ensembler.update(actions)
             return action
 
         # Action queue logic for n_action_steps > 1. When the action_queue is depleted, populate it by
         # querying the policy.
         if len(self._action_queue) == 0:
-            actions = self.model(batch)[0][:, : self.config.n_action_steps]
-
-            # TODO(rcadene): make _forward return output dictionary?
-            actions = self.unnormalize_outputs({"action": actions})["action"]
+            actions = self.predict_action_chunk(batch)[:, : self.config.n_action_steps]
 
             # `self.model.forward` returns a (batch_size, n_action_steps, action_dim) tensor, but the queue
             # effectively has shape (n_action_steps, batch_size, *), hence the transpose.
             self._action_queue.extend(actions.transpose(0, 1))
         return self._action_queue.popleft()
 
+    @torch.no_grad
+    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+        """Predict a chunk of actions given environment observations."""
+        self.eval()
+
+        batch = self.normalize_inputs(batch)
+        if self.config.image_features:
+            batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
+            batch[OBS_IMAGES] = [batch[key] for key in self.config.image_features]
+
+        actions = self.model(batch)[0]
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
+        return actions
+
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
         """Run the batch through the model and compute the loss for training or validation."""
         batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
-            batch["observation.images"] = [batch[key] for key in self.config.image_features]
+            batch[OBS_IMAGES] = [batch[key] for key in self.config.image_features]
 
         batch = self.normalize_targets(batch)
         actions_hat, (mu_hat, log_sigma_x2_hat) = self.model(batch)
 
         l1_loss = (
-            F.l1_loss(batch["action"], actions_hat, reduction="none") * ~batch["action_is_pad"].unsqueeze(-1)
+            F.l1_loss(batch[ACTION], actions_hat, reduction="none") * ~batch["action_is_pad"].unsqueeze(-1)
         ).mean()
 
         loss_dict = {"l1_loss": l1_loss.item()}
diff --git a/lerobot/common/policies/diffusion/configuration_diffusion.py b/src/lerobot/policies/diffusion/configuration_diffusion.py
similarity index 99%
rename from lerobot/common/policies/diffusion/configuration_diffusion.py
rename to src/lerobot/policies/diffusion/configuration_diffusion.py
index c8841f06b..ce2de7052 100644
--- a/lerobot/common/policies/diffusion/configuration_diffusion.py
+++ b/src/lerobot/policies/diffusion/configuration_diffusion.py
@@ -16,10 +16,10 @@
 # limitations under the License.
 from dataclasses import dataclass, field
 
-from lerobot.common.optim.optimizers import AdamConfig
-from lerobot.common.optim.schedulers import DiffuserSchedulerConfig
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
+from lerobot.optim.optimizers import AdamConfig
+from lerobot.optim.schedulers import DiffuserSchedulerConfig
 
 
 @PreTrainedConfig.register_subclass("diffusion")
diff --git a/lerobot/common/policies/diffusion/modeling_diffusion.py b/src/lerobot/policies/diffusion/modeling_diffusion.py
similarity index 96%
rename from lerobot/common/policies/diffusion/modeling_diffusion.py
rename to src/lerobot/policies/diffusion/modeling_diffusion.py
index 446e2cb6e..af40f7a86 100644
--- a/lerobot/common/policies/diffusion/modeling_diffusion.py
+++ b/src/lerobot/policies/diffusion/modeling_diffusion.py
@@ -33,11 +33,11 @@ from diffusers.schedulers.scheduling_ddim import DDIMScheduler
 from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
 from torch import Tensor, nn
 
-from lerobot.common.constants import OBS_ENV_STATE, OBS_STATE
-from lerobot.common.policies.diffusion.configuration_diffusion import DiffusionConfig
-from lerobot.common.policies.normalize import Normalize, Unnormalize
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.policies.utils import (
+from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
+from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
+from lerobot.policies.normalize import Normalize, Unnormalize
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.utils import (
     get_device_from_parameters,
     get_dtype_from_parameters,
     get_output_shape,
@@ -99,6 +99,18 @@ class DiffusionPolicy(PreTrainedPolicy):
         if self.config.env_state_feature:
             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:
+        """Predict a chunk of actions given environment observations."""
+        # stack n latest observations from the queue
+        batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
+        actions = self.diffusion.generate_actions(batch)
+
+        # TODO(rcadene): make above methods return output dictionary?
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
+
+        return actions
+
     @torch.no_grad
     def select_action(self, batch: dict[str, Tensor]) -> Tensor:
         """Select a single action given environment observations.
@@ -124,23 +136,15 @@ class DiffusionPolicy(PreTrainedPolicy):
         batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
-            batch["observation.images"] = torch.stack(
-                [batch[key] for key in self.config.image_features], dim=-4
-            )
+            batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
         # Note: It's important that this happens after stacking the images into a single key.
         self._queues = populate_queues(self._queues, batch)
 
-        if len(self._queues["action"]) == 0:
-            # stack n latest observations from the queue
-            batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
-            actions = self.diffusion.generate_actions(batch)
+        if len(self._queues[ACTION]) == 0:
+            actions = self.predict_action_chunk(batch)
+            self._queues[ACTION].extend(actions.transpose(0, 1))
 
-            # TODO(rcadene): make above methods return output dictionary?
-            actions = self.unnormalize_outputs({"action": actions})["action"]
-
-            self._queues["action"].extend(actions.transpose(0, 1))
-
-        action = self._queues["action"].popleft()
+        action = self._queues[ACTION].popleft()
         return action
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, None]:
@@ -148,9 +152,7 @@ class DiffusionPolicy(PreTrainedPolicy):
         batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
-            batch["observation.images"] = torch.stack(
-                [batch[key] for key in self.config.image_features], dim=-4
-            )
+            batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
         batch = self.normalize_targets(batch)
         loss = self.diffusion.compute_loss(batch)
         # no output_dict so returning None
diff --git a/lerobot/common/policies/factory.py b/src/lerobot/policies/factory.py
similarity index 74%
rename from lerobot/common/policies/factory.py
rename to src/lerobot/policies/factory.py
index 68bc73ca4..ef56bdb61 100644
--- a/lerobot/common/policies/factory.py
+++ b/src/lerobot/policies/factory.py
@@ -18,67 +18,62 @@ import logging
 
 from torch import nn
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata
-from lerobot.common.datasets.utils import dataset_to_policy_features
-from lerobot.common.envs.configs import EnvConfig
-from lerobot.common.envs.utils import env_to_policy_features
-from lerobot.common.policies.act.configuration_act import ACTConfig
-from lerobot.common.policies.diffusion.configuration_diffusion import DiffusionConfig
-from lerobot.common.policies.pi0.configuration_pi0 import PI0Config
-from lerobot.common.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.policies.sac.configuration_sac import SACConfig
-from lerobot.common.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
-from lerobot.common.policies.smolvla.configuration_smolvla import SmolVLAConfig
-from lerobot.common.policies.smolvla2.configuration_smolvla2 import SmolVLA2Config
-from lerobot.common.policies.tdmpc.configuration_tdmpc import TDMPCConfig
-from lerobot.common.policies.vqbet.configuration_vqbet import VQBeTConfig
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType
+from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
+from lerobot.datasets.utils import dataset_to_policy_features
+from lerobot.envs.configs import EnvConfig
+from lerobot.envs.utils import env_to_policy_features
+from lerobot.policies.act.configuration_act import ACTConfig
+from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
+from lerobot.policies.pi0.configuration_pi0 import PI0Config
+from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.sac.configuration_sac import SACConfig
+from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
+from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
+from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
+from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
 
 
 def get_policy_class(name: str) -> PreTrainedPolicy:
     """Get the policy's class and config class given a name (matching the policy class' `name` attribute)."""
     if name == "tdmpc":
-        from lerobot.common.policies.tdmpc.modeling_tdmpc import TDMPCPolicy
+        from lerobot.policies.tdmpc.modeling_tdmpc import TDMPCPolicy
 
         return TDMPCPolicy
     elif name == "diffusion":
-        from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy
+        from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
 
         return DiffusionPolicy
     elif name == "act":
-        from lerobot.common.policies.act.modeling_act import ACTPolicy
+        from lerobot.policies.act.modeling_act import ACTPolicy
 
         return ACTPolicy
     elif name == "vqbet":
-        from lerobot.common.policies.vqbet.modeling_vqbet import VQBeTPolicy
+        from lerobot.policies.vqbet.modeling_vqbet import VQBeTPolicy
 
         return VQBeTPolicy
     elif name == "pi0":
-        from lerobot.common.policies.pi0.modeling_pi0 import PI0Policy
+        from lerobot.policies.pi0.modeling_pi0 import PI0Policy
 
         return PI0Policy
     elif name == "pi0fast":
-        from lerobot.common.policies.pi0fast.modeling_pi0fast import PI0FASTPolicy
+        from lerobot.policies.pi0fast.modeling_pi0fast import PI0FASTPolicy
 
         return PI0FASTPolicy
     elif name == "sac":
-        from lerobot.common.policies.sac.modeling_sac import SACPolicy
+        from lerobot.policies.sac.modeling_sac import SACPolicy
 
         return SACPolicy
     elif name == "reward_classifier":
-        from lerobot.common.policies.sac.reward_model.modeling_classifier import Classifier
+        from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
 
         return Classifier
     elif name == "smolvla":
-        from lerobot.common.policies.smolvla.modeling_smolvla import SmolVLAPolicy
+        from lerobot.policies.smolvla.modeling_smolvla import SmolVLAPolicy
 
         return SmolVLAPolicy
-    elif name == "smolvla2":
-        from lerobot.common.policies.smolvla2.modeling_smolvla2 import SmolVLA2Policy
-
-        return SmolVLA2Policy
     else:
         raise NotImplementedError(f"Policy with name {name} is not implemented.")
 
@@ -100,8 +95,6 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         return SACConfig(**kwargs)
     elif policy_type == "smolvla":
         return SmolVLAConfig(**kwargs)
-    elif policy_type == "smolvla2":
-        return SmolVLA2Config(**kwargs)
     elif policy_type == "reward_classifier":
         return RewardClassifierConfig(**kwargs)
     else:
diff --git a/lerobot/common/policies/normalize.py b/src/lerobot/policies/normalize.py
similarity index 100%
rename from lerobot/common/policies/normalize.py
rename to src/lerobot/policies/normalize.py
diff --git a/lerobot/common/policies/pi0/configuration_pi0.py b/src/lerobot/policies/pi0/configuration_pi0.py
similarity index 98%
rename from lerobot/common/policies/pi0/configuration_pi0.py
rename to src/lerobot/policies/pi0/configuration_pi0.py
index 8c7cc1305..c9728e418 100644
--- a/lerobot/common/policies/pi0/configuration_pi0.py
+++ b/src/lerobot/policies/pi0/configuration_pi0.py
@@ -14,12 +14,12 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.common.optim.optimizers import AdamWConfig
-from lerobot.common.optim.schedulers import (
-    CosineDecayWithWarmupSchedulerConfig,
-)
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim.optimizers import AdamWConfig
+from lerobot.optim.schedulers import (
+    CosineDecayWithWarmupSchedulerConfig,
+)
 
 
 @PreTrainedConfig.register_subclass("pi0")
diff --git a/lerobot/common/policies/pi0/conversion_scripts/benchmark.py b/src/lerobot/policies/pi0/conversion_scripts/benchmark.py
similarity index 95%
rename from lerobot/common/policies/pi0/conversion_scripts/benchmark.py
rename to src/lerobot/policies/pi0/conversion_scripts/benchmark.py
index cb3c0e9ba..c1a488244 100644
--- a/lerobot/common/policies/pi0/conversion_scripts/benchmark.py
+++ b/src/lerobot/policies/pi0/conversion_scripts/benchmark.py
@@ -14,9 +14,9 @@
 
 import torch
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.policies.factory import make_policy
 from lerobot.configs.policies import PreTrainedConfig
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.policies.factory import make_policy
 
 torch.backends.cudnn.benchmark = True
 
diff --git a/lerobot/common/policies/pi0/conversion_scripts/compare_with_jax.py b/src/lerobot/policies/pi0/conversion_scripts/compare_with_jax.py
similarity index 96%
rename from lerobot/common/policies/pi0/conversion_scripts/compare_with_jax.py
rename to src/lerobot/policies/pi0/conversion_scripts/compare_with_jax.py
index 6bd7c91f7..c0c2e4816 100644
--- a/lerobot/common/policies/pi0/conversion_scripts/compare_with_jax.py
+++ b/src/lerobot/policies/pi0/conversion_scripts/compare_with_jax.py
@@ -18,9 +18,9 @@ from pathlib import Path
 
 import torch
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDatasetMetadata
-from lerobot.common.policies.factory import make_policy
 from lerobot.configs.policies import PreTrainedConfig
+from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata
+from lerobot.policies.factory import make_policy
 
 
 def display(tensor: torch.Tensor):
@@ -97,7 +97,7 @@ def main():
 
     noise = torch.from_numpy(noise).to(device=device, dtype=torch.float32)
 
-    from lerobot.common import policies  # noqa
+    from lerobot import policies  # noqa
 
     cfg = PreTrainedConfig.from_pretrained(ckpt_torch_dir)
     cfg.pretrained_path = ckpt_torch_dir
diff --git a/lerobot/common/policies/pi0/conversion_scripts/conversion_utils.py b/src/lerobot/policies/pi0/conversion_scripts/conversion_utils.py
similarity index 100%
rename from lerobot/common/policies/pi0/conversion_scripts/conversion_utils.py
rename to src/lerobot/policies/pi0/conversion_scripts/conversion_utils.py
diff --git a/lerobot/common/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.py b/src/lerobot/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.py
similarity index 98%
rename from lerobot/common/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.py
rename to src/lerobot/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.py
index 73ff506ff..742c9ab3f 100644
--- a/lerobot/common/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.py
+++ b/src/lerobot/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.py
@@ -33,13 +33,13 @@ python
 
 Converting pi0_base:
 ```python
-python lerobot/common/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.py \
+python -m lerobot.policies.pi0.conversion_scripts.convert_pi0_to_hf_lerobot \
     --checkpoint_dir /home/remi_cadene/.cache/openpi/openpi-assets/checkpoints/pi0_base/params \
     --output_path /home/remi_cadene/.cache/openpi/openpi-assets/checkpoints/pi0_base_pytorch
 ```
 
 ```python
-python lerobot/common/policies/pi0/conversion_scripts/convert_pi0_to_hf_lerobot.py \
+python -m lerobot.policies.pi0.conversion_scripts.convert_pi0_to_hf_lerobot \
     --checkpoint_dir /home/remi_cadene/.cache/openpi/openpi-assets/checkpoints/pi0_aloha_sim/params \
     --output_path /home/remi_cadene/.cache/openpi/openpi-assets/checkpoints/pi0_aloha_sim_pytorch
 ```
@@ -54,12 +54,12 @@ import orbax.checkpoint as ocp
 import torch
 from jax.sharding import SingleDeviceSharding
 
-from lerobot.common.policies.pi0.configuration_pi0 import PI0Config
-from lerobot.common.policies.pi0.conversion_scripts.conversion_utils import (
+from lerobot.policies.pi0.configuration_pi0 import PI0Config
+from lerobot.policies.pi0.conversion_scripts.conversion_utils import (
     get_gemma_config,
     get_paligemma_config,
 )
-from lerobot.common.policies.pi0.modeling_pi0 import PI0Policy
+from lerobot.policies.pi0.modeling_pi0 import PI0Policy
 
 PRECISIONS = {"bfloat16": torch.bfloat16, "float32": torch.float32, "float16": torch.float16}
 
diff --git a/lerobot/common/policies/pi0/flex_attention.py b/src/lerobot/policies/pi0/flex_attention.py
similarity index 100%
rename from lerobot/common/policies/pi0/flex_attention.py
rename to src/lerobot/policies/pi0/flex_attention.py
diff --git a/lerobot/common/policies/pi0/modeling_pi0.py b/src/lerobot/policies/pi0/modeling_pi0.py
similarity index 97%
rename from lerobot/common/policies/pi0/modeling_pi0.py
rename to src/lerobot/policies/pi0/modeling_pi0.py
index 1d8a50559..241509d0b 100644
--- a/lerobot/common/policies/pi0/modeling_pi0.py
+++ b/src/lerobot/policies/pi0/modeling_pi0.py
@@ -29,7 +29,7 @@ pip install -e ".[pi0]"
 
 Example of finetuning the pi0 pretrained model (`pi0_base` in `openpi`):
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
 --policy.path=lerobot/pi0 \
 --dataset.repo_id=danaaubakirova/koch_test
 ```
@@ -37,7 +37,7 @@ python lerobot/scripts/train.py \
 Example of finetuning the pi0 neural network with PaliGemma and expert Gemma
 pretrained with VLM default parameters before pi0 finetuning:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
 --policy.type=pi0 \
 --dataset.repo_id=danaaubakirova/koch_test
 ```
@@ -57,15 +57,15 @@ import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
 from transformers import AutoTokenizer
 
-from lerobot.common.constants import ACTION, OBS_STATE
-from lerobot.common.policies.normalize import Normalize, Unnormalize
-from lerobot.common.policies.pi0.configuration_pi0 import PI0Config
-from lerobot.common.policies.pi0.paligemma_with_expert import (
+from lerobot.constants import ACTION, OBS_STATE
+from lerobot.policies.normalize import Normalize, Unnormalize
+from lerobot.policies.pi0.configuration_pi0 import PI0Config
+from lerobot.policies.pi0.paligemma_with_expert import (
     PaliGemmaWithExpertConfig,
     PaliGemmaWithExpertModel,
 )
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.utils.utils import get_safe_dtype
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.utils.utils import get_safe_dtype
 
 
 def create_sinusoidal_pos_embedding(
@@ -260,6 +260,11 @@ class PI0Policy(PreTrainedPolicy):
     def get_optim_params(self) -> dict:
         return self.parameters()
 
+    @torch.no_grad
+    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+        """Predict a chunk of actions given environment observations."""
+        raise NotImplementedError("Currently not implemented for PI0")
+
     @torch.no_grad
     def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
         """Select a single action given environment observations.
diff --git a/lerobot/common/policies/pi0/paligemma_with_expert.py b/src/lerobot/policies/pi0/paligemma_with_expert.py
similarity index 99%
rename from lerobot/common/policies/pi0/paligemma_with_expert.py
rename to src/lerobot/policies/pi0/paligemma_with_expert.py
index fb5077fb2..f0f5713e5 100644
--- a/lerobot/common/policies/pi0/paligemma_with_expert.py
+++ b/src/lerobot/policies/pi0/paligemma_with_expert.py
@@ -27,7 +27,7 @@ from transformers import (
 )
 from transformers.models.auto import CONFIG_MAPPING
 
-from lerobot.common.policies.pi0.flex_attention import flex_attention_forward
+from lerobot.policies.pi0.flex_attention import flex_attention_forward
 
 
 def apply_rope(x, positions, max_wavelength=10_000):
diff --git a/lerobot/common/policies/pi0fast/configuration_pi0fast.py b/src/lerobot/policies/pi0fast/configuration_pi0fast.py
similarity index 97%
rename from lerobot/common/policies/pi0fast/configuration_pi0fast.py
rename to src/lerobot/policies/pi0fast/configuration_pi0fast.py
index 29c856e06..b72bcd735 100644
--- a/lerobot/common/policies/pi0fast/configuration_pi0fast.py
+++ b/src/lerobot/policies/pi0fast/configuration_pi0fast.py
@@ -1,11 +1,11 @@
 from dataclasses import dataclass, field
 
-from lerobot.common.optim.optimizers import AdamWConfig
-from lerobot.common.optim.schedulers import (
-    CosineDecayWithWarmupSchedulerConfig,
-)
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim.optimizers import AdamWConfig
+from lerobot.optim.schedulers import (
+    CosineDecayWithWarmupSchedulerConfig,
+)
 
 
 @PreTrainedConfig.register_subclass("pi0fast")
diff --git a/lerobot/common/policies/pi0fast/modeling_pi0fast.py b/src/lerobot/policies/pi0fast/modeling_pi0fast.py
similarity index 98%
rename from lerobot/common/policies/pi0fast/modeling_pi0fast.py
rename to src/lerobot/policies/pi0fast/modeling_pi0fast.py
index 7102bdded..d3e576d1c 100644
--- a/lerobot/common/policies/pi0fast/modeling_pi0fast.py
+++ b/src/lerobot/policies/pi0fast/modeling_pi0fast.py
@@ -24,14 +24,14 @@ Designed by Physical Intelligence. Ported from Jax by Hugging Face.
 
 Example of finetuning the pi0+FAST pretrained model (`pi0_fast_base` in `openpi`):
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
 --policy.path=lerobot/pi0fast_base \
 --dataset.repo_id=danaaubakirova/koch_test
 ```
 
 Example of training the pi0+FAST neural network with from scratch:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
 --policy.type=pi0fast \
 --dataset.repo_id=danaaubakirova/koch_test
 ```
@@ -56,10 +56,10 @@ from transformers import AutoProcessor, AutoTokenizer, PaliGemmaForConditionalGe
 from transformers.cache_utils import HybridCache, StaticCache
 from transformers.models.auto import CONFIG_MAPPING
 
-from lerobot.common.constants import ACTION, OBS_STATE
-from lerobot.common.policies.normalize import Normalize, Unnormalize
-from lerobot.common.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
-from lerobot.common.policies.pretrained import PreTrainedPolicy
+from lerobot.constants import ACTION, OBS_STATE
+from lerobot.policies.normalize import Normalize, Unnormalize
+from lerobot.policies.pi0fast.configuration_pi0fast import PI0FASTConfig
+from lerobot.policies.pretrained import PreTrainedPolicy
 
 PRECISION = {
     "float16": torch.float16,
@@ -192,6 +192,11 @@ class PI0FASTPolicy(PreTrainedPolicy):
             actions[:, :, motor_idx] = aloha_gripper_from_angular_inv(actions[:, :, motor_idx])
         return actions
 
+    @torch.no_grad
+    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+        """Predict a chunk of actions given environment observations."""
+        raise NotImplementedError("Currently not implemented for PI0FAST")
+
     @torch.no_grad
     def select_action(self, batch: dict[str, Tensor]) -> Tensor:
         """Select a single action given environment observations.
diff --git a/lerobot/common/policies/pretrained.py b/src/lerobot/policies/pretrained.py
similarity index 71%
rename from lerobot/common/policies/pretrained.py
rename to src/lerobot/policies/pretrained.py
index da4ef1572..d18b798a8 100644
--- a/lerobot/common/policies/pretrained.py
+++ b/src/lerobot/policies/pretrained.py
@@ -14,34 +14,26 @@
 import abc
 import logging
 import os
+from importlib.resources import files
 from pathlib import Path
-from typing import Type, TypeVar
+from tempfile import TemporaryDirectory
+from typing import List, Type, TypeVar
 
 import packaging
 import safetensors
-from huggingface_hub import hf_hub_download
+from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from huggingface_hub.errors import HfHubHTTPError
 from safetensors.torch import load_model as load_model_as_safetensor
 from safetensors.torch import save_model as save_model_as_safetensor
 from torch import Tensor, nn
 
-from lerobot.common.utils.hub import HubMixin
 from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.train import TrainPipelineConfig
+from lerobot.utils.hub import HubMixin
 
 T = TypeVar("T", bound="PreTrainedPolicy")
 
-DEFAULT_POLICY_CARD = """
----
-# For reference on model card metadata, see the spec: https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1
-# Doc / guide: https://huggingface.co/docs/hub/model-cards
-{{ card_data }}
----
-
-This policy has been pushed to the Hub using [LeRobot](https://github.com/huggingface/lerobot):
-- Docs: {{ docs_url | default("[More Information Needed]", true) }}
-"""
-
 
 class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
     """
@@ -150,16 +142,6 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
             safetensors.torch.load_model(model, model_file, strict=strict, device=map_location)
         return model
 
-    # def generate_model_card(self, *args, **kwargs) -> ModelCard:
-    #     card = ModelCard.from_template(
-    #         card_data=self._hub_mixin_info.model_card_data,
-    #         template_str=self._hub_mixin_info.model_card_template,
-    #         repo_url=self._hub_mixin_info.repo_url,
-    #         docs_url=self._hub_mixin_info.docs_url,
-    #         **kwargs,
-    #     )
-    #     return card
-
     @abc.abstractmethod
     def get_optim_params(self) -> dict:
         """
@@ -189,6 +171,15 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         """
         raise NotImplementedError
 
+    @abc.abstractmethod
+    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+        """Returns the action chunk (for action chunking policies) for a given observation, potentially in batch mode.
+
+        Child classes using action chunking should use this method within `select_action` to form the action chunk
+        cached for selection.
+        """
+        raise NotImplementedError
+
     @abc.abstractmethod
     def select_action(self, batch: dict[str, Tensor]) -> Tensor:
         """Return one action to run in the environment (potentially in batch mode).
@@ -197,3 +188,56 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         with caching.
         """
         raise NotImplementedError
+
+    def push_model_to_hub(
+        self,
+        cfg: TrainPipelineConfig,
+    ):
+        api = HfApi()
+        repo_id = api.create_repo(
+            repo_id=self.config.repo_id, private=self.config.private, exist_ok=True
+        ).repo_id
+
+        # Push the files to the repo in a single commit
+        with TemporaryDirectory(ignore_cleanup_errors=True) as tmp:
+            saved_path = Path(tmp) / repo_id
+
+            self.save_pretrained(saved_path)  # Calls _save_pretrained and stores model tensors
+
+            card = self.generate_model_card(
+                cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags
+            )
+            card.save(str(saved_path / "README.md"))
+
+            cfg.save_pretrained(saved_path)  # Calls _save_pretrained and stores train config
+
+            commit_info = api.upload_folder(
+                repo_id=repo_id,
+                repo_type="model",
+                folder_path=saved_path,
+                commit_message="Upload policy weights, train config and readme",
+                allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.md"],
+                ignore_patterns=["*.tmp", "*.log"],
+            )
+
+            logging.info(f"Model pushed to {commit_info.repo_url.url}")
+
+    def generate_model_card(
+        self, dataset_repo_id: str, model_type: str, license: str | None, tags: List[str] | None
+    ) -> ModelCard:
+        base_model = "lerobot/smolvla_base" if model_type == "smolvla" else None  # Set a base model
+
+        card_data = ModelCardData(
+            license=license or "apache-2.0",
+            library_name="lerobot",
+            pipeline_tag="robotics",
+            tags=list(set(tags or []).union({"robotics", "lerobot", model_type})),
+            model_name=model_type,
+            datasets=dataset_repo_id,
+            base_model=base_model,
+        )
+
+        template_card = files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text()
+        card = ModelCard.from_template(card_data, template_str=template_card)
+        card.validate()
+        return card
diff --git a/lerobot/common/policies/sac/configuration_sac.py b/src/lerobot/policies/sac/configuration_sac.py
similarity index 98%
rename from lerobot/common/policies/sac/configuration_sac.py
rename to src/lerobot/policies/sac/configuration_sac.py
index db58beb2f..c57eeeb72 100644
--- a/lerobot/common/policies/sac/configuration_sac.py
+++ b/src/lerobot/policies/sac/configuration_sac.py
@@ -17,10 +17,10 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.common.constants import ACTION, OBS_IMAGE, OBS_STATE
-from lerobot.common.optim.optimizers import MultiAdamConfig
 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
 
 
 def is_image_feature(key: str) -> bool:
diff --git a/lerobot/common/policies/sac/modeling_sac.py b/src/lerobot/policies/sac/modeling_sac.py
similarity index 98%
rename from lerobot/common/policies/sac/modeling_sac.py
rename to src/lerobot/policies/sac/modeling_sac.py
index b588115ea..54ea122a8 100644
--- a/lerobot/common/policies/sac/modeling_sac.py
+++ b/src/lerobot/policies/sac/modeling_sac.py
@@ -27,10 +27,10 @@ import torch.nn.functional as F  # noqa: N812
 from torch import Tensor
 from torch.distributions import MultivariateNormal, TanhTransform, Transform, TransformedDistribution
 
-from lerobot.common.policies.normalize import NormalizeBuffer
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.policies.sac.configuration_sac import SACConfig, is_image_feature
-from lerobot.common.policies.utils import get_device_from_parameters
+from lerobot.policies.normalize import NormalizeBuffer
+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
 
 DISCRETE_DIMENSION_INDEX = -1  # Gripper is always the last dimension
 
@@ -76,6 +76,11 @@ class SACPolicy(
         """Reset the policy"""
         pass
 
+    @torch.no_grad
+    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+        """Predict a chunk of actions given environment observations."""
+        raise NotImplementedError("SACPolicy does not support action chunking. It returns single actions!")
+
     @torch.no_grad()
     def select_action(self, batch: dict[str, Tensor]) -> Tensor:
         """Select action for inference/evaluation"""
diff --git a/lerobot/common/policies/sac/reward_model/configuration_classifier.py b/src/lerobot/policies/sac/reward_model/configuration_classifier.py
similarity index 94%
rename from lerobot/common/policies/sac/reward_model/configuration_classifier.py
rename to src/lerobot/policies/sac/reward_model/configuration_classifier.py
index 6e2a551d4..fc53283b3 100644
--- a/lerobot/common/policies/sac/reward_model/configuration_classifier.py
+++ b/src/lerobot/policies/sac/reward_model/configuration_classifier.py
@@ -15,10 +15,10 @@
 # limitations under the License.
 from dataclasses import dataclass, field
 
-from lerobot.common.optim.optimizers import AdamWConfig, OptimizerConfig
-from lerobot.common.optim.schedulers import LRSchedulerConfig
 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
 
 
 @PreTrainedConfig.register_subclass(name="reward_classifier")
diff --git a/lerobot/common/policies/sac/reward_model/modeling_classifier.py b/src/lerobot/policies/sac/reward_model/modeling_classifier.py
similarity index 94%
rename from lerobot/common/policies/sac/reward_model/modeling_classifier.py
rename to src/lerobot/policies/sac/reward_model/modeling_classifier.py
index f537e3aef..cadd1c9f2 100644
--- a/lerobot/common/policies/sac/reward_model/modeling_classifier.py
+++ b/src/lerobot/policies/sac/reward_model/modeling_classifier.py
@@ -19,10 +19,10 @@ import logging
 import torch
 from torch import Tensor, nn
 
-from lerobot.common.constants import OBS_IMAGE, REWARD
-from lerobot.common.policies.normalize import Normalize, Unnormalize
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
+from lerobot.constants import OBS_IMAGE, REWARD
+from lerobot.policies.normalize import Normalize, Unnormalize
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
 
 
 class ClassifierOutput:
@@ -308,6 +308,13 @@ class Classifier(PreTrainedPolicy):
         """
         raise NotImplementedError("Reward classifiers do not select actions")
 
+    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+        """
+        This method is required by PreTrainedPolicy but not used for reward classifiers.
+        The reward classifier is not an actor and does not produce action chunks.
+        """
+        raise NotImplementedError("Reward classifiers do not predict action chunks")
+
     def reset(self):
         """
         This method is required by PreTrainedPolicy but not used for reward classifiers.
diff --git a/lerobot/common/policies/smolvla/configuration_smolvla.py b/src/lerobot/policies/smolvla/configuration_smolvla.py
similarity index 98%
rename from lerobot/common/policies/smolvla/configuration_smolvla.py
rename to src/lerobot/policies/smolvla/configuration_smolvla.py
index 5996cf2e7..571900c4a 100644
--- a/lerobot/common/policies/smolvla/configuration_smolvla.py
+++ b/src/lerobot/policies/smolvla/configuration_smolvla.py
@@ -14,12 +14,12 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.common.optim.optimizers import AdamWConfig
-from lerobot.common.optim.schedulers import (
-    CosineDecayWithWarmupSchedulerConfig,
-)
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.optim.optimizers import AdamWConfig
+from lerobot.optim.schedulers import (
+    CosineDecayWithWarmupSchedulerConfig,
+)
 
 
 @PreTrainedConfig.register_subclass("smolvla")
diff --git a/lerobot/common/policies/smolvla/modeling_smolvla.py b/src/lerobot/policies/smolvla/modeling_smolvla.py
similarity index 94%
rename from lerobot/common/policies/smolvla/modeling_smolvla.py
rename to src/lerobot/policies/smolvla/modeling_smolvla.py
index 5e0a9622e..11bb8bf52 100644
--- a/lerobot/common/policies/smolvla/modeling_smolvla.py
+++ b/src/lerobot/policies/smolvla/modeling_smolvla.py
@@ -28,7 +28,7 @@ pip install -e ".[smolvla]"
 
 Example of finetuning the smolvla pretrained model (`smolvla_base`):
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
 --policy.path=lerobot/smolvla_base \
 --dataset.repo_id=danaaubakirova/svla_so100_task1_v3 \
 --batch_size=64 \
@@ -38,7 +38,7 @@ python lerobot/scripts/train.py \
 Example of finetuning a smolVLA. SmolVLA is composed of a pretrained VLM,
 and an action expert.
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
 --policy.type=smolvla \
 --dataset.repo_id=danaaubakirova/svla_so100_task1_v3 \
 --batch_size=64 \
@@ -63,18 +63,18 @@ import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
 from transformers import AutoProcessor
 
-from lerobot.common.constants import ACTION, OBS_STATE
-from lerobot.common.policies.normalize import (
+from lerobot.constants import ACTION, OBS_STATE
+from lerobot.policies.normalize import (
     Normalize,
     Unnormalize,
 )
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.policies.smolvla.configuration_smolvla import SmolVLAConfig
-from lerobot.common.policies.smolvla.smolvlm_with_expert import SmolVLMWithExpertModel
-from lerobot.common.policies.utils import (
+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.common.utils.utils import get_safe_dtype
+from lerobot.utils.utils import get_safe_dtype
 
 # Matches ".soNNN", optionally followed by "-something", up to the "_buffer_" marker
 _VARIANT_RE = re.compile(r"\.so\d+(?:-[\w]+)?_buffer_")
@@ -383,6 +383,45 @@ class SmolVLAPolicy(PreTrainedPolicy):
     def get_optim_params(self) -> dict:
         return self.parameters()
 
+    def _get_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
+        for k in batch:
+            if k in self._queues:
+                batch[k] = torch.stack(list(self._queues[k]), dim=1)
+
+        images, img_masks = self.prepare_images(batch)
+        state = self.prepare_state(batch)
+        lang_tokens, lang_masks = self.prepare_language(batch)
+
+        actions = self.model.sample_actions(images, img_masks, lang_tokens, lang_masks, state, noise=noise)
+
+        # Unpad actions
+        original_action_dim = self.config.action_feature.shape[0]
+        actions = actions[:, :, :original_action_dim]
+
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
+
+        if self.config.adapt_to_pi_aloha:
+            actions = self._pi_aloha_encode_actions(actions)
+
+        return actions
+
+    def _prepare_batch(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
+        if self.config.adapt_to_pi_aloha:
+            batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
+
+        batch = self.normalize_inputs(batch)
+
+        return batch
+
+    def predict_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
+        self.eval()
+
+        batch = self._prepare_batch(batch)
+        self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
+
+        actions = self._get_action_chunk(batch, noise)
+        return actions
+
     @torch.no_grad
     def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
         """Select a single action given environment observations.
@@ -392,38 +431,18 @@ class SmolVLAPolicy(PreTrainedPolicy):
         queue is empty.
         """
         self.eval()
-
-        if self.config.adapt_to_pi_aloha:
-            batch[OBS_STATE] = self._pi_aloha_decode_state(batch[OBS_STATE])
-
-        batch = self.normalize_inputs(batch)
-
+        batch = self._prepare_batch(batch)
         self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
+
         # Action queue logic for n_action_steps > 1. When the action_queue is depleted, populate it by
         # querying the policy.
         if len(self._queues[ACTION]) == 0:
-            for k in batch:
-                if k in self._queues:
-                    batch[k] = torch.stack(list(self._queues[k]), dim=1)
-            images, img_masks = self.prepare_images(batch)
-            state = self.prepare_state(batch)
-            lang_tokens, lang_masks = self.prepare_language(batch)
+            actions = self._get_action_chunk(batch, noise)
 
-            actions = self.model.sample_actions(
-                images, img_masks, lang_tokens, lang_masks, state, noise=noise
-            )
-            # Unpad actions
-            original_action_dim = self.config.action_feature.shape[0]
-            actions = actions[:, :, :original_action_dim]
-
-            actions = self.unnormalize_outputs({"action": actions})["action"]
-
-            if self.config.adapt_to_pi_aloha:
-                actions = self._pi_aloha_encode_actions(actions)
-
-            # `self.model.forward` returns a (batch_size, n_action_steps, action_dim) tensor, but the queue
+            # `self.predict_action_chunk` returns a (batch_size, n_action_steps, action_dim) tensor, but the queue
             # effectively has shape (n_action_steps, batch_size, *), hence the transpose.
             self._queues[ACTION].extend(actions.transpose(0, 1)[: self.config.n_action_steps])
+
         return self._queues[ACTION].popleft()
 
     def forward(self, batch: dict[str, Tensor], noise=None, time=None) -> dict[str, Tensor]:
diff --git a/lerobot/common/policies/smolvla/smolvlm_with_expert.py b/src/lerobot/policies/smolvla/smolvlm_with_expert.py
similarity index 100%
rename from lerobot/common/policies/smolvla/smolvlm_with_expert.py
rename to src/lerobot/policies/smolvla/smolvlm_with_expert.py
diff --git a/lerobot/common/policies/tdmpc/configuration_tdmpc.py b/src/lerobot/policies/tdmpc/configuration_tdmpc.py
similarity index 99%
rename from lerobot/common/policies/tdmpc/configuration_tdmpc.py
rename to src/lerobot/policies/tdmpc/configuration_tdmpc.py
index 3fce01df9..3c1a29932 100644
--- a/lerobot/common/policies/tdmpc/configuration_tdmpc.py
+++ b/src/lerobot/policies/tdmpc/configuration_tdmpc.py
@@ -16,9 +16,9 @@
 # limitations under the License.
 from dataclasses import dataclass, field
 
-from lerobot.common.optim.optimizers import AdamConfig
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
+from lerobot.optim.optimizers import AdamConfig
 
 
 @PreTrainedConfig.register_subclass("tdmpc")
diff --git a/lerobot/common/policies/tdmpc/modeling_tdmpc.py b/src/lerobot/policies/tdmpc/modeling_tdmpc.py
similarity index 93%
rename from lerobot/common/policies/tdmpc/modeling_tdmpc.py
rename to src/lerobot/policies/tdmpc/modeling_tdmpc.py
index 476e6decd..8b70b265d 100644
--- a/lerobot/common/policies/tdmpc/modeling_tdmpc.py
+++ b/src/lerobot/policies/tdmpc/modeling_tdmpc.py
@@ -35,11 +35,11 @@ import torch.nn as nn
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor
 
-from lerobot.common.constants import OBS_ENV_STATE, OBS_STATE
-from lerobot.common.policies.normalize import Normalize, Unnormalize
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.policies.tdmpc.configuration_tdmpc import TDMPCConfig
-from lerobot.common.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
+from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGE, OBS_STATE, REWARD
+from lerobot.policies.normalize import Normalize, Unnormalize
+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
 
 
 class TDMPCPolicy(PreTrainedPolicy):
@@ -110,52 +110,58 @@ class TDMPCPolicy(PreTrainedPolicy):
         # CEM for the next step.
         self._prev_mean: torch.Tensor | None = None
 
+    @torch.no_grad
+    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+        """Predict a chunk of actions given environment observations."""
+        batch = {key: torch.stack(list(self._queues[key]), dim=1) for key in batch if key in self._queues}
+
+        # Remove the time dimensions as it is not handled yet.
+        for key in batch:
+            assert batch[key].shape[1] == 1
+            batch[key] = batch[key][:, 0]
+
+        # NOTE: Order of observations matters here.
+        encode_keys = []
+        if self.config.image_features:
+            encode_keys.append(OBS_IMAGE)
+        if self.config.env_state_feature:
+            encode_keys.append(OBS_ENV_STATE)
+        encode_keys.append(OBS_STATE)
+        z = self.model.encode({k: batch[k] for k in encode_keys})
+        if self.config.use_mpc:  # noqa: SIM108
+            actions = self.plan(z)  # (horizon, batch, action_dim)
+        else:
+            # Plan with the policy (π) alone. This always returns one action so unsqueeze to get a
+            # sequence dimension like in the MPC branch.
+            actions = self.model.pi(z).unsqueeze(0)
+
+        actions = torch.clamp(actions, -1, +1)
+
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
+        return actions
+
     @torch.no_grad()
     def select_action(self, batch: dict[str, Tensor]) -> Tensor:
         """Select a single action given environment observations."""
         batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
-            batch["observation.image"] = batch[next(iter(self.config.image_features))]
+            batch[OBS_IMAGE] = batch[next(iter(self.config.image_features))]
 
         self._queues = populate_queues(self._queues, batch)
 
         # When the action queue is depleted, populate it again by querying the policy.
-        if len(self._queues["action"]) == 0:
-            batch = {key: torch.stack(list(self._queues[key]), dim=1) for key in batch if key in self._queues}
-
-            # Remove the time dimensions as it is not handled yet.
-            for key in batch:
-                assert batch[key].shape[1] == 1
-                batch[key] = batch[key][:, 0]
-
-            # NOTE: Order of observations matters here.
-            encode_keys = []
-            if self.config.image_features:
-                encode_keys.append("observation.image")
-            if self.config.env_state_feature:
-                encode_keys.append("observation.environment_state")
-            encode_keys.append("observation.state")
-            z = self.model.encode({k: batch[k] for k in encode_keys})
-            if self.config.use_mpc:  # noqa: SIM108
-                actions = self.plan(z)  # (horizon, batch, action_dim)
-            else:
-                # Plan with the policy (π) alone. This always returns one action so unsqueeze to get a
-                # sequence dimension like in the MPC branch.
-                actions = self.model.pi(z).unsqueeze(0)
-
-            actions = torch.clamp(actions, -1, +1)
-
-            actions = self.unnormalize_outputs({"action": actions})["action"]
+        if len(self._queues[ACTION]) == 0:
+            actions = self.predict_action_chunk(batch)
 
             if self.config.n_action_repeats > 1:
                 for _ in range(self.config.n_action_repeats):
-                    self._queues["action"].append(actions[0])
+                    self._queues[ACTION].append(actions[0])
             else:
                 # Action queue is (n_action_steps, batch_size, action_dim), so we transpose the action.
-                self._queues["action"].extend(actions[: self.config.n_action_steps])
+                self._queues[ACTION].extend(actions[: self.config.n_action_steps])
 
-        action = self._queues["action"].popleft()
+        action = self._queues[ACTION].popleft()
         return action
 
     @torch.no_grad()
@@ -312,7 +318,7 @@ class TDMPCPolicy(PreTrainedPolicy):
         batch = self.normalize_inputs(batch)
         if self.config.image_features:
             batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
-            batch["observation.image"] = batch[next(iter(self.config.image_features))]
+            batch[OBS_IMAGE] = batch[next(iter(self.config.image_features))]
         batch = self.normalize_targets(batch)
 
         info = {}
@@ -322,15 +328,15 @@ class TDMPCPolicy(PreTrainedPolicy):
             if isinstance(batch[key], torch.Tensor) and batch[key].ndim > 1:
                 batch[key] = batch[key].transpose(1, 0)
 
-        action = batch["action"]  # (t, b, action_dim)
-        reward = batch["next.reward"]  # (t, b)
+        action = batch[ACTION]  # (t, b, action_dim)
+        reward = batch[REWARD]  # (t, b)
         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:
-            observations["observation.image"] = flatten_forward_unflatten(
+            observations[OBS_IMAGE] = flatten_forward_unflatten(
                 partial(random_shifts_aug, max_random_shift_ratio=self.config.max_random_shift_ratio),
-                observations["observation.image"],
+                observations[OBS_IMAGE],
             )
 
         # Get the current observation for predicting trajectories, and all future observations for use in
@@ -340,7 +346,7 @@ class TDMPCPolicy(PreTrainedPolicy):
             current_observation[k] = observations[k][0]
             next_observations[k] = observations[k][1:]
         horizon, batch_size = next_observations[
-            "observation.image" if self.config.image_features else "observation.environment_state"
+            OBS_IMAGE if self.config.image_features else OBS_ENV_STATE
         ].shape[:2]
 
         # Run latent rollout using the latent dynamics model and policy model.
diff --git a/lerobot/common/policies/utils.py b/src/lerobot/policies/utils.py
similarity index 100%
rename from lerobot/common/policies/utils.py
rename to src/lerobot/policies/utils.py
diff --git a/lerobot/common/policies/vqbet/configuration_vqbet.py b/src/lerobot/policies/vqbet/configuration_vqbet.py
similarity index 98%
rename from lerobot/common/policies/vqbet/configuration_vqbet.py
rename to src/lerobot/policies/vqbet/configuration_vqbet.py
index 28e9c4338..d7a79f189 100644
--- a/lerobot/common/policies/vqbet/configuration_vqbet.py
+++ b/src/lerobot/policies/vqbet/configuration_vqbet.py
@@ -18,10 +18,10 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.common.optim.optimizers import AdamConfig
-from lerobot.common.optim.schedulers import VQBeTSchedulerConfig
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import NormalizationMode
+from lerobot.optim.optimizers import AdamConfig
+from lerobot.optim.schedulers import VQBeTSchedulerConfig
 
 
 @PreTrainedConfig.register_subclass("vqbet")
diff --git a/lerobot/common/policies/vqbet/modeling_vqbet.py b/src/lerobot/policies/vqbet/modeling_vqbet.py
similarity index 96%
rename from lerobot/common/policies/vqbet/modeling_vqbet.py
rename to src/lerobot/policies/vqbet/modeling_vqbet.py
index 44006a5b2..c045ccbd2 100644
--- a/lerobot/common/policies/vqbet/modeling_vqbet.py
+++ b/src/lerobot/policies/vqbet/modeling_vqbet.py
@@ -27,11 +27,12 @@ import torch.nn.functional as F  # noqa: N812
 import torchvision
 from torch import Tensor, nn
 
-from lerobot.common.policies.normalize import Normalize, Unnormalize
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.policies.utils import get_device_from_parameters, get_output_shape, populate_queues
-from lerobot.common.policies.vqbet.configuration_vqbet import VQBeTConfig
-from lerobot.common.policies.vqbet.vqbet_utils import GPT, ResidualVQ
+from lerobot.constants import ACTION, OBS_IMAGES, OBS_STATE
+from lerobot.policies.normalize import Normalize, Unnormalize
+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
 
 # ruff: noqa: N806
 
@@ -118,11 +119,18 @@ class VQBeTPolicy(PreTrainedPolicy):
         queues are populated during rollout of the policy, they contain the n latest observations and actions
         """
         self._queues = {
-            "observation.images": deque(maxlen=self.config.n_obs_steps),
-            "observation.state": deque(maxlen=self.config.n_obs_steps),
-            "action": deque(maxlen=self.config.action_chunk_size),
+            OBS_IMAGES: deque(maxlen=self.config.n_obs_steps),
+            OBS_STATE: deque(maxlen=self.config.n_obs_steps),
+            ACTION: deque(maxlen=self.config.action_chunk_size),
         }
 
+    @torch.no_grad
+    def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
+        batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
+        actions = self.vqbet(batch, rollout=True)[:, : self.config.action_chunk_size]
+        actions = self.unnormalize_outputs({ACTION: actions})[ACTION]
+        return actions
+
     @torch.no_grad
     def select_action(self, batch: dict[str, Tensor]) -> Tensor:
         """Select a single action given environment observations.
@@ -144,23 +152,19 @@ class VQBeTPolicy(PreTrainedPolicy):
                 stacklevel=1,
             )
 
-        if len(self._queues["action"]) == 0:
-            batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues}
-            actions = self.vqbet(batch, rollout=True)[:, : self.config.action_chunk_size]
-
-            # the dimension of returned action is (batch_size, action_chunk_size, action_dim)
-            actions = self.unnormalize_outputs({"action": actions})["action"]
+        if len(self._queues[ACTION]) == 0:
+            actions = self.predict_action_chunk(batch)
             # since the data in the action queue's dimension is (action_chunk_size, batch_size, action_dim), we transpose the action and fill the queue
-            self._queues["action"].extend(actions.transpose(0, 1))
+            self._queues[ACTION].extend(actions.transpose(0, 1))
 
-        action = self._queues["action"].popleft()
+        action = self._queues[ACTION].popleft()
         return action
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
         """Run the batch through the model and compute the loss for training or validation."""
         batch = self.normalize_inputs(batch)
         batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
-        batch["observation.images"] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
+        batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
         batch = self.normalize_targets(batch)
         # VQ-BeT discretizes action using VQ-VAE before training BeT (please refer to section 3.2 in the VQ-BeT paper https://huggingface.co/papers/2403.03181)
         if not self.vqbet.action_head.vqvae_model.discretized.item():
@@ -168,7 +172,7 @@ class VQBeTPolicy(PreTrainedPolicy):
             # n_different_codes: how many of the total possible VQ codes are being used in single batch (how many of them have at least one encoder embedding as a nearest neighbor). This can be at most `vqvae_n_embed * number of layers of RVQ (=2)`.
             # n_different_combinations: how many different code combinations are being used out of all possible combinations in single batch. This can be at most `vqvae_n_embed ^ number of layers of RVQ (=2)` (hint consider the RVQ as a decision tree).
             loss, n_different_codes, n_different_combinations, recon_l1_error = (
-                self.vqbet.action_head.discretize(self.config.n_vqvae_training_steps, batch["action"])
+                self.vqbet.action_head.discretize(self.config.n_vqvae_training_steps, batch[ACTION])
             )
             return loss, {
                 "n_different_codes": n_different_codes,
@@ -404,7 +408,7 @@ class VQBeTModel(nn.Module):
             )
         # else, it calculate overall loss (bin prediction loss, and offset loss)
         else:
-            output = batch["action"][:, self.select_target_actions_indices]
+            output = batch[ACTION][:, self.select_target_actions_indices]
             loss = self.action_head.loss_fn(action_head_output, output, reduction="mean")
             return action_head_output, loss
 
diff --git a/lerobot/common/policies/vqbet/vqbet_utils.py b/src/lerobot/policies/vqbet/vqbet_utils.py
similarity index 99%
rename from lerobot/common/policies/vqbet/vqbet_utils.py
rename to src/lerobot/policies/vqbet/vqbet_utils.py
index 09a86c07b..03b02a280 100644
--- a/lerobot/common/policies/vqbet/vqbet_utils.py
+++ b/src/lerobot/policies/vqbet/vqbet_utils.py
@@ -30,7 +30,7 @@ from torch import einsum, nn
 from torch.cuda.amp import autocast
 from torch.optim import Optimizer
 
-from lerobot.common.policies.vqbet.configuration_vqbet import VQBeTConfig
+from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
 
 # ruff: noqa: N806
 
diff --git a/lerobot/record.py b/src/lerobot/record.py
similarity index 81%
rename from lerobot/record.py
rename to src/lerobot/record.py
index 2f443c208..635bdf1e4 100644
--- a/lerobot/record.py
+++ b/src/lerobot/record.py
@@ -40,21 +40,21 @@ import time
 from dataclasses import asdict, dataclass
 from pathlib import Path
 from pprint import pformat
+from typing import List
 
-import numpy as np
-import rerun as rr
-
-from lerobot.common.cameras import (  # noqa: F401
+from lerobot.cameras import (  # noqa: F401
     CameraConfig,  # noqa: F401
 )
-from lerobot.common.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
-from lerobot.common.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
-from lerobot.common.datasets.image_writer import safe_stop_image_writer
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.datasets.utils import build_dataset_frame, hw_to_dataset_features
-from lerobot.common.policies.factory import make_policy
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.robots import (  # noqa: F401
+from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
+from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
+from lerobot.configs import parser
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.datasets.image_writer import safe_stop_image_writer
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import build_dataset_frame, hw_to_dataset_features
+from lerobot.policies.factory import make_policy
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.robots import (  # noqa: F401
     Robot,
     RobotConfig,
     koch_follower,
@@ -62,29 +62,29 @@ from lerobot.common.robots import (  # noqa: F401
     so100_follower,
     so101_follower,
 )
-from lerobot.common.teleoperators import (  # noqa: F401
+from lerobot.teleoperators import (  # noqa: F401
     Teleoperator,
     TeleoperatorConfig,
+    koch_leader,
     make_teleoperator_from_config,
+    so100_leader,
+    so101_leader,
 )
-from lerobot.common.utils.control_utils import (
+from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardTeleop
+from lerobot.utils.control_utils import (
     init_keyboard_listener,
     is_headless,
     predict_action,
     sanity_check_dataset_name,
     sanity_check_dataset_robot_compatibility,
 )
-from lerobot.common.utils.robot_utils import busy_wait
-from lerobot.common.utils.utils import (
+from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.utils import (
     get_safe_torch_device,
     init_logging,
     log_say,
 )
-from lerobot.common.utils.visualization_utils import _init_rerun
-from lerobot.configs import parser
-from lerobot.configs.policies import PreTrainedConfig
-
-from .common.teleoperators import koch_leader, so100_leader, so101_leader  # noqa: F401
+from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data
 
 
 @dataclass
@@ -164,7 +164,7 @@ def record_loop(
     events: dict,
     fps: int,
     dataset: LeRobotDataset | None = None,
-    teleop: Teleoperator | None = None,
+    teleop: Teleoperator | List[Teleoperator] | None = None,
     policy: PreTrainedPolicy | None = None,
     control_time_s: int | None = None,
     single_task: str | None = None,
@@ -173,6 +173,23 @@ def record_loop(
     if dataset is not None and dataset.fps != fps:
         raise ValueError(f"The dataset fps should be equal to requested fps ({dataset.fps} != {fps}).")
 
+    teleop_arm = teleop_keyboard = None
+    if isinstance(teleop, list):
+        teleop_keyboard = next((t for t in teleop if isinstance(t, KeyboardTeleop)), None)
+        teleop_arm = next(
+            (
+                t
+                for t in teleop
+                if isinstance(t, (so100_leader.SO100Leader, so101_leader.SO101Leader, koch_leader.KochLeader))
+            ),
+            None,
+        )
+
+        if not (teleop_arm and teleop_keyboard and len(teleop) == 2 and robot.name == "lekiwi_client"):
+            raise ValueError(
+                "For multi-teleop, the list must contain exactly one KeyboardTeleop and one arm teleoperator. Currently only supported for LeKiwi robot."
+            )
+
     # if policy is given it needs cleaning up
     if policy is not None:
         policy.reset()
@@ -201,8 +218,17 @@ def record_loop(
                 robot_type=robot.robot_type,
             )
             action = {key: action_values[i].item() for i, key in enumerate(robot.action_features)}
-        elif policy is None and teleop is not None:
+        elif policy is None and isinstance(teleop, Teleoperator):
             action = teleop.get_action()
+        elif policy is None and isinstance(teleop, list):
+            # TODO(pepijn, steven): clean the record loop for use of multiple robots (possibly with pipeline)
+            arm_action = teleop_arm.get_action()
+            arm_action = {f"arm_{k}": v for k, v in arm_action.items()}
+
+            keyboard_action = teleop_keyboard.get_action()
+            base_action = robot._from_keyboard_to_base_action(keyboard_action)
+
+            action = {**arm_action, **base_action} if len(base_action) > 0 else arm_action
         else:
             logging.info(
                 "No policy or teleoperator provided, skipping action generation."
@@ -221,14 +247,7 @@ def record_loop(
             dataset.add_frame(frame, task=single_task)
 
         if display_data:
-            for obs, val in observation.items():
-                if isinstance(val, float):
-                    rr.log(f"observation.{obs}", rr.Scalar(val))
-                elif isinstance(val, np.ndarray):
-                    rr.log(f"observation.{obs}", rr.Image(val), static=True)
-            for act, val in action.items():
-                if isinstance(val, float):
-                    rr.log(f"action.{act}", rr.Scalar(val))
+            log_rerun_data(observation, action)
 
         dt_s = time.perf_counter() - start_loop_t
         busy_wait(1 / fps - dt_s)
@@ -285,7 +304,8 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
 
     listener, events = init_keyboard_listener()
 
-    for recorded_episodes in range(cfg.dataset.num_episodes):
+    recorded_episodes = 0
+    while recorded_episodes < cfg.dataset.num_episodes and not events["stop_recording"]:
         log_say(f"Recording episode {dataset.num_episodes}", cfg.play_sounds)
         record_loop(
             robot=robot,
@@ -323,9 +343,7 @@ def record(cfg: RecordConfig) -> LeRobotDataset:
             continue
 
         dataset.save_episode()
-
-        if events["stop_recording"]:
-            break
+        recorded_episodes += 1
 
     log_say("Stop recording", cfg.play_sounds, blocking=True)
 
diff --git a/lerobot/replay.py b/src/lerobot/replay.py
similarity index 92%
rename from lerobot/replay.py
rename to src/lerobot/replay.py
index 36eb0864d..ef20c28ef 100644
--- a/lerobot/replay.py
+++ b/src/lerobot/replay.py
@@ -35,8 +35,8 @@ from pprint import pformat
 
 import draccus
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.robots import (  # noqa: F401
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.robots import (  # noqa: F401
     Robot,
     RobotConfig,
     koch_follower,
@@ -44,8 +44,8 @@ from lerobot.common.robots import (  # noqa: F401
     so100_follower,
     so101_follower,
 )
-from lerobot.common.utils.robot_utils import busy_wait
-from lerobot.common.utils.utils import (
+from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.utils import (
     init_logging,
     log_say,
 )
diff --git a/lerobot/common/robots/__init__.py b/src/lerobot/robots/__init__.py
similarity index 100%
rename from lerobot/common/robots/__init__.py
rename to src/lerobot/robots/__init__.py
diff --git a/lerobot/common/robots/config.py b/src/lerobot/robots/config.py
similarity index 100%
rename from lerobot/common/robots/config.py
rename to src/lerobot/robots/config.py
diff --git a/lerobot/common/robots/koch_follower/__init__.py b/src/lerobot/robots/koch_follower/__init__.py
similarity index 100%
rename from lerobot/common/robots/koch_follower/__init__.py
rename to src/lerobot/robots/koch_follower/__init__.py
diff --git a/lerobot/common/robots/koch_follower/config_koch_follower.py b/src/lerobot/robots/koch_follower/config_koch_follower.py
similarity index 96%
rename from lerobot/common/robots/koch_follower/config_koch_follower.py
rename to src/lerobot/robots/koch_follower/config_koch_follower.py
index 6ac164726..a7c9249ae 100644
--- a/lerobot/common/robots/koch_follower/config_koch_follower.py
+++ b/src/lerobot/robots/koch_follower/config_koch_follower.py
@@ -14,7 +14,7 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.common.cameras import CameraConfig
+from lerobot.cameras import CameraConfig
 
 from ..config import RobotConfig
 
diff --git a/lerobot/common/robots/koch_follower/koch.mdx b/src/lerobot/robots/koch_follower/koch.mdx
similarity index 96%
rename from lerobot/common/robots/koch_follower/koch.mdx
rename to src/lerobot/robots/koch_follower/koch.mdx
index c39865944..f70a1802c 100644
--- a/lerobot/common/robots/koch_follower/koch.mdx
+++ b/src/lerobot/robots/koch_follower/koch.mdx
@@ -29,7 +29,7 @@ pip install -e ".[dynamixel]"
 
 To find the port for each bus servo adapter, run this script:
 ```bash
-python lerobot/find_port.py
+python -m lerobot.find_port
 ```
 
 <hfoptions id="example">
@@ -103,7 +103,7 @@ python -m lerobot.setup_motors \
 <hfoption id="API example">
 
 ```python
-from lerobot.common.robots.koch_follower import KochFollower, KochFollowerConfig
+from lerobot.robots.koch_follower import KochFollower, KochFollowerConfig
 
 config = KochFollowerConfig(
     port="/dev/tty.usbmodem575E0031751",
@@ -169,7 +169,7 @@ python -m lerobot.setup_motors \
 <hfoption id="API example">
 
 ```python
-from lerobot.common.teleoperators.koch_leader import KochLeader, KochLeaderConfig
+from lerobot.teleoperators.koch_leader import KochLeader, KochLeaderConfig
 
 config = KochLeaderConfig(
     port="/dev/tty.usbmodem575E0031751",
@@ -203,7 +203,7 @@ python -m lerobot.calibrate \
 <hfoption id="API example">
 
 ```python
-from lerobot.common.robots.koch_follower import KochFollowerConfig, KochFollower
+from lerobot.robots.koch_follower import KochFollowerConfig, KochFollower
 
 config = KochFollowerConfig(
     port="/dev/tty.usbmodem585A0076891",
@@ -237,7 +237,7 @@ python -m lerobot.calibrate \
 <hfoption id="API example">
 
 ```python
-from lerobot.common.teleoperators.koch_leader import KochLeaderConfig, KochLeader
+from lerobot.teleoperators.koch_leader import KochLeaderConfig, KochLeader
 
 config = KochLeaderConfig(
     port="/dev/tty.usbmodem575E0031751",
diff --git a/lerobot/common/robots/koch_follower/koch_follower.py b/src/lerobot/robots/koch_follower/koch_follower.py
similarity index 97%
rename from lerobot/common/robots/koch_follower/koch_follower.py
rename to src/lerobot/robots/koch_follower/koch_follower.py
index 64ece25f2..1cfc6cf08 100644
--- a/lerobot/common/robots/koch_follower/koch_follower.py
+++ b/src/lerobot/robots/koch_follower/koch_follower.py
@@ -19,10 +19,10 @@ import time
 from functools import cached_property
 from typing import Any
 
-from lerobot.common.cameras.utils import make_cameras_from_configs
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
-from lerobot.common.motors import Motor, MotorCalibration, MotorNormMode
-from lerobot.common.motors.dynamixel import (
+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,
 )
diff --git a/lerobot/common/robots/lekiwi/__init__.py b/src/lerobot/robots/lekiwi/__init__.py
similarity index 100%
rename from lerobot/common/robots/lekiwi/__init__.py
rename to src/lerobot/robots/lekiwi/__init__.py
diff --git a/lerobot/common/robots/lekiwi/config_lekiwi.py b/src/lerobot/robots/lekiwi/config_lekiwi.py
similarity index 95%
rename from lerobot/common/robots/lekiwi/config_lekiwi.py
rename to src/lerobot/robots/lekiwi/config_lekiwi.py
index 022d09cdd..f0f8c24b3 100644
--- a/lerobot/common/robots/lekiwi/config_lekiwi.py
+++ b/src/lerobot/robots/lekiwi/config_lekiwi.py
@@ -14,8 +14,8 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.common.cameras.configs import CameraConfig, Cv2Rotation
-from lerobot.common.cameras.opencv.configuration_opencv import OpenCVCameraConfig
+from lerobot.cameras.configs import CameraConfig, Cv2Rotation
+from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig
 
 from ..config import RobotConfig
 
diff --git a/lerobot/common/robots/lekiwi/lekiwi.mdx b/src/lerobot/robots/lekiwi/lekiwi.mdx
similarity index 97%
rename from lerobot/common/robots/lekiwi/lekiwi.mdx
rename to src/lerobot/robots/lekiwi/lekiwi.mdx
index 6eaebce79..61b1c05c1 100644
--- a/lerobot/common/robots/lekiwi/lekiwi.mdx
+++ b/src/lerobot/robots/lekiwi/lekiwi.mdx
@@ -47,7 +47,7 @@ First, we will assemble the two SO100/SO101 arms. One to attach to the mobile ba
 
 To find the port for each bus servo adapter, run this script:
 ```bash
-python lerobot/find_port.py
+python -m lerobot.find_port
 ```
 
 <hfoptions id="example">
@@ -175,7 +175,7 @@ python -m lerobot.calibrate \
 <hfoption id="API example">
 
 ```python
-from lerobot.common.teleoperators.so100_leader import SO100LeaderConfig, SO100Leader
+from lerobot.teleoperators.so100_leader import SO100LeaderConfig, SO100Leader
 
 config = SO100LeaderConfig(
     port="/dev/tty.usbmodem58760431551",
@@ -197,7 +197,7 @@ leader.disconnect()
 
 To teleoperate, SSH into your Raspberry Pi, and run `conda activate lerobot` and this command:
 ```bash
-python -m lerobot.common.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi
+python -m lerobot.robots.lekiwi.lekiwi_host --robot.id=my_awesome_kiwi
 ```
 
 Then on your laptop, also run `conda activate lerobot` and run the API example, make sure you set the correct `remote_ip` and `port` in `examples/lekiwi/teleoperate.py`.
@@ -227,7 +227,7 @@ You should see on your laptop something like this: ```[INFO] Connected to remote
 | F   | Decrease speed |
 
 > [!TIP]
->  If you use a different keyboard, you can change the keys for each command in the [`LeKiwiConfig`](../lerobot/common/robot_devices/robots/configs.py).
+>  If you use a different keyboard, you can change the keys for each command in the [`LeKiwiConfig`](../src/lerobot/robot_devices/robots/configs.py).
 
 ### Wired version
 If you have the **wired** LeKiwi version, please run all commands on your laptop.
diff --git a/lerobot/common/robots/lekiwi/lekiwi.py b/src/lerobot/robots/lekiwi/lekiwi.py
similarity index 98%
rename from lerobot/common/robots/lekiwi/lekiwi.py
rename to src/lerobot/robots/lekiwi/lekiwi.py
index f6a9b8bf1..ff1465d8b 100644
--- a/lerobot/common/robots/lekiwi/lekiwi.py
+++ b/src/lerobot/robots/lekiwi/lekiwi.py
@@ -22,10 +22,10 @@ from typing import Any
 
 import numpy as np
 
-from lerobot.common.cameras.utils import make_cameras_from_configs
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
-from lerobot.common.motors import Motor, MotorCalibration, MotorNormMode
-from lerobot.common.motors.feetech import (
+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,
 )
diff --git a/lerobot/common/robots/lekiwi/lekiwi_client.py b/src/lerobot/robots/lekiwi/lekiwi_client.py
similarity index 94%
rename from lerobot/common/robots/lekiwi/lekiwi_client.py
rename to src/lerobot/robots/lekiwi/lekiwi_client.py
index f79b7f81a..0ce259bb6 100644
--- a/lerobot/common/robots/lekiwi/lekiwi_client.py
+++ b/src/lerobot/robots/lekiwi/lekiwi_client.py
@@ -22,10 +22,9 @@ from typing import Any, Dict, Optional, Tuple
 
 import cv2
 import numpy as np
-import torch
 import zmq
 
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 from ..robot import Robot
 from .config_lekiwi import LeKiwiClientConfig
@@ -195,26 +194,23 @@ class LeKiwiClient(Robot):
         self, observation: Dict[str, Any]
     ) -> Tuple[Dict[str, np.ndarray], Dict[str, Any]]:
         """Extracts frames, and state from the parsed observation."""
-        flat_state = {key: value for key, value in observation.items() if key in self._state_ft}
 
-        state_vec = np.array(
-            [flat_state.get(k, 0.0) for k in self._state_order],
-            dtype=np.float32,
-        )
+        flat_state = {key: observation.get(key, 0.0) for key in self._state_order}
+
+        state_vec = np.array([flat_state[key] for key in self._state_order], dtype=np.float32)
+
+        obs_dict: Dict[str, Any] = {**flat_state, "observation.state": state_vec}
 
         # Decode images
-        image_observation = {
-            f"observation.images.{key}": value
-            for key, value in observation.items()
-            if key in self._cameras_ft
-        }
         current_frames: Dict[str, np.ndarray] = {}
-        for cam_name, image_b64 in image_observation.items():
+        for cam_name, image_b64 in observation.items():
+            if cam_name not in self._cameras_ft:
+                continue
             frame = self._decode_image_from_b64(image_b64)
             if frame is not None:
                 current_frames[cam_name] = frame
 
-        return current_frames, {"observation.state": state_vec}
+        return current_frames, obs_dict
 
     def _get_data(self) -> Tuple[Dict[str, np.ndarray], Dict[str, Any], Dict[str, Any]]:
         """
@@ -267,7 +263,7 @@ class LeKiwiClient(Robot):
             if frame is None:
                 logging.warning("Frame is None")
                 frame = np.zeros((640, 480, 3), dtype=np.uint8)
-            obs_dict[cam_name] = torch.from_numpy(frame)
+            obs_dict[cam_name] = frame
 
         return obs_dict
 
@@ -327,7 +323,10 @@ class LeKiwiClient(Robot):
 
         # TODO(Steven): Remove the np conversion when it is possible to record a non-numpy array value
         actions = np.array([action.get(k, 0.0) for k in self._state_order], dtype=np.float32)
-        return {"action": actions}
+
+        action_sent = {key: actions[i] for i, key in enumerate(self._state_order)}
+        action_sent["action"] = actions
+        return action_sent
 
     def disconnect(self):
         """Cleans ZMQ comms"""
diff --git a/lerobot/common/robots/lekiwi/lekiwi_host.py b/src/lerobot/robots/lekiwi/lekiwi_host.py
similarity index 100%
rename from lerobot/common/robots/lekiwi/lekiwi_host.py
rename to src/lerobot/robots/lekiwi/lekiwi_host.py
diff --git a/lerobot/common/robots/robot.py b/src/lerobot/robots/robot.py
similarity index 98%
rename from lerobot/common/robots/robot.py
rename to src/lerobot/robots/robot.py
index 76c57faf4..6820645cc 100644
--- a/lerobot/common/robots/robot.py
+++ b/src/lerobot/robots/robot.py
@@ -18,8 +18,8 @@ from typing import Any, Type
 
 import draccus
 
-from lerobot.common.constants import HF_LEROBOT_CALIBRATION, ROBOTS
-from lerobot.common.motors import MotorCalibration
+from lerobot.constants import HF_LEROBOT_CALIBRATION, ROBOTS
+from lerobot.motors import MotorCalibration
 
 from .config import RobotConfig
 
diff --git a/lerobot/common/robots/so100_follower/__init__.py b/src/lerobot/robots/so100_follower/__init__.py
similarity index 100%
rename from lerobot/common/robots/so100_follower/__init__.py
rename to src/lerobot/robots/so100_follower/__init__.py
diff --git a/lerobot/common/robots/so100_follower/config_so100_follower.py b/src/lerobot/robots/so100_follower/config_so100_follower.py
similarity index 84%
rename from lerobot/common/robots/so100_follower/config_so100_follower.py
rename to src/lerobot/robots/so100_follower/config_so100_follower.py
index b76675d26..7cd23d340 100644
--- a/lerobot/common/robots/so100_follower/config_so100_follower.py
+++ b/src/lerobot/robots/so100_follower/config_so100_follower.py
@@ -14,7 +14,7 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.common.cameras import CameraConfig
+from lerobot.cameras import CameraConfig
 
 from ..config import RobotConfig
 
@@ -44,6 +44,14 @@ class SO100FollowerConfig(RobotConfig):
 class SO100FollowerEndEffectorConfig(SO100FollowerConfig):
     """Configuration for the SO100FollowerEndEffector robot."""
 
+    # Path to URDF file for kinematics
+    # NOTE: It is highly recommended to use the urdf in the SO-ARM100 repo:
+    # https://github.com/TheRobotStudio/SO-ARM100/blob/main/Simulation/SO101/so101_new_calib.urdf
+    urdf_path: str | None = None
+
+    # End-effector frame name in the URDF
+    target_frame_name: str = "gripper_frame_link"
+
     # Default bounds for the end-effector position (in meters)
     end_effector_bounds: dict[str, list[float]] = field(
         default_factory=lambda: {
diff --git a/lerobot/common/robots/so100_follower/so100.mdx b/src/lerobot/robots/so100_follower/so100.mdx
similarity index 97%
rename from lerobot/common/robots/so100_follower/so100.mdx
rename to src/lerobot/robots/so100_follower/so100.mdx
index d6149b5b8..f5eea6aef 100644
--- a/lerobot/common/robots/so100_follower/so100.mdx
+++ b/src/lerobot/robots/so100_follower/so100.mdx
@@ -15,6 +15,166 @@ In addition to these instructions, you need to install the Feetech SDK:
 pip install -e ".[feetech]"
 ```
 
+## Configure the motors
+
+**Note:**
+Unlike the SO-101, the motor connectors are not easily accessible once the arm is assembled, so the configuration step must be done beforehand.
+
+### 1. Find the USB ports associated with each arm
+
+To find the port for each bus servo adapter, run this script:
+```bash
+python -m lerobot.find_port
+```
+
+<hfoptions id="example">
+<hfoption id="Mac">
+
+Example output:
+
+```
+Finding all available ports for the MotorBus.
+['/dev/tty.usbmodem575E0032081', '/dev/tty.usbmodem575E0031751']
+Remove the USB cable from your MotorsBus and press Enter when done.
+
+[...Disconnect corresponding leader or follower arm and press Enter...]
+
+The port of this MotorsBus is /dev/tty.usbmodem575E0032081
+Reconnect the USB cable.
+```
+
+Where the found port is: `/dev/tty.usbmodem575E0032081` corresponding to your leader or follower arm.
+
+</hfoption>
+<hfoption id="Linux">
+
+On Linux, you might need to give access to the USB ports by running:
+```bash
+sudo chmod 666 /dev/ttyACM0
+sudo chmod 666 /dev/ttyACM1
+```
+
+Example output:
+
+```
+Finding all available ports for the MotorBus.
+['/dev/ttyACM0', '/dev/ttyACM1']
+Remove the usb cable from your MotorsBus and press Enter when done.
+
+[...Disconnect corresponding leader or follower arm and press Enter...]
+
+The port of this MotorsBus is /dev/ttyACM1
+Reconnect the USB cable.
+```
+
+Where the found port is: `/dev/ttyACM1` corresponding to your leader or follower arm.
+
+</hfoption>
+</hfoptions>
+
+### 2. Set the motors ids and baudrates
+
+Each motor is identified by a unique id on the bus. When brand new, motors usually come with a default id of `1`. For the communication to work properly between the motors and the controller, we first need to set a unique, different id to each motor. Additionally, the speed at which data is transmitted on the bus is determined by the baudrate. In order to talk to each other, the controller and all the motors need to be configured with the same baudrate.
+
+To that end, we first need to connect to each motor individually with the controller in order to set these. Since we will write these parameters in the non-volatile section of the motors' internal memory (EEPROM), we'll only need to do this once.
+
+If you are repurposing motors from another robot, you will probably also need to perform this step as the ids and baudrate likely won't match.
+
+#### Follower
+
+Connect the usb cable from your computer and the power supply to the follower arm's controller board. Then, run the following command or run the API example with the port you got from the previous step. You'll also need to give your leader arm a name with the `id` parameter.
+
+For a visual reference on how to set the motor ids please refer to [this video](https://huggingface.co/docs/lerobot/en/so101#setup-motors-video) where we follow the process for the SO101 arm.
+
+<hfoptions id="setup_motors">
+<hfoption id="Command">
+
+```bash
+python -m lerobot.setup_motors \
+    --robot.type=so100_follower \
+    --robot.port=/dev/tty.usbmodem585A0076841  # <- paste here the port found at previous step
+```
+</hfoption>
+<hfoption id="API example">
+
+```python
+from lerobot.robots.so100_follower import SO100Follower, SO100FollowerConfig
+
+config = SO100FollowerConfig(
+    port="/dev/tty.usbmodem585A0076841",
+    id="my_awesome_follower_arm",
+)
+follower = SO100Follower(config)
+follower.setup_motors()
+```
+</hfoption>
+</hfoptions>
+
+You should see the following instruction
+```
+Connect the controller board to the 'gripper' motor only and press enter.
+```
+
+As instructed, plug the gripper's motor. Make sure it's the only motor connected to the board, and that the motor itself is not yet daisy-chained to any other motor. As you press `[Enter]`, the script will automatically set the id and baudrate for that motor.
+
+<details>
+<summary>Troubleshooting</summary>
+
+  If you get an error at that point, check your cables and make sure they are plugged in properly:
+  <ul>
+    <li>Power supply</li>
+    <li>USB cable between your computer and the controller board</li>
+    <li>The 3-pin cable from the controller board to the motor</li>
+  </ul>
+
+If you are using a Waveshare controller board, make sure that the two jumpers are set on the `B` channel (USB).
+</details>
+
+You should then see the following message:
+```
+'gripper' motor id set to 6
+```
+
+Followed by the next instruction:
+```
+Connect the controller board to the 'wrist_roll' motor only and press enter.
+```
+
+You can disconnect the 3-pin cable from the controller board, but you can leave it connected to the gripper motor on the other end, as it will already be in the right place. Now, plug in another 3-pin cable to the wrist roll motor and connect it to the controller board. As with the previous motor, make sure it is the only motor connected to the board and that the motor itself isn't connected to any other one.
+
+Repeat the operation for each motor as instructed.
+
+> [!TIP]
+> Check your cabling at each step before pressing Enter. For instance, the power supply cable might disconnect as you manipulate the board.
+
+When you are done, the script will simply finish, at which point the motors are ready to be used. You can now plug the 3-pin cable from each motor to the next one, and the cable from the first motor (the 'shoulder pan' with id=1) to the controller board, which can now be attached to the base of the arm.
+
+#### Leader
+Do the same steps for the leader arm.
+
+<hfoptions id="setup_motors">
+<hfoption id="Command">
+```bash
+python -m lerobot.setup_motors \
+    --teleop.type=so100_leader \
+    --teleop.port=/dev/tty.usbmodem575E0031751  # <- paste here the port found at previous step
+```
+</hfoption>
+<hfoption id="API example">
+
+```python
+from lerobot.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
+
+config = SO100LeaderConfig(
+    port="/dev/tty.usbmodem585A0076841",
+    id="my_awesome_leader_arm",
+)
+leader = SO100Leader(config)
+leader.setup_motors()
+```
+</hfoption>
+</hfoptions>
+
 ## Step-by-Step Assembly Instructions
 
 ## Remove the gears of the 6 leader motors
@@ -252,163 +412,6 @@ For the leader configuration, perform **Steps 1–23**. Make sure that you remov
   <img src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/so100_assembly_28.webp" style="height:300px;"/>
 </div>
 
-## Configure the motors
-
-### 1. Find the USB ports associated with each arm
-
-To find the port for each bus servo adapter, run this script:
-```bash
-python lerobot/find_port.py
-```
-
-<hfoptions id="example">
-<hfoption id="Mac">
-
-Example output:
-
-```
-Finding all available ports for the MotorBus.
-['/dev/tty.usbmodem575E0032081', '/dev/tty.usbmodem575E0031751']
-Remove the USB cable from your MotorsBus and press Enter when done.
-
-[...Disconnect corresponding leader or follower arm and press Enter...]
-
-The port of this MotorsBus is /dev/tty.usbmodem575E0032081
-Reconnect the USB cable.
-```
-
-Where the found port is: `/dev/tty.usbmodem575E0032081` corresponding to your leader or follower arm.
-
-</hfoption>
-<hfoption id="Linux">
-
-On Linux, you might need to give access to the USB ports by running:
-```bash
-sudo chmod 666 /dev/ttyACM0
-sudo chmod 666 /dev/ttyACM1
-```
-
-Example output:
-
-```
-Finding all available ports for the MotorBus.
-['/dev/ttyACM0', '/dev/ttyACM1']
-Remove the usb cable from your MotorsBus and press Enter when done.
-
-[...Disconnect corresponding leader or follower arm and press Enter...]
-
-The port of this MotorsBus is /dev/ttyACM1
-Reconnect the USB cable.
-```
-
-Where the found port is: `/dev/ttyACM1` corresponding to your leader or follower arm.
-
-</hfoption>
-</hfoptions>
-
-### 2. Set the motors ids and baudrates
-
-Each motor is identified by a unique id on the bus. When brand new, motors usually come with a default id of `1`. For the communication to work properly between the motors and the controller, we first need to set a unique, different id to each motor. Additionally, the speed at which data is transmitted on the bus is determined by the baudrate. In order to talk to each other, the controller and all the motors need to be configured with the same baudrate.
-
-To that end, we first need to connect to each motor individually with the controller in order to set these. Since we will write these parameters in the non-volatile section of the motors' internal memory (EEPROM), we'll only need to do this once.
-
-If you are repurposing motors from another robot, you will probably also need to perform this step as the ids and baudrate likely won't match.
-
-#### Follower
-
-Connect the usb cable from your computer and the power supply to the follower arm's controller board. Then, run the following command or run the API example with the port you got from the previous step. You'll also need to give your leader arm a name with the `id` parameter.
-
-For a visual reference on how to set the motor ids please refer to [this video](https://huggingface.co/docs/lerobot/en/so101#setup-motors-video) where we follow the process for the SO101 arm.
-
-<hfoptions id="setup_motors">
-<hfoption id="Command">
-
-```bash
-python -m lerobot.setup_motors \
-    --robot.type=so100_follower \
-    --robot.port=/dev/tty.usbmodem585A0076841  # <- paste here the port found at previous step
-```
-</hfoption>
-<hfoption id="API example">
-
-```python
-from lerobot.common.robots.so100_follower import SO100Follower, SO100FollowerConfig
-
-config = SO100FollowerConfig(
-    port="/dev/tty.usbmodem585A0076841",
-    id="my_awesome_follower_arm",
-)
-follower = SO100Follower(config)
-follower.setup_motors()
-```
-</hfoption>
-</hfoptions>
-
-You should see the following instruction
-```
-Connect the controller board to the 'gripper' motor only and press enter.
-```
-
-As instructed, plug the gripper's motor. Make sure it's the only motor connected to the board, and that the motor itself is not yet daisy-chained to any other motor. As you press `[Enter]`, the script will automatically set the id and baudrate for that motor.
-
-<details>
-<summary>Troubleshooting</summary>
-
-  If you get an error at that point, check your cables and make sure they are plugged in properly:
-  <ul>
-    <li>Power supply</li>
-    <li>USB cable between your computer and the controller board</li>
-    <li>The 3-pin cable from the controller board to the motor</li>
-  </ul>
-
-If you are using a Waveshare controller board, make sure that the two jumpers are set on the `B` channel (USB).
-</details>
-
-You should then see the following message:
-```
-'gripper' motor id set to 6
-```
-
-Followed by the next instruction:
-```
-Connect the controller board to the 'wrist_roll' motor only and press enter.
-```
-
-You can disconnect the 3-pin cable from the controller board, but you can leave it connected to the gripper motor on the other end, as it will already be in the right place. Now, plug in another 3-pin cable to the wrist roll motor and connect it to the controller board. As with the previous motor, make sure it is the only motor connected to the board and that the motor itself isn't connected to any other one.
-
-Repeat the operation for each motor as instructed.
-
-> [!TIP]
-> Check your cabling at each step before pressing Enter. For instance, the power supply cable might disconnect as you manipulate the board.
-
-When you are done, the script will simply finish, at which point the motors are ready to be used. You can now plug the 3-pin cable from each motor to the next one, and the cable from the first motor (the 'shoulder pan' with id=1) to the controller board, which can now be attached to the base of the arm.
-
-#### Leader
-Do the same steps for the leader arm.
-
-<hfoptions id="setup_motors">
-<hfoption id="Command">
-```bash
-python -m lerobot.setup_motors \
-    --teleop.type=so100_leader \
-    --teleop.port=/dev/tty.usbmodem575E0031751  # <- paste here the port found at previous step
-```
-</hfoption>
-<hfoption id="API example">
-
-```python
-from lerobot.common.teleoperators.so100_leader import SO100Leader, SO100LeaderConfig
-
-config = SO100LeaderConfig(
-    port="/dev/tty.usbmodem585A0076841",
-    id="my_awesome_leader_arm",
-)
-leader = SO100Leader(config)
-leader.setup_motors()
-```
-</hfoption>
-</hfoptions>
-
 ## Calibrate
 
 Next, you'll need to calibrate your robot to ensure that the leader and follower arms have the same position values when they are in the same physical position.
@@ -431,7 +434,7 @@ python -m lerobot.calibrate \
 <hfoption id="API example">
 
 ```python
-from lerobot.common.robots.so100_follower import SO100FollowerConfig, SO100Follower
+from lerobot.robots.so100_follower import SO100FollowerConfig, SO100Follower
 
 config = SO100FollowerConfig(
     port="/dev/tty.usbmodem585A0076891",
@@ -465,7 +468,7 @@ python -m lerobot.calibrate \
 <hfoption id="API example">
 
 ```python
-from lerobot.common.teleoperators.so100_leader import SO100LeaderConfig, SO100Leader
+from lerobot.teleoperators.so100_leader import SO100LeaderConfig, SO100Leader
 
 config = SO100LeaderConfig(
     port="/dev/tty.usbmodem58760431551",
diff --git a/lerobot/common/robots/so100_follower/so100_follower.py b/src/lerobot/robots/so100_follower/so100_follower.py
similarity index 96%
rename from lerobot/common/robots/so100_follower/so100_follower.py
rename to src/lerobot/robots/so100_follower/so100_follower.py
index 952049940..e5da6bc1a 100644
--- a/lerobot/common/robots/so100_follower/so100_follower.py
+++ b/src/lerobot/robots/so100_follower/so100_follower.py
@@ -19,10 +19,10 @@ import time
 from functools import cached_property
 from typing import Any
 
-from lerobot.common.cameras.utils import make_cameras_from_configs
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
-from lerobot.common.motors import Motor, MotorCalibration, MotorNormMode
-from lerobot.common.motors.feetech import (
+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,
 )
diff --git a/lerobot/common/robots/so100_follower/so100_follower_end_effector.py b/src/lerobot/robots/so100_follower/so100_follower_end_effector.py
similarity index 89%
rename from lerobot/common/robots/so100_follower/so100_follower_end_effector.py
rename to src/lerobot/robots/so100_follower/so100_follower_end_effector.py
index 82e89305b..5fe2993cb 100644
--- a/lerobot/common/robots/so100_follower/so100_follower_end_effector.py
+++ b/src/lerobot/robots/so100_follower/so100_follower_end_effector.py
@@ -20,17 +20,16 @@ from typing import Any
 
 import numpy as np
 
-from lerobot.common.cameras import make_cameras_from_configs
-from lerobot.common.errors import DeviceNotConnectedError
-from lerobot.common.model.kinematics import RobotKinematics
-from lerobot.common.motors import Motor, MotorNormMode
-from lerobot.common.motors.feetech import FeetechMotorsBus
+from lerobot.cameras import make_cameras_from_configs
+from lerobot.errors import DeviceNotConnectedError
+from lerobot.model.kinematics import RobotKinematics
+from lerobot.motors import Motor, MotorNormMode
+from lerobot.motors.feetech import FeetechMotorsBus
 
 from . import SO100Follower
 from .config_so100_follower import SO100FollowerEndEffectorConfig
 
 logger = logging.getLogger(__name__)
-EE_FRAME = "gripper_tip"
 
 
 class SO100FollowerEndEffector(SO100Follower):
@@ -64,7 +63,16 @@ class SO100FollowerEndEffector(SO100Follower):
         self.config = config
 
         # Initialize the kinematics module for the so100 robot
-        self.kinematics = RobotKinematics(robot_type="so_new_calibration")
+        if self.config.urdf_path is None:
+            raise ValueError(
+                "urdf_path must be provided in the configuration for end-effector control. "
+                "Please set urdf_path in your SO100FollowerEndEffectorConfig."
+            )
+
+        self.kinematics = RobotKinematics(
+            urdf_path=self.config.urdf_path,
+            target_frame_name=self.config.target_frame_name,
+        )
 
         # Store the bounds for end-effector position
         self.end_effector_bounds = self.config.end_effector_bounds
@@ -126,7 +134,7 @@ class SO100FollowerEndEffector(SO100Follower):
 
         # Calculate current end-effector position using forward kinematics
         if self.current_ee_pos is None:
-            self.current_ee_pos = self.kinematics.forward_kinematics(self.current_joint_pos, frame=EE_FRAME)
+            self.current_ee_pos = self.kinematics.forward_kinematics(self.current_joint_pos)
 
         # Set desired end-effector position by adding delta
         desired_ee_pos = np.eye(4)
@@ -142,11 +150,10 @@ class SO100FollowerEndEffector(SO100Follower):
             )
 
         # Compute inverse kinematics to get joint positions
-        target_joint_values_in_degrees = self.kinematics.ik(
-            self.current_joint_pos, desired_ee_pos, position_only=True, frame=EE_FRAME
+        target_joint_values_in_degrees = self.kinematics.inverse_kinematics(
+            self.current_joint_pos, desired_ee_pos
         )
 
-        target_joint_values_in_degrees = np.clip(target_joint_values_in_degrees, -180.0, 180.0)
         # Create joint space action dictionary
         joint_action = {
             f"{key}.pos": target_joint_values_in_degrees[i] for i, key in enumerate(self.bus.motors.keys())
diff --git a/lerobot/common/robots/so101_follower/__init__.py b/src/lerobot/robots/so101_follower/__init__.py
similarity index 100%
rename from lerobot/common/robots/so101_follower/__init__.py
rename to src/lerobot/robots/so101_follower/__init__.py
diff --git a/lerobot/common/robots/so101_follower/config_so101_follower.py b/src/lerobot/robots/so101_follower/config_so101_follower.py
similarity index 96%
rename from lerobot/common/robots/so101_follower/config_so101_follower.py
rename to src/lerobot/robots/so101_follower/config_so101_follower.py
index 6dbf21fd5..be630e6ac 100644
--- a/lerobot/common/robots/so101_follower/config_so101_follower.py
+++ b/src/lerobot/robots/so101_follower/config_so101_follower.py
@@ -16,7 +16,7 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.common.cameras import CameraConfig
+from lerobot.cameras import CameraConfig
 
 from ..config import RobotConfig
 
diff --git a/lerobot/common/robots/so101_follower/so101.mdx b/src/lerobot/robots/so101_follower/so101.mdx
similarity index 95%
rename from lerobot/common/robots/so101_follower/so101.mdx
rename to src/lerobot/robots/so101_follower/so101.mdx
index 5d39a1780..c49807d93 100644
--- a/lerobot/common/robots/so101_follower/so101.mdx
+++ b/src/lerobot/robots/so101_follower/so101.mdx
@@ -22,11 +22,11 @@ The follower arm uses 6x STS3215 motors with 1/345 gearing. The leader, however,
 
 | Leader-Arm Axis | Motor | Gear Ratio |
 |-----------------|:-------:|:----------:|
-| Base / Shoulder Yaw | 1 | 1 / 191 |
-| Shoulder Pitch      | 2 | 1 / 345 |
-| Elbow               | 3 | 1 / 191 |
-| Wrist Roll          | 4 | 1 / 147 |
-| Wrist Pitch         | 5 | 1 / 147 |
+| Base / Shoulder Pan | 1 | 1 / 191 |
+| Shoulder Lift       | 2 | 1 / 345 |
+| Elbow Flex          | 3 | 1 / 191 |
+| Wrist Flex          | 4 | 1 / 147 |
+| Wrist Roll          | 5 | 1 / 147 |
 | Gripper             | 6 | 1 / 147 |
 
 ### Clean Parts
@@ -136,7 +136,7 @@ Remove all support material from the 3D-printed parts. The easiest way to do thi
 
 To find the port for each bus servo adapter, run this script:
 ```bash
-python lerobot/find_port.py
+python -m lerobot.find_port
 ```
 
 <hfoptions id="example">
@@ -218,7 +218,7 @@ python -m lerobot.setup_motors \
 <hfoption id="API example">
 
 ```python
-from lerobot.common.robots.so101_follower import SO101Follower, SO101FollowerConfig
+from lerobot.robots.so101_follower import SO101Follower, SO101FollowerConfig
 
 config = SO101FollowerConfig(
     port="/dev/tty.usbmodem585A0076841",
@@ -284,7 +284,7 @@ python -m lerobot.setup_motors \
 <hfoption id="API example">
 
 ```python
-from lerobot.common.teleoperators.so101_leader import SO101Leader, SO101LeaderConfig
+from lerobot.teleoperators.so101_leader import SO101Leader, SO101LeaderConfig
 
 config = SO101LeaderConfig(
     port="/dev/tty.usbmodem585A0076841",
@@ -318,7 +318,7 @@ python -m lerobot.calibrate \
 <hfoption id="API example">
 
 ```python
-from lerobot.common.robots.so101_follower import SO101FollowerConfig, SO101Follower
+from lerobot.robots.so101_follower import SO101FollowerConfig, SO101Follower
 
 config = SO101FollowerConfig(
     port="/dev/tty.usbmodem585A0076891",
@@ -360,7 +360,7 @@ python -m lerobot.calibrate \
 <hfoption id="API example">
 
 ```python
-from lerobot.common.teleoperators.so101_leader import SO101LeaderConfig, SO101Leader
+from lerobot.teleoperators.so101_leader import SO101LeaderConfig, SO101Leader
 
 config = SO101LeaderConfig(
     port="/dev/tty.usbmodem58760431551",
diff --git a/lerobot/common/robots/so101_follower/so101_follower.py b/src/lerobot/robots/so101_follower/so101_follower.py
similarity index 96%
rename from lerobot/common/robots/so101_follower/so101_follower.py
rename to src/lerobot/robots/so101_follower/so101_follower.py
index a3c7aa0c2..3ae3c3967 100644
--- a/lerobot/common/robots/so101_follower/so101_follower.py
+++ b/src/lerobot/robots/so101_follower/so101_follower.py
@@ -19,10 +19,10 @@ import time
 from functools import cached_property
 from typing import Any
 
-from lerobot.common.cameras.utils import make_cameras_from_configs
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
-from lerobot.common.motors import Motor, MotorCalibration, MotorNormMode
-from lerobot.common.motors.feetech import (
+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,
 )
diff --git a/lerobot/common/robots/stretch3/README.md b/src/lerobot/robots/stretch3/README.md
similarity index 100%
rename from lerobot/common/robots/stretch3/README.md
rename to src/lerobot/robots/stretch3/README.md
diff --git a/lerobot/common/robots/stretch3/__init__.py b/src/lerobot/robots/stretch3/__init__.py
similarity index 100%
rename from lerobot/common/robots/stretch3/__init__.py
rename to src/lerobot/robots/stretch3/__init__.py
diff --git a/lerobot/common/robots/stretch3/configuration_stretch3.py b/src/lerobot/robots/stretch3/configuration_stretch3.py
similarity index 91%
rename from lerobot/common/robots/stretch3/configuration_stretch3.py
rename to src/lerobot/robots/stretch3/configuration_stretch3.py
index e62e4fa01..9fcf8f742 100644
--- a/lerobot/common/robots/stretch3/configuration_stretch3.py
+++ b/src/lerobot/robots/stretch3/configuration_stretch3.py
@@ -14,9 +14,9 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.common.cameras import CameraConfig
-from lerobot.common.cameras.opencv import OpenCVCameraConfig
-from lerobot.common.cameras.realsense import RealSenseCameraConfig
+from lerobot.cameras import CameraConfig
+from lerobot.cameras.opencv import OpenCVCameraConfig
+from lerobot.cameras.realsense import RealSenseCameraConfig
 
 from ..config import RobotConfig
 
diff --git a/lerobot/common/robots/stretch3/robot_stretch3.py b/src/lerobot/robots/stretch3/robot_stretch3.py
similarity index 96%
rename from lerobot/common/robots/stretch3/robot_stretch3.py
rename to src/lerobot/robots/stretch3/robot_stretch3.py
index 048db381f..b907d6a3f 100644
--- a/lerobot/common/robots/stretch3/robot_stretch3.py
+++ b/src/lerobot/robots/stretch3/robot_stretch3.py
@@ -21,9 +21,9 @@ from stretch_body.gamepad_teleop import GamePadTeleop
 from stretch_body.robot import Robot as StretchAPI
 from stretch_body.robot_params import RobotParams
 
-from lerobot.common.cameras.utils import make_cameras_from_configs
-from lerobot.common.constants import OBS_IMAGES, OBS_STATE
-from lerobot.common.datasets.utils import get_nested_item
+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 ..robot import Robot
 from .configuration_stretch3 import Stretch3RobotConfig
diff --git a/lerobot/common/robots/utils.py b/src/lerobot/robots/utils.py
similarity index 98%
rename from lerobot/common/robots/utils.py
rename to src/lerobot/robots/utils.py
index ccc1c58e8..435303c6e 100644
--- a/lerobot/common/robots/utils.py
+++ b/src/lerobot/robots/utils.py
@@ -15,7 +15,7 @@
 import logging
 from pprint import pformat
 
-from lerobot.common.robots import RobotConfig
+from lerobot.robots import RobotConfig
 
 from .robot import Robot
 
diff --git a/lerobot/common/robots/viperx/README.md b/src/lerobot/robots/viperx/README.md
similarity index 88%
rename from lerobot/common/robots/viperx/README.md
rename to src/lerobot/robots/viperx/README.md
index be2a323b6..445368e7a 100644
--- a/lerobot/common/robots/viperx/README.md
+++ b/src/lerobot/robots/viperx/README.md
@@ -58,7 +58,7 @@ python lerobot/scripts/control_robot.py \
   --control.type=teleoperate
 ```
 
-By adding `--robot.max_relative_target=5`, we override the default value for `max_relative_target` defined in [`AlohaRobotConfig`](lerobot/common/robot_devices/robots/configs.py). It is expected to be `5` to limit the magnitude of the movement for more safety, but the teleoperation won't be smooth. When you feel confident, you can disable this limit by adding `--robot.max_relative_target=null` to the command line:
+By adding `--robot.max_relative_target=5`, we override the default value for `max_relative_target` defined in [`AlohaRobotConfig`](lerobot/robot_devices/robots/configs.py). It is expected to be `5` to limit the magnitude of the movement for more safety, but the teleoperation won't be smooth. When you feel confident, you can disable this limit by adding `--robot.max_relative_target=null` to the command line:
 ```bash
 python lerobot/scripts/control_robot.py \
   --robot.type=aloha \
@@ -107,7 +107,7 @@ echo ${HF_USER}/aloha_test
 
 If you didn't upload with `--control.push_to_hub=false`, you can also visualize it locally with:
 ```bash
-python lerobot/scripts/visualize_dataset_html.py \
+python -m lerobot.scripts.visualize_dataset_html \
   --repo-id ${HF_USER}/aloha_test
 ```
 
@@ -129,9 +129,9 @@ python lerobot/scripts/control_robot.py \
 
 ## Train a policy
 
-To train a policy to control your robot, use the [`python lerobot/scripts/train.py`](../lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
+To train a policy to control your robot, use the [`python -m lerobot.scripts.train`](../src/lerobot/scripts/train.py) script. A few arguments are required. Here is an example command:
 ```bash
-python lerobot/scripts/train.py \
+python -m lerobot.scripts.train \
   --dataset.repo_id=${HF_USER}/aloha_test \
   --policy.type=act \
   --output_dir=outputs/train/act_aloha_test \
@@ -142,7 +142,7 @@ python lerobot/scripts/train.py \
 
 Let's explain it:
 1. We provided the dataset as argument with `--dataset.repo_id=${HF_USER}/aloha_test`.
-2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../lerobot/common/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor states, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
+2. We provided the policy with `policy.type=act`. This loads configurations from [`configuration_act.py`](../src/lerobot/policies/act/configuration_act.py). Importantly, this policy will automatically adapt to the number of motor states, motor actions and cameras of your robot (e.g. `laptop` and `phone`) which have been saved in your dataset.
 4. We provided `policy.device=cuda` since we are training on a Nvidia GPU, but you could use `policy.device=mps` to train on Apple silicon.
 5. We provided `wandb.enable=true` to use [Weights and Biases](https://docs.wandb.ai/quickstart) for visualizing training plots. This is optional but if you use it, make sure you are logged in by running `wandb login`.
 
@@ -152,7 +152,7 @@ Training should take several hours. You will find checkpoints in `outputs/train/
 
 ## Evaluate your policy
 
-You can use the `record` function from [`lerobot/scripts/control_robot.py`](../lerobot/scripts/control_robot.py) but with a policy checkpoint as input. For instance, run this command to record 10 evaluation episodes:
+You can use the `record` function from [`lerobot/scripts/control_robot.py`](../src/lerobot/scripts/control_robot.py) but with a policy checkpoint as input. For instance, run this command to record 10 evaluation episodes:
 ```bash
 python lerobot/scripts/control_robot.py \
   --robot.type=aloha \
diff --git a/lerobot/common/robots/viperx/__init__.py b/src/lerobot/robots/viperx/__init__.py
similarity index 100%
rename from lerobot/common/robots/viperx/__init__.py
rename to src/lerobot/robots/viperx/__init__.py
diff --git a/lerobot/common/robots/viperx/config_viperx.py b/src/lerobot/robots/viperx/config_viperx.py
similarity index 97%
rename from lerobot/common/robots/viperx/config_viperx.py
rename to src/lerobot/robots/viperx/config_viperx.py
index 6c7e2cc84..4922f1d18 100644
--- a/lerobot/common/robots/viperx/config_viperx.py
+++ b/src/lerobot/robots/viperx/config_viperx.py
@@ -14,7 +14,7 @@
 
 from dataclasses import dataclass, field
 
-from lerobot.common.cameras import CameraConfig
+from lerobot.cameras import CameraConfig
 
 from ..config import RobotConfig
 
diff --git a/lerobot/common/robots/viperx/viperx.py b/src/lerobot/robots/viperx/viperx.py
similarity index 96%
rename from lerobot/common/robots/viperx/viperx.py
rename to src/lerobot/robots/viperx/viperx.py
index 8ed8ef74c..881640cd5 100644
--- a/lerobot/common/robots/viperx/viperx.py
+++ b/src/lerobot/robots/viperx/viperx.py
@@ -17,11 +17,11 @@ import time
 from functools import cached_property
 from typing import Any
 
-from lerobot.common.cameras.utils import make_cameras_from_configs
-from lerobot.common.constants import OBS_STATE
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
-from lerobot.common.motors import Motor, MotorCalibration, MotorNormMode
-from lerobot.common.motors.dynamixel import (
+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,
 )
diff --git a/lerobot/scripts/display_sys_info.py b/src/lerobot/scripts/display_sys_info.py
similarity index 100%
rename from lerobot/scripts/display_sys_info.py
rename to src/lerobot/scripts/display_sys_info.py
diff --git a/lerobot/scripts/eval.py b/src/lerobot/scripts/eval.py
similarity index 97%
rename from lerobot/scripts/eval.py
rename to src/lerobot/scripts/eval.py
index 58275f666..d85ac27b3 100644
--- a/lerobot/scripts/eval.py
+++ b/src/lerobot/scripts/eval.py
@@ -21,7 +21,7 @@ You want to evaluate a model from the hub (eg: https://huggingface.co/lerobot/di
 for 10 episodes.
 
 ```
-python lerobot/scripts/eval.py \
+python -m lerobot.scripts.eval \
     --policy.path=lerobot/diffusion_pusht \
     --env.type=pusht \
     --eval.batch_size=10 \
@@ -32,7 +32,7 @@ python lerobot/scripts/eval.py \
 
 OR, you want to evaluate a model checkpoint from the LeRobot training script for 10 episodes.
 ```
-python lerobot/scripts/eval.py \
+python -m lerobot.scripts.eval \
     --policy.path=outputs/train/diffusion_pusht/checkpoints/005000/pretrained_model \
     --env.type=pusht \
     --eval.batch_size=10 \
@@ -65,20 +65,20 @@ from termcolor import colored
 from torch import Tensor, nn
 from tqdm import trange
 
-from lerobot.common.envs.factory import make_env
-from lerobot.common.envs.utils import add_envs_task, check_env_attributes_and_types, preprocess_observation
-from lerobot.common.policies.factory import make_policy
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.policies.utils import get_device_from_parameters
-from lerobot.common.utils.io_utils import write_video
-from lerobot.common.utils.random_utils import set_seed
-from lerobot.common.utils.utils import (
+from lerobot.configs import parser
+from lerobot.configs.eval import EvalPipelineConfig
+from lerobot.envs.factory import make_env
+from lerobot.envs.utils import add_envs_task, check_env_attributes_and_types, preprocess_observation
+from lerobot.policies.factory import make_policy
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.utils import get_device_from_parameters
+from lerobot.utils.io_utils import write_video
+from lerobot.utils.random_utils import set_seed
+from lerobot.utils.utils import (
     get_safe_torch_device,
     init_logging,
     inside_slurm,
 )
-from lerobot.configs import parser
-from lerobot.configs.eval import EvalPipelineConfig
 
 
 def rollout(
diff --git a/lerobot/scripts/find_joint_limits.py b/src/lerobot/scripts/find_joint_limits.py
similarity index 88%
rename from lerobot/scripts/find_joint_limits.py
rename to src/lerobot/scripts/find_joint_limits.py
index 95676dd35..f7e07514f 100644
--- a/lerobot/scripts/find_joint_limits.py
+++ b/src/lerobot/scripts/find_joint_limits.py
@@ -36,20 +36,21 @@ from dataclasses import dataclass
 import draccus
 import numpy as np
 
-from lerobot.common.model.kinematics import RobotKinematics
-from lerobot.common.robots import (  # noqa: F401
+from lerobot.model.kinematics import RobotKinematics
+from lerobot.robots import (  # noqa: F401
     RobotConfig,
     koch_follower,
     make_robot_from_config,
     so100_follower,
 )
-from lerobot.common.teleoperators import (  # noqa: F401
+from lerobot.teleoperators import (  # noqa: F401
     TeleoperatorConfig,
     gamepad,
     koch_leader,
     make_teleoperator_from_config,
     so100_leader,
 )
+from lerobot.utils.robot_utils import busy_wait
 
 
 @dataclass
@@ -76,12 +77,12 @@ def find_joint_and_ee_bounds(cfg: FindJointLimitsConfig):
         # Note to be compatible with the rest of the codebase,
         # we are using the new calibration method for so101 and so100
         robot_type = "so_new_calibration"
-    kinematics = RobotKinematics(robot_type=robot_type)
+    kinematics = RobotKinematics(cfg.robot.urdf_path, cfg.robot.target_frame_name)
 
     # Initialize min/max values
     observation = robot.get_observation()
     joint_positions = np.array([observation[f"{key}.pos"] for key in robot.bus.motors])
-    ee_pos = kinematics.forward_kinematics(joint_positions, frame="gripper_tip")[:3, 3]
+    ee_pos = kinematics.forward_kinematics(joint_positions)[:3, 3]
 
     max_pos = joint_positions.copy()
     min_pos = joint_positions.copy()
@@ -94,7 +95,7 @@ def find_joint_and_ee_bounds(cfg: FindJointLimitsConfig):
 
         observation = robot.get_observation()
         joint_positions = np.array([observation[f"{key}.pos"] for key in robot.bus.motors])
-        ee_pos = kinematics.forward_kinematics(joint_positions, frame="gripper_tip")[:3, 3]
+        ee_pos = kinematics.forward_kinematics(joint_positions)[:3, 3]
 
         # Skip initial warmup period
         if (time.perf_counter() - start_episode_t) < 5:
@@ -113,6 +114,8 @@ def find_joint_and_ee_bounds(cfg: FindJointLimitsConfig):
             print(f"Min joint pos position {np.round(min_pos, 4).tolist()}")
             break
 
+        busy_wait(0.01)
+
 
 if __name__ == "__main__":
     find_joint_and_ee_bounds()
diff --git a/lerobot/scripts/rl/actor.py b/src/lerobot/scripts/rl/actor.py
similarity index 96%
rename from lerobot/scripts/rl/actor.py
rename to src/lerobot/scripts/rl/actor.py
index da24d0dc5..0e96d3354 100644
--- a/lerobot/scripts/rl/actor.py
+++ b/src/lerobot/scripts/rl/actor.py
@@ -24,7 +24,7 @@ Examples of usage:
 
 - Start an actor server for real robot training with human-in-the-loop intervention:
 ```bash
-python lerobot/scripts/rl/actor.py --config_path lerobot/configs/train_config_hilserl_so100.json
+python -m lerobot.scripts.rl.actor --config_path src/lerobot/configs/train_config_hilserl_so100.json
 ```
 
 **NOTE**: The actor server requires a running learner server to connect to. Ensure the learner
@@ -57,37 +57,37 @@ import torch
 from torch import nn
 from torch.multiprocessing import Event, Queue
 
-from lerobot.common.cameras import opencv  # noqa: F401
-from lerobot.common.policies.factory import make_policy
-from lerobot.common.policies.sac.modeling_sac import SACPolicy
-from lerobot.common.robots import so100_follower  # noqa: F401
-from lerobot.common.teleoperators import gamepad, so101_leader  # noqa: F401
-from lerobot.common.transport import services_pb2, services_pb2_grpc
-from lerobot.common.transport.utils import (
+from lerobot.cameras import opencv  # noqa: F401
+from lerobot.configs import parser
+from lerobot.configs.train import TrainRLServerPipelineConfig
+from lerobot.policies.factory import make_policy
+from lerobot.policies.sac.modeling_sac import SACPolicy
+from lerobot.robots import so100_follower  # noqa: F401
+from lerobot.scripts.rl import learner_service
+from lerobot.scripts.rl.gym_manipulator import make_robot_env
+from lerobot.teleoperators import gamepad, so101_leader  # noqa: F401
+from lerobot.transport import services_pb2, services_pb2_grpc
+from lerobot.transport.utils import (
     bytes_to_state_dict,
     python_object_to_bytes,
     receive_bytes_in_chunks,
     send_bytes_in_chunks,
     transitions_to_bytes,
 )
-from lerobot.common.utils.process import ProcessSignalHandler
-from lerobot.common.utils.queue import get_last_item_from_queue
-from lerobot.common.utils.random_utils import set_seed
-from lerobot.common.utils.robot_utils import busy_wait
-from lerobot.common.utils.transition import (
+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 (
     Transition,
     move_state_dict_to_device,
     move_transition_to_device,
 )
-from lerobot.common.utils.utils import (
+from lerobot.utils.utils import (
     TimerManager,
     get_safe_torch_device,
     init_logging,
 )
-from lerobot.configs import parser
-from lerobot.configs.train import TrainRLServerPipelineConfig
-from lerobot.scripts.rl import learner_service
-from lerobot.scripts.rl.gym_manipulator import make_robot_env
 
 ACTOR_SHUTDOWN_TIMEOUT = 30
 
diff --git a/lerobot/scripts/rl/crop_dataset_roi.py b/src/lerobot/scripts/rl/crop_dataset_roi.py
similarity index 97%
rename from lerobot/scripts/rl/crop_dataset_roi.py
rename to src/lerobot/scripts/rl/crop_dataset_roi.py
index 5b7038de3..0b71b5363 100644
--- a/lerobot/scripts/rl/crop_dataset_roi.py
+++ b/src/lerobot/scripts/rl/crop_dataset_roi.py
@@ -21,12 +21,11 @@ from pathlib import Path
 from typing import Dict, Tuple
 
 import cv2
-
-# import torch.nn.functional as F  # noqa: N812
+import torch
 import torchvision.transforms.functional as F  # type: ignore  # noqa: N812
 from tqdm import tqdm  # type: ignore
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
 
 
 def select_rect_roi(img):
@@ -224,7 +223,8 @@ def convert_lerobot_dataset_to_cropper_lerobot_dataset(
                 cropped = F.crop(value, top, left, height, width)
                 value = F.resize(cropped, resize_size)
                 value = value.clamp(0, 1)
-
+            if key.startswith("complementary_info") and isinstance(value, torch.Tensor) and value.dim() == 0:
+                value = value.unsqueeze(0)
             new_frame[key] = value
 
         new_dataset.add_frame(new_frame, task=task)
@@ -265,8 +265,7 @@ if __name__ == "__main__":
     )
     parser.add_argument(
         "--push-to-hub",
-        type=bool,
-        default=False,
+        action="store_true",
         help="Whether to push the new dataset to the hub.",
     )
     parser.add_argument(
diff --git a/lerobot/scripts/rl/eval_policy.py b/src/lerobot/scripts/rl/eval_policy.py
similarity index 88%
rename from lerobot/scripts/rl/eval_policy.py
rename to src/lerobot/scripts/rl/eval_policy.py
index 3762719bf..aa97483b6 100644
--- a/lerobot/scripts/rl/eval_policy.py
+++ b/src/lerobot/scripts/rl/eval_policy.py
@@ -15,21 +15,21 @@
 # limitations under the License.
 import logging
 
-from lerobot.common.cameras import opencv  # noqa: F401
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.policies.factory import make_policy
-from lerobot.common.robots import (  # noqa: F401
+from lerobot.cameras import opencv  # noqa: F401
+from lerobot.configs import parser
+from lerobot.configs.train import TrainRLServerPipelineConfig
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.policies.factory import make_policy
+from lerobot.robots import (  # noqa: F401
     RobotConfig,
     make_robot_from_config,
     so100_follower,
 )
-from lerobot.common.teleoperators import (
+from lerobot.scripts.rl.gym_manipulator import make_robot_env
+from lerobot.teleoperators import (
     gamepad,  # noqa: F401
     so101_leader,  # noqa: F401
 )
-from lerobot.configs import parser
-from lerobot.configs.train import TrainRLServerPipelineConfig
-from lerobot.scripts.rl.gym_manipulator import make_robot_env
 
 logging.basicConfig(level=logging.INFO)
 
diff --git a/lerobot/scripts/rl/gym_manipulator.py b/src/lerobot/scripts/rl/gym_manipulator.py
similarity index 96%
rename from lerobot/scripts/rl/gym_manipulator.py
rename to src/lerobot/scripts/rl/gym_manipulator.py
index e7327d96d..673043b6e 100644
--- a/lerobot/scripts/rl/gym_manipulator.py
+++ b/src/lerobot/scripts/rl/gym_manipulator.py
@@ -47,26 +47,26 @@ import numpy as np
 import torch
 import torchvision.transforms.functional as F  # noqa: N812
 
-from lerobot.common.cameras import opencv  # noqa: F401
-from lerobot.common.envs.configs import EnvConfig
-from lerobot.common.envs.utils import preprocess_observation
-from lerobot.common.model.kinematics import RobotKinematics
-from lerobot.common.robots import (  # noqa: F401
+from lerobot.cameras import opencv  # noqa: F401
+from lerobot.configs import parser
+from lerobot.envs.configs import EnvConfig
+from lerobot.envs.utils import preprocess_observation
+from lerobot.model.kinematics import RobotKinematics
+from lerobot.robots import (  # noqa: F401
     RobotConfig,
     make_robot_from_config,
     so100_follower,
 )
-from lerobot.common.teleoperators import (
+from lerobot.teleoperators import (
     gamepad,  # noqa: F401
     keyboard,  # noqa: F401
     make_teleoperator_from_config,
     so101_leader,  # noqa: F401
 )
-from lerobot.common.teleoperators.gamepad.teleop_gamepad import GamepadTeleop
-from lerobot.common.teleoperators.keyboard.teleop_keyboard import KeyboardEndEffectorTeleop
-from lerobot.common.utils.robot_utils import busy_wait
-from lerobot.common.utils.utils import log_say
-from lerobot.configs import parser
+from lerobot.teleoperators.gamepad.teleop_gamepad import GamepadTeleop
+from lerobot.teleoperators.keyboard.teleop_keyboard import KeyboardEndEffectorTeleop
+from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.utils import log_say
 
 logging.basicConfig(level=logging.INFO)
 
@@ -254,20 +254,19 @@ class RobotEnv(gym.Env):
         self._joint_names = [f"{key}.pos" for key in self.robot.bus.motors]
         self._image_keys = self.robot.cameras.keys()
 
-        # Read initial joint positions using the bus
-        self.current_joint_positions = self._get_observation()["agent_pos"]
+        self.current_observation = None
 
         self.use_gripper = use_gripper
 
         self._setup_spaces()
 
-    def _get_observation(self) -> np.ndarray:
+    def _get_observation(self) -> dict[str, np.ndarray]:
         """Helper to convert a dictionary from bus.sync_read to an ordered numpy array."""
         obs_dict = self.robot.get_observation()
-        joint_positions = np.array([obs_dict[name] for name in self._joint_names], dtype=np.float32)
+        joint_positions = np.array([obs_dict[name] for name in self._joint_names])
 
         images = {key: obs_dict[key] for key in self._image_keys}
-        return {"agent_pos": joint_positions, "pixels": images}
+        self.current_observation = {"agent_pos": joint_positions, "pixels": images}
 
     def _setup_spaces(self):
         """
@@ -281,24 +280,24 @@ class RobotEnv(gym.Env):
             - The action space is defined as a Box space representing joint position commands. It is defined as relative (delta)
               or absolute, based on the configuration.
         """
-        example_obs = self._get_observation()
+        self._get_observation()
 
         observation_spaces = {}
 
         # Define observation spaces for images and other states.
-        if "pixels" in example_obs:
-            prefix = "observation.images" if len(example_obs["pixels"]) > 1 else "observation.image"
+        if "pixels" in self.current_observation:
+            prefix = "observation.images"
             observation_spaces = {
                 f"{prefix}.{key}": gym.spaces.Box(
-                    low=0, high=255, shape=example_obs["pixels"][key].shape, dtype=np.uint8
+                    low=0, high=255, shape=self.current_observation["pixels"][key].shape, dtype=np.uint8
                 )
-                for key in example_obs["pixels"]
+                for key in self.current_observation["pixels"]
             }
 
         observation_spaces["observation.state"] = gym.spaces.Box(
             low=0,
             high=10,
-            shape=example_obs["agent_pos"].shape,
+            shape=self.current_observation["agent_pos"].shape,
             dtype=np.float32,
         )
 
@@ -340,14 +339,12 @@ class RobotEnv(gym.Env):
 
         self.robot.reset()
 
-        # Capture the initial observation.
-        observation = self._get_observation()
-
         # Reset episode tracking variables.
         self.current_step = 0
         self.episode_data = None
-
-        return observation, {"is_intervention": False}
+        self.current_observation = None
+        self._get_observation()
+        return self.current_observation, {"is_intervention": False}
 
     def step(self, action) -> tuple[dict[str, np.ndarray], float, bool, bool, dict[str, Any]]:
         """
@@ -367,8 +364,6 @@ class RobotEnv(gym.Env):
                 - truncated (bool): True if the episode was truncated (e.g., time constraints).
                 - info (dict): Additional debugging information including intervention status.
         """
-        self.current_joint_positions = self._get_observation()["agent_pos"]
-
         action_dict = {"delta_x": action[0], "delta_y": action[1], "delta_z": action[2]}
 
         # 1.0 action corresponds to no-op action
@@ -376,6 +371,8 @@ class RobotEnv(gym.Env):
 
         self.robot.send_action(action_dict)
 
+        self._get_observation()
+
         if self.display_cameras:
             self.render()
 
@@ -386,7 +383,7 @@ class RobotEnv(gym.Env):
         truncated = False
 
         return (
-            self._get_observation(),
+            self.current_observation,
             reward,
             terminated,
             truncated,
@@ -399,11 +396,10 @@ class RobotEnv(gym.Env):
         """
         import cv2
 
-        observation = self._get_observation()
-        image_keys = [key for key in observation if "image" in key]
+        image_keys = [key for key in self.current_observation if "image" in key]
 
         for key in image_keys:
-            cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
+            cv2.imshow(key, cv2.cvtColor(self.current_observation[key].numpy(), cv2.COLOR_RGB2BGR))
             cv2.waitKey(1)
 
     def close(self):
@@ -520,7 +516,10 @@ class AddCurrentToObservation(gym.ObservationWrapper):
         Returns:
             The modified observation with current values.
         """
-        present_current_observation = self.unwrapped._get_observation()["agent_pos"]
+        present_current_dict = self.env.unwrapped.robot.bus.sync_read("Present_Current")
+        present_current_observation = np.array(
+            [present_current_dict[name] for name in self.env.unwrapped.robot.bus.motors]
+        )
         observation["agent_pos"] = np.concatenate(
             [observation["agent_pos"], present_current_observation], axis=-1
         )
@@ -1090,13 +1089,10 @@ class EEObservationWrapper(gym.ObservationWrapper):
             dtype=np.float32,
         )
 
-        # Initialize kinematics instance for the appropriate robot type
-        robot_type = getattr(env.unwrapped.robot.config, "robot_type", "so101")
-        if "so100" in robot_type or "so101" in robot_type:
-            # Note to be compatible with the rest of the codebase,
-            # we are using the new calibration method for so101 and so100
-            robot_type = "so_new_calibration"
-        self.kinematics = RobotKinematics(robot_type)
+        self.kinematics = RobotKinematics(
+            urdf_path=env.unwrapped.robot.config.urdf_path,
+            target_frame_name=env.unwrapped.robot.config.target_frame_name,
+        )
 
     def observation(self, observation):
         """
@@ -1108,9 +1104,9 @@ class EEObservationWrapper(gym.ObservationWrapper):
         Returns:
             Enhanced observation with end-effector pose information.
         """
-        current_joint_pos = self.unwrapped._get_observation()["agent_pos"]
+        current_joint_pos = self.unwrapped.current_observation["agent_pos"]
 
-        current_ee_pos = self.kinematics.forward_kinematics(current_joint_pos, frame="gripper_tip")[:3, 3]
+        current_ee_pos = self.kinematics.forward_kinematics(current_joint_pos)[:3, 3]
         observation["agent_pos"] = np.concatenate([observation["agent_pos"], current_ee_pos], -1)
         return observation
 
@@ -1157,12 +1153,10 @@ class BaseLeaderControlWrapper(gym.Wrapper):
         self.event_lock = Lock()  # Thread-safe access to events
 
         # Initialize robot control
-        robot_type = getattr(env.unwrapped.robot.config, "robot_type", "so101")
-        if "so100" in robot_type or "so101" in robot_type:
-            # Note to be compatible with the rest of the codebase,
-            # we are using the new calibration method for so101 and so100
-            robot_type = "so_new_calibration"
-        self.kinematics = RobotKinematics(robot_type)
+        self.kinematics = RobotKinematics(
+            urdf_path=env.unwrapped.robot.config.urdf_path,
+            target_frame_name=env.unwrapped.robot.config.target_frame_name,
+        )
         self.leader_torque_enabled = True
         self.prev_leader_gripper = None
 
@@ -1260,14 +1254,14 @@ class BaseLeaderControlWrapper(gym.Wrapper):
         leader_pos_dict = self.robot_leader.bus.sync_read("Present_Position")
         follower_pos_dict = self.robot_follower.bus.sync_read("Present_Position")
 
-        leader_pos = np.array([leader_pos_dict[name] for name in leader_pos_dict], dtype=np.float32)
-        follower_pos = np.array([follower_pos_dict[name] for name in follower_pos_dict], dtype=np.float32)
+        leader_pos = np.array([leader_pos_dict[name] for name in leader_pos_dict])
+        follower_pos = np.array([follower_pos_dict[name] for name in follower_pos_dict])
 
         self.leader_tracking_error_queue.append(np.linalg.norm(follower_pos[:-1] - leader_pos[:-1]))
 
         # [:3, 3] Last column of the transformation matrix corresponds to the xyz translation
-        leader_ee = self.kinematics.forward_kinematics(leader_pos, frame="gripper_tip")[:3, 3]
-        follower_ee = self.kinematics.forward_kinematics(follower_pos, frame="gripper_tip")[:3, 3]
+        leader_ee = self.kinematics.forward_kinematics(leader_pos)[:3, 3]
+        follower_ee = self.kinematics.forward_kinematics(follower_pos)[:3, 3]
 
         action = np.clip(leader_ee - follower_ee, -self.end_effector_step_sizes, self.end_effector_step_sizes)
         # Normalize the action to the range [-1, 1]
@@ -1341,6 +1335,9 @@ class BaseLeaderControlWrapper(gym.Wrapper):
         # NOTE:
         obs, reward, terminated, truncated, info = self.env.step(action)
 
+        if isinstance(action, np.ndarray):
+            action = torch.from_numpy(action)
+
         # Add intervention info
         info["is_intervention"] = is_intervention
         info["action_intervention"] = action
@@ -1877,7 +1874,6 @@ def make_robot_env(cfg: EnvConfig) -> gym.Env:
     if cfg.robot is None:
         raise ValueError("RobotConfig (cfg.robot) must be provided for gym_manipulator environment.")
     robot = make_robot_from_config(cfg.robot)
-
     teleop_device = make_teleoperator_from_config(cfg.teleop)
     teleop_device.connect()
 
@@ -1980,7 +1976,7 @@ def init_reward_classifier(cfg):
     if cfg.reward_classifier_pretrained_path is None:
         return None
 
-    from lerobot.common.policies.sac.reward_model.modeling_classifier import Classifier
+    from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
 
     # Get device from config or default to CUDA
     device = getattr(cfg, "device", "cpu")
@@ -2023,7 +2019,7 @@ def record_dataset(env, policy, cfg):
                                   a success (reward=1) is detected. This helps collect
                                   more positive examples for reward classifier training.
     """
-    from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+    from lerobot.datasets.lerobot_dataset import LeRobotDataset
 
     # Setup initial action (zero action if using teleop)
     action = env.action_space.sample() * 0.0
@@ -2177,7 +2173,7 @@ def replay_episode(env, cfg):
             - dataset_root: Local root directory for dataset
             - episode: Episode ID to replay
     """
-    from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+    from lerobot.datasets.lerobot_dataset import LeRobotDataset
 
     dataset = LeRobotDataset(cfg.repo_id, root=cfg.dataset_root, episodes=[cfg.episode])
     env.reset()
@@ -2210,7 +2206,7 @@ def main(cfg: EnvConfig):
     if cfg.mode == "record":
         policy = None
         if cfg.pretrained_policy_name_or_path is not None:
-            from lerobot.common.policies.sac.modeling_sac import SACPolicy
+            from lerobot.policies.sac.modeling_sac import SACPolicy
 
             policy = SACPolicy.from_pretrained(cfg.pretrained_policy_name_or_path)
             policy.to(cfg.device)
diff --git a/lerobot/scripts/rl/learner.py b/src/lerobot/scripts/rl/learner.py
similarity index 97%
rename from lerobot/scripts/rl/learner.py
rename to src/lerobot/scripts/rl/learner.py
index 663dbe918..d8830d83e 100644
--- a/lerobot/scripts/rl/learner.py
+++ b/src/lerobot/scripts/rl/learner.py
@@ -25,7 +25,7 @@ Examples of usage:
 
 - Start a learner server for training:
 ```bash
-python lerobot/scripts/rl/learner.py --config_path lerobot/configs/train_config_hilserl_so100.json
+python -m lerobot.scripts.rl.learner --config_path src/lerobot/configs/train_config_hilserl_so100.json
 ```
 
 **NOTE**: Start the learner server before launching the actor server. The learner opens a gRPC server
@@ -59,46 +59,46 @@ from torch import nn
 from torch.multiprocessing import Queue
 from torch.optim.optimizer import Optimizer
 
-from lerobot.common.cameras import opencv  # noqa: F401
-from lerobot.common.constants import (
+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.common.datasets.factory import make_dataset
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.policies.factory import make_policy
-from lerobot.common.policies.sac.modeling_sac import SACPolicy
-from lerobot.common.robots import so100_follower  # noqa: F401
-from lerobot.common.teleoperators import gamepad, so101_leader  # noqa: F401
-from lerobot.common.transport import services_pb2_grpc
-from lerobot.common.transport.utils import (
+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.robots import so100_follower  # noqa: F401
+from lerobot.scripts.rl import learner_service
+from lerobot.teleoperators import gamepad, so101_leader  # noqa: F401
+from lerobot.transport import services_pb2_grpc
+from lerobot.transport.utils import (
     bytes_to_python_object,
     bytes_to_transitions,
     state_to_bytes,
 )
-from lerobot.common.utils.buffer import ReplayBuffer, concatenate_batch_transitions
-from lerobot.common.utils.process import ProcessSignalHandler
-from lerobot.common.utils.random_utils import set_seed
-from lerobot.common.utils.train_utils import (
+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,
     save_checkpoint,
     update_last_checkpoint,
 )
-from lerobot.common.utils.train_utils import (
+from lerobot.utils.train_utils import (
     load_training_state as utils_load_training_state,
 )
-from lerobot.common.utils.transition import move_state_dict_to_device, move_transition_to_device
-from lerobot.common.utils.utils import (
+from lerobot.utils.transition import move_state_dict_to_device, move_transition_to_device
+from lerobot.utils.utils import (
     format_big_number,
     get_safe_torch_device,
     init_logging,
 )
-from lerobot.common.utils.wandb_utils import WandBLogger
-from lerobot.configs import parser
-from lerobot.configs.train import TrainRLServerPipelineConfig
-from lerobot.scripts.rl import learner_service
+from lerobot.utils.wandb_utils import WandBLogger
 
 LOG_PREFIX = "[LEARNER]"
 
@@ -157,7 +157,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.common.utils.wandb_utils import WandBLogger
+        from lerobot.utils.wandb_utils import WandBLogger
 
         wandb_logger = WandBLogger(cfg)
     else:
diff --git a/lerobot/scripts/rl/learner_service.py b/src/lerobot/scripts/rl/learner_service.py
similarity index 94%
rename from lerobot/scripts/rl/learner_service.py
rename to src/lerobot/scripts/rl/learner_service.py
index f967d812c..198e52945 100644
--- a/lerobot/scripts/rl/learner_service.py
+++ b/src/lerobot/scripts/rl/learner_service.py
@@ -19,9 +19,9 @@ import logging
 import time
 from multiprocessing import Event, Queue
 
-from lerobot.common.transport import services_pb2, services_pb2_grpc
-from lerobot.common.transport.utils import receive_bytes_in_chunks, send_bytes_in_chunks
-from lerobot.common.utils.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_MESSAGE_SIZE = 4 * 1024 * 1024  # 4 MB
 MAX_WORKERS = 3  # Stream parameters, send transitions and interactions
diff --git a/lerobot/scripts/train.py b/src/lerobot/scripts/train.py
similarity index 92%
rename from lerobot/scripts/train.py
rename to src/lerobot/scripts/train.py
index 0de247be9..2f2e88de6 100644
--- a/lerobot/scripts/train.py
+++ b/src/lerobot/scripts/train.py
@@ -24,33 +24,33 @@ from termcolor import colored
 from torch.amp import GradScaler
 from torch.optim import Optimizer
 
-from lerobot.common.datasets.factory import make_dataset
-from lerobot.common.datasets.sampler import EpisodeAwareSampler
-from lerobot.common.datasets.utils import cycle
-from lerobot.common.envs.factory import make_env
-from lerobot.common.optim.factory import make_optimizer_and_scheduler
-from lerobot.common.policies.factory import make_policy
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.policies.utils import get_device_from_parameters
-from lerobot.common.utils.logging_utils import AverageMeter, MetricsTracker
-from lerobot.common.utils.random_utils import set_seed
-from lerobot.common.utils.train_utils import (
+from lerobot.configs import parser
+from lerobot.configs.train import TrainPipelineConfig
+from lerobot.datasets.factory import make_dataset
+from lerobot.datasets.sampler import EpisodeAwareSampler
+from lerobot.datasets.utils import cycle
+from lerobot.envs.factory import make_env
+from lerobot.optim.factory import make_optimizer_and_scheduler
+from lerobot.policies.factory import make_policy
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.utils import get_device_from_parameters
+from lerobot.scripts.eval import eval_policy
+from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
+from lerobot.utils.random_utils import set_seed
+from lerobot.utils.train_utils import (
     get_step_checkpoint_dir,
     get_step_identifier,
     load_training_state,
     save_checkpoint,
     update_last_checkpoint,
 )
-from lerobot.common.utils.utils import (
+from lerobot.utils.utils import (
     format_big_number,
     get_safe_torch_device,
     has_method,
     init_logging,
 )
-from lerobot.common.utils.wandb_utils import WandBLogger
-from lerobot.configs import parser
-from lerobot.configs.train import TrainPipelineConfig
-from lerobot.scripts.eval import eval_policy
+from lerobot.utils.wandb_utils import WandBLogger
 
 
 def update_policy(
@@ -282,6 +282,9 @@ def train(cfg: TrainPipelineConfig):
         eval_env.close()
     logging.info("End of training")
 
+    if cfg.policy.push_to_hub:
+        policy.push_model_to_hub(cfg)
+
 
 if __name__ == "__main__":
     init_logging()
diff --git a/lerobot/scripts/visualize_dataset.py b/src/lerobot/scripts/visualize_dataset.py
similarity index 97%
rename from lerobot/scripts/visualize_dataset.py
rename to src/lerobot/scripts/visualize_dataset.py
index cdfea6b8b..37db66ddf 100644
--- a/lerobot/scripts/visualize_dataset.py
+++ b/src/lerobot/scripts/visualize_dataset.py
@@ -29,14 +29,14 @@ Examples:
 
 - Visualize data stored on a local machine:
 ```
-local$ python lerobot/scripts/visualize_dataset.py \
+local$ python -m lerobot.scripts.visualize_dataset \
     --repo-id lerobot/pusht \
     --episode-index 0
 ```
 
 - Visualize data stored on a distant machine with a local viewer:
 ```
-distant$ python lerobot/scripts/visualize_dataset.py \
+distant$ python -m lerobot.scripts.visualize_dataset \
     --repo-id lerobot/pusht \
     --episode-index 0 \
     --save 1 \
@@ -50,7 +50,7 @@ local$ rerun lerobot_pusht_episode_0.rrd
 (You need to forward the websocket port to the distant machine, with
 `ssh -L 9087:localhost:9087 username@remote-host`)
 ```
-distant$ python lerobot/scripts/visualize_dataset.py \
+distant$ python -m lerobot.scripts.visualize_dataset \
     --repo-id lerobot/pusht \
     --episode-index 0 \
     --mode distant \
@@ -74,7 +74,7 @@ import torch
 import torch.utils.data
 import tqdm
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
 
 
 class EpisodeSampler(torch.utils.data.Sampler):
diff --git a/lerobot/scripts/visualize_dataset_html.py b/src/lerobot/scripts/visualize_dataset_html.py
similarity index 97%
rename from lerobot/scripts/visualize_dataset_html.py
rename to src/lerobot/scripts/visualize_dataset_html.py
index d0c8f1ace..a722da603 100644
--- a/lerobot/scripts/visualize_dataset_html.py
+++ b/src/lerobot/scripts/visualize_dataset_html.py
@@ -29,7 +29,7 @@ Example of usage:
 
 - Visualize data stored on a local machine:
 ```bash
-local$ python lerobot/scripts/visualize_dataset_html.py \
+local$ python -m lerobot.scripts.visualize_dataset_html \
     --repo-id lerobot/pusht
 
 local$ open http://localhost:9090
@@ -37,7 +37,7 @@ local$ open http://localhost:9090
 
 - Visualize data stored on a distant machine with a local viewer:
 ```bash
-distant$ python lerobot/scripts/visualize_dataset_html.py \
+distant$ python -m lerobot.scripts.visualize_dataset_html \
     --repo-id lerobot/pusht
 
 local$ ssh -L 9090:localhost:9090 distant  # create a ssh tunnel
@@ -46,7 +46,7 @@ local$ open http://localhost:9090
 
 - Select episodes to visualize:
 ```bash
-python lerobot/scripts/visualize_dataset_html.py \
+python -m lerobot.scripts.visualize_dataset_html \
     --repo-id lerobot/pusht \
     --episodes 7 3 5 1 4
 ```
@@ -68,9 +68,9 @@ import requests
 from flask import Flask, redirect, render_template, request, url_for
 
 from lerobot import available_datasets
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.datasets.utils import IterableNamespace
-from lerobot.common.utils.utils import init_logging
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import IterableNamespace
+from lerobot.utils.utils import init_logging
 
 
 def run_server(
diff --git a/lerobot/scripts/visualize_image_transforms.py b/src/lerobot/scripts/visualize_image_transforms.py
similarity index 96%
rename from lerobot/scripts/visualize_image_transforms.py
rename to src/lerobot/scripts/visualize_image_transforms.py
index 80935d327..14caf89df 100644
--- a/lerobot/scripts/visualize_image_transforms.py
+++ b/src/lerobot/scripts/visualize_image_transforms.py
@@ -20,7 +20,7 @@ Additionally, each individual transform can be visualized separately as well as
 
 Example:
 ```bash
-python lerobot/scripts/visualize_image_transforms.py \
+python -m lerobot.scripts.visualize_image_transforms \
     --repo_id=lerobot/pusht \
     --episodes='[0]' \
     --image_transforms.enable=True
@@ -35,13 +35,13 @@ from pathlib import Path
 import draccus
 from torchvision.transforms import ToPILImage
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.datasets.transforms import (
+from lerobot.configs.default import DatasetConfig
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.transforms import (
     ImageTransforms,
     ImageTransformsConfig,
     make_transform_from_config,
 )
-from lerobot.configs.default import DatasetConfig
 
 OUTPUT_DIR = Path("outputs/image_transforms")
 to_pil = ToPILImage()
diff --git a/lerobot/setup_motors.py b/src/lerobot/setup_motors.py
similarity index 95%
rename from lerobot/setup_motors.py
rename to src/lerobot/setup_motors.py
index 7909dc68d..c54582a1d 100644
--- a/lerobot/setup_motors.py
+++ b/src/lerobot/setup_motors.py
@@ -28,7 +28,7 @@ from dataclasses import dataclass
 
 import draccus
 
-from .common.robots import (  # noqa: F401
+from lerobot.robots import (  # noqa: F401
     RobotConfig,
     koch_follower,
     lekiwi,
@@ -36,7 +36,7 @@ from .common.robots import (  # noqa: F401
     so100_follower,
     so101_follower,
 )
-from .common.teleoperators import (  # noqa: F401
+from lerobot.teleoperators import (  # noqa: F401
     TeleoperatorConfig,
     koch_leader,
     make_teleoperator_from_config,
diff --git a/lerobot/teleoperate.py b/src/lerobot/teleoperate.py
similarity index 76%
rename from lerobot/teleoperate.py
rename to src/lerobot/teleoperate.py
index 6080dfb40..e2819345b 100644
--- a/lerobot/teleoperate.py
+++ b/src/lerobot/teleoperate.py
@@ -36,12 +36,11 @@ from dataclasses import asdict, dataclass
 from pprint import pformat
 
 import draccus
-import numpy as np
 import rerun as rr
 
-from lerobot.common.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
-from lerobot.common.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
-from lerobot.common.robots import (  # noqa: F401
+from lerobot.cameras.opencv.configuration_opencv import OpenCVCameraConfig  # noqa: F401
+from lerobot.cameras.realsense.configuration_realsense import RealSenseCameraConfig  # noqa: F401
+from lerobot.robots import (  # noqa: F401
     Robot,
     RobotConfig,
     koch_follower,
@@ -49,20 +48,23 @@ from lerobot.common.robots import (  # noqa: F401
     so100_follower,
     so101_follower,
 )
-from lerobot.common.teleoperators import (
+from lerobot.teleoperators import (  # noqa: F401
     Teleoperator,
     TeleoperatorConfig,
+    gamepad,
+    koch_leader,
     make_teleoperator_from_config,
+    so100_leader,
+    so101_leader,
 )
-from lerobot.common.utils.robot_utils import busy_wait
-from lerobot.common.utils.utils import init_logging, move_cursor_up
-from lerobot.common.utils.visualization_utils import _init_rerun
-
-from .common.teleoperators import gamepad, koch_leader, so100_leader, so101_leader  # noqa: F401
+from lerobot.utils.robot_utils import busy_wait
+from lerobot.utils.utils import init_logging, move_cursor_up
+from lerobot.utils.visualization_utils import _init_rerun, log_rerun_data
 
 
 @dataclass
 class TeleoperateConfig:
+    # TODO: pepijn, steven: if more robots require multiple teleoperators (like lekiwi) its good to make this possibele in teleop.py and record.py with List[Teleoperator]
     teleop: TeleoperatorConfig
     robot: RobotConfig
     # Limit the maximum frames per second.
@@ -82,14 +84,7 @@ def teleop_loop(
         action = teleop.get_action()
         if display_data:
             observation = robot.get_observation()
-            for obs, val in observation.items():
-                if isinstance(val, float):
-                    rr.log(f"observation_{obs}", rr.Scalar(val))
-                elif isinstance(val, np.ndarray):
-                    rr.log(f"observation_{obs}", rr.Image(val), static=True)
-            for act, val in action.items():
-                if isinstance(val, float):
-                    rr.log(f"action_{act}", rr.Scalar(val))
+            log_rerun_data(observation, action)
 
         robot.send_action(action)
         dt_s = time.perf_counter() - loop_start
diff --git a/lerobot/common/teleoperators/__init__.py b/src/lerobot/teleoperators/__init__.py
similarity index 100%
rename from lerobot/common/teleoperators/__init__.py
rename to src/lerobot/teleoperators/__init__.py
diff --git a/lerobot/common/teleoperators/config.py b/src/lerobot/teleoperators/config.py
similarity index 100%
rename from lerobot/common/teleoperators/config.py
rename to src/lerobot/teleoperators/config.py
diff --git a/lerobot/common/teleoperators/gamepad/__init__.py b/src/lerobot/teleoperators/gamepad/__init__.py
similarity index 100%
rename from lerobot/common/teleoperators/gamepad/__init__.py
rename to src/lerobot/teleoperators/gamepad/__init__.py
diff --git a/lerobot/common/teleoperators/gamepad/configuration_gamepad.py b/src/lerobot/teleoperators/gamepad/configuration_gamepad.py
similarity index 100%
rename from lerobot/common/teleoperators/gamepad/configuration_gamepad.py
rename to src/lerobot/teleoperators/gamepad/configuration_gamepad.py
diff --git a/lerobot/common/teleoperators/gamepad/gamepad_utils.py b/src/lerobot/teleoperators/gamepad/gamepad_utils.py
similarity index 97%
rename from lerobot/common/teleoperators/gamepad/gamepad_utils.py
rename to src/lerobot/teleoperators/gamepad/gamepad_utils.py
index 21a293c77..9b62dc666 100644
--- a/lerobot/common/teleoperators/gamepad/gamepad_utils.py
+++ b/src/lerobot/teleoperators/gamepad/gamepad_utils.py
@@ -295,8 +295,8 @@ class GamepadController(InputController):
         try:
             # Read joystick axes
             # Left stick X and Y (typically axes 0 and 1)
-            x_input = self.joystick.get_axis(0)  # Left/Right
-            y_input = self.joystick.get_axis(1)  # Up/Down (often inverted)
+            y_input = self.joystick.get_axis(0)  # Left/Right
+            x_input = self.joystick.get_axis(1)  # Up/Down (often inverted)
 
             # Right stick Y (typically axis 3 or 4)
             z_input = self.joystick.get_axis(3)  # Up/Down for Z
@@ -307,8 +307,8 @@ class GamepadController(InputController):
             z_input = 0 if abs(z_input) < self.deadzone else z_input
 
             # Calculate deltas (note: may need to invert axes depending on controller)
-            delta_x = -y_input * self.y_step_size  # Forward/backward
-            delta_y = -x_input * self.x_step_size  # Left/right
+            delta_x = -x_input * self.x_step_size  # Forward/backward
+            delta_y = -y_input * self.y_step_size  # Left/right
             delta_z = -z_input * self.z_step_size  # Up/down
 
             return delta_x, delta_y, delta_z
@@ -424,14 +424,14 @@ class GamepadControllerHID(InputController):
             # These offsets are for the Logitech RumblePad 2
             if data and len(data) >= 8:
                 # Normalize joystick values from 0-255 to -1.0-1.0
-                self.left_x = (data[1] - 128) / 128.0
-                self.left_y = (data[2] - 128) / 128.0
+                self.left_y = (data[1] - 128) / 128.0
+                self.left_x = (data[2] - 128) / 128.0
                 self.right_x = (data[3] - 128) / 128.0
                 self.right_y = (data[4] - 128) / 128.0
 
                 # Apply deadzone
-                self.left_x = 0 if abs(self.left_x) < self.deadzone else self.left_x
                 self.left_y = 0 if abs(self.left_y) < self.deadzone else self.left_y
+                self.left_x = 0 if abs(self.left_x) < self.deadzone else self.left_x
                 self.right_x = 0 if abs(self.right_x) < self.deadzone else self.right_x
                 self.right_y = 0 if abs(self.right_y) < self.deadzone else self.right_y
 
@@ -465,8 +465,8 @@ class GamepadControllerHID(InputController):
     def get_deltas(self):
         """Get the current movement deltas from gamepad state."""
         # Calculate deltas - invert as needed based on controller orientation
-        delta_x = -self.left_y * self.x_step_size  # Forward/backward
-        delta_y = -self.left_x * self.y_step_size  # Left/right
+        delta_x = -self.left_x * self.x_step_size  # Forward/backward
+        delta_y = -self.left_y * self.y_step_size  # Left/right
         delta_z = -self.right_y * self.z_step_size  # Up/down
 
         return delta_x, delta_y, delta_z
diff --git a/lerobot/common/teleoperators/gamepad/teleop_gamepad.py b/src/lerobot/teleoperators/gamepad/teleop_gamepad.py
similarity index 100%
rename from lerobot/common/teleoperators/gamepad/teleop_gamepad.py
rename to src/lerobot/teleoperators/gamepad/teleop_gamepad.py
diff --git a/lerobot/common/teleoperators/keyboard/__init__.py b/src/lerobot/teleoperators/keyboard/__init__.py
similarity index 100%
rename from lerobot/common/teleoperators/keyboard/__init__.py
rename to src/lerobot/teleoperators/keyboard/__init__.py
diff --git a/lerobot/common/teleoperators/keyboard/configuration_keyboard.py b/src/lerobot/teleoperators/keyboard/configuration_keyboard.py
similarity index 100%
rename from lerobot/common/teleoperators/keyboard/configuration_keyboard.py
rename to src/lerobot/teleoperators/keyboard/configuration_keyboard.py
diff --git a/lerobot/common/teleoperators/keyboard/teleop_keyboard.py b/src/lerobot/teleoperators/keyboard/teleop_keyboard.py
similarity index 98%
rename from lerobot/common/teleoperators/keyboard/teleop_keyboard.py
rename to src/lerobot/teleoperators/keyboard/teleop_keyboard.py
index bd3ab903e..d034982f1 100644
--- a/lerobot/common/teleoperators/keyboard/teleop_keyboard.py
+++ b/src/lerobot/teleoperators/keyboard/teleop_keyboard.py
@@ -21,7 +21,7 @@ import time
 from queue import Queue
 from typing import Any
 
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 from ..teleoperator import Teleoperator
 from .configuration_keyboard import KeyboardEndEffectorTeleopConfig, KeyboardTeleopConfig
@@ -196,17 +196,18 @@ class KeyboardEndEffectorTeleop(KeyboardTeleop):
         delta_x = 0.0
         delta_y = 0.0
         delta_z = 0.0
+        gripper_action = 1.0
 
         # Generate action based on current key states
         for key, val in self.current_pressed.items():
             if key == keyboard.Key.up:
-                delta_x = int(val)
-            elif key == keyboard.Key.down:
-                delta_x = -int(val)
-            elif key == keyboard.Key.left:
-                delta_y = int(val)
-            elif key == keyboard.Key.right:
                 delta_y = -int(val)
+            elif key == keyboard.Key.down:
+                delta_y = int(val)
+            elif key == keyboard.Key.left:
+                delta_x = int(val)
+            elif key == keyboard.Key.right:
+                delta_x = -int(val)
             elif key == keyboard.Key.shift:
                 delta_z = -int(val)
             elif key == keyboard.Key.shift_r:
@@ -230,7 +231,6 @@ class KeyboardEndEffectorTeleop(KeyboardTeleop):
             "delta_z": delta_z,
         }
 
-        gripper_action = 1  # default gripper action is to stay
         if self.config.use_gripper:
             action_dict["gripper"] = gripper_action
 
diff --git a/lerobot/common/teleoperators/koch_leader/__init__.py b/src/lerobot/teleoperators/koch_leader/__init__.py
similarity index 100%
rename from lerobot/common/teleoperators/koch_leader/__init__.py
rename to src/lerobot/teleoperators/koch_leader/__init__.py
diff --git a/lerobot/common/teleoperators/koch_leader/config_koch_leader.py b/src/lerobot/teleoperators/koch_leader/config_koch_leader.py
similarity index 100%
rename from lerobot/common/teleoperators/koch_leader/config_koch_leader.py
rename to src/lerobot/teleoperators/koch_leader/config_koch_leader.py
diff --git a/lerobot/common/teleoperators/koch_leader/koch_leader.py b/src/lerobot/teleoperators/koch_leader/koch_leader.py
similarity index 97%
rename from lerobot/common/teleoperators/koch_leader/koch_leader.py
rename to src/lerobot/teleoperators/koch_leader/koch_leader.py
index 820acc87c..8eb076fae 100644
--- a/lerobot/common/teleoperators/koch_leader/koch_leader.py
+++ b/src/lerobot/teleoperators/koch_leader/koch_leader.py
@@ -17,9 +17,9 @@
 import logging
 import time
 
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
-from lerobot.common.motors import Motor, MotorCalibration, MotorNormMode
-from lerobot.common.motors.dynamixel import (
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.motors import Motor, MotorCalibration, MotorNormMode
+from lerobot.motors.dynamixel import (
     DriveMode,
     DynamixelMotorsBus,
     OperatingMode,
diff --git a/lerobot/common/teleoperators/so100_leader/__init__.py b/src/lerobot/teleoperators/so100_leader/__init__.py
similarity index 100%
rename from lerobot/common/teleoperators/so100_leader/__init__.py
rename to src/lerobot/teleoperators/so100_leader/__init__.py
diff --git a/lerobot/common/teleoperators/so100_leader/config_so100_leader.py b/src/lerobot/teleoperators/so100_leader/config_so100_leader.py
similarity index 100%
rename from lerobot/common/teleoperators/so100_leader/config_so100_leader.py
rename to src/lerobot/teleoperators/so100_leader/config_so100_leader.py
diff --git a/lerobot/common/teleoperators/so100_leader/so100_leader.py b/src/lerobot/teleoperators/so100_leader/so100_leader.py
similarity index 96%
rename from lerobot/common/teleoperators/so100_leader/so100_leader.py
rename to src/lerobot/teleoperators/so100_leader/so100_leader.py
index 59b083e3f..18dad44d4 100644
--- a/lerobot/common/teleoperators/so100_leader/so100_leader.py
+++ b/src/lerobot/teleoperators/so100_leader/so100_leader.py
@@ -17,9 +17,9 @@
 import logging
 import time
 
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
-from lerobot.common.motors import Motor, MotorCalibration, MotorNormMode
-from lerobot.common.motors.feetech import (
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.motors import Motor, MotorCalibration, MotorNormMode
+from lerobot.motors.feetech import (
     FeetechMotorsBus,
     OperatingMode,
 )
diff --git a/lerobot/common/teleoperators/so101_leader/__init__.py b/src/lerobot/teleoperators/so101_leader/__init__.py
similarity index 100%
rename from lerobot/common/teleoperators/so101_leader/__init__.py
rename to src/lerobot/teleoperators/so101_leader/__init__.py
diff --git a/lerobot/common/teleoperators/so101_leader/config_so101_leader.py b/src/lerobot/teleoperators/so101_leader/config_so101_leader.py
similarity index 100%
rename from lerobot/common/teleoperators/so101_leader/config_so101_leader.py
rename to src/lerobot/teleoperators/so101_leader/config_so101_leader.py
diff --git a/lerobot/common/teleoperators/so101_leader/so101_leader.py b/src/lerobot/teleoperators/so101_leader/so101_leader.py
similarity index 96%
rename from lerobot/common/teleoperators/so101_leader/so101_leader.py
rename to src/lerobot/teleoperators/so101_leader/so101_leader.py
index 80ddfbb1d..2ce28d2e4 100644
--- a/lerobot/common/teleoperators/so101_leader/so101_leader.py
+++ b/src/lerobot/teleoperators/so101_leader/so101_leader.py
@@ -17,9 +17,9 @@
 import logging
 import time
 
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
-from lerobot.common.motors import Motor, MotorCalibration, MotorNormMode
-from lerobot.common.motors.feetech import (
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.motors import Motor, MotorCalibration, MotorNormMode
+from lerobot.motors.feetech import (
     FeetechMotorsBus,
     OperatingMode,
 )
diff --git a/lerobot/common/teleoperators/stretch3_gamepad/__init__.py b/src/lerobot/teleoperators/stretch3_gamepad/__init__.py
similarity index 100%
rename from lerobot/common/teleoperators/stretch3_gamepad/__init__.py
rename to src/lerobot/teleoperators/stretch3_gamepad/__init__.py
diff --git a/lerobot/common/teleoperators/stretch3_gamepad/configuration_stretch3.py b/src/lerobot/teleoperators/stretch3_gamepad/configuration_stretch3.py
similarity index 100%
rename from lerobot/common/teleoperators/stretch3_gamepad/configuration_stretch3.py
rename to src/lerobot/teleoperators/stretch3_gamepad/configuration_stretch3.py
diff --git a/lerobot/common/teleoperators/stretch3_gamepad/stretch3_gamepad.py b/src/lerobot/teleoperators/stretch3_gamepad/stretch3_gamepad.py
similarity index 98%
rename from lerobot/common/teleoperators/stretch3_gamepad/stretch3_gamepad.py
rename to src/lerobot/teleoperators/stretch3_gamepad/stretch3_gamepad.py
index 1e9768c7e..bdcb57d40 100644
--- a/lerobot/common/teleoperators/stretch3_gamepad/stretch3_gamepad.py
+++ b/src/lerobot/teleoperators/stretch3_gamepad/stretch3_gamepad.py
@@ -20,7 +20,7 @@ import numpy as np
 from stretch_body.gamepad_teleop import GamePadTeleop
 from stretch_body.robot_params import RobotParams
 
-from lerobot.common.errors import DeviceAlreadyConnectedError
+from lerobot.errors import DeviceAlreadyConnectedError
 
 from ..teleoperator import Teleoperator
 from .configuration_stretch3 import Stretch3GamePadConfig
diff --git a/lerobot/common/teleoperators/teleoperator.py b/src/lerobot/teleoperators/teleoperator.py
similarity index 97%
rename from lerobot/common/teleoperators/teleoperator.py
rename to src/lerobot/teleoperators/teleoperator.py
index 6a20a3a8a..49f259c17 100644
--- a/lerobot/common/teleoperators/teleoperator.py
+++ b/src/lerobot/teleoperators/teleoperator.py
@@ -18,8 +18,8 @@ from typing import Any, Type
 
 import draccus
 
-from lerobot.common.constants import HF_LEROBOT_CALIBRATION, TELEOPERATORS
-from lerobot.common.motors.motors_bus import MotorCalibration
+from lerobot.constants import HF_LEROBOT_CALIBRATION, TELEOPERATORS
+from lerobot.motors.motors_bus import MotorCalibration
 
 from .config import TeleoperatorConfig
 
diff --git a/lerobot/common/teleoperators/utils.py b/src/lerobot/teleoperators/utils.py
similarity index 100%
rename from lerobot/common/teleoperators/utils.py
rename to src/lerobot/teleoperators/utils.py
diff --git a/lerobot/common/teleoperators/widowx/__init__.py b/src/lerobot/teleoperators/widowx/__init__.py
similarity index 100%
rename from lerobot/common/teleoperators/widowx/__init__.py
rename to src/lerobot/teleoperators/widowx/__init__.py
diff --git a/lerobot/common/teleoperators/widowx/config_widowx.py b/src/lerobot/teleoperators/widowx/config_widowx.py
similarity index 100%
rename from lerobot/common/teleoperators/widowx/config_widowx.py
rename to src/lerobot/teleoperators/widowx/config_widowx.py
diff --git a/lerobot/common/teleoperators/widowx/widowx.py b/src/lerobot/teleoperators/widowx/widowx.py
similarity index 96%
rename from lerobot/common/teleoperators/widowx/widowx.py
rename to src/lerobot/teleoperators/widowx/widowx.py
index 8a42c9063..6becd767f 100644
--- a/lerobot/common/teleoperators/widowx/widowx.py
+++ b/src/lerobot/teleoperators/widowx/widowx.py
@@ -17,9 +17,9 @@
 import logging
 import time
 
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
-from lerobot.common.motors import Motor, MotorCalibration, MotorNormMode
-from lerobot.common.motors.dynamixel import (
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.motors import Motor, MotorCalibration, MotorNormMode
+from lerobot.motors.dynamixel import (
     DriveMode,
     DynamixelMotorsBus,
     OperatingMode,
diff --git a/src/lerobot/templates/lerobot_modelcard_template.md b/src/lerobot/templates/lerobot_modelcard_template.md
new file mode 100644
index 000000000..64ad7196c
--- /dev/null
+++ b/src/lerobot/templates/lerobot_modelcard_template.md
@@ -0,0 +1,74 @@
+---
+# For reference on model card metadata, see the spec: https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1
+# Doc / guide: https://huggingface.co/docs/hub/model-cards
+{{ card_data }}
+---
+
+# Model Card for {{ model_name | default("Model ID", true) }}
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+{% if model_name == "smolvla" %}
+[SmolVLA](https://huggingface.co/papers/2506.01844) is a compact, efficient vision-language-action model that achieves competitive performance at reduced computational costs and can be deployed on consumer-grade hardware.
+{% elif model_name == "act" %}
+[Action Chunking with Transformers (ACT)](https://huggingface.co/papers/2304.13705) is an imitation-learning method that predicts short action chunks instead of single steps. It learns from teleoperated data and often achieves high success rates.
+{% elif model_name == "tdmpc" %}
+[TD-MPC](https://huggingface.co/papers/2203.04955) combines model-free and model-based approaches to improve sample efficiency and performance in continuous control tasks by using a learned latent dynamics model and terminal value function.
+{% elif model_name == "diffusion" %}
+[Diffusion Policy](https://huggingface.co/papers/2303.04137) treats visuomotor control as a generative diffusion process, producing smooth, multi-step action trajectories that excel at contact-rich manipulation.
+{% elif model_name == "vqbet" %}
+[VQ-BET](https://huggingface.co/papers/2403.03181) combines vector-quantised action tokens with Behaviour Transformers to discretise control and achieve data-efficient imitation across diverse skills.
+{% elif model_name == "pi0" %}
+[Pi0](https://huggingface.co/papers/2410.24164) is a generalist vision-language-action transformer that converts multimodal observations and text instructions into robot actions for zero-shot task transfer.
+{% elif model_name == "pi0fast" %}
+[Pi0-Fast](https://huggingface.co/papers/2501.09747) is a variant of Pi0 that uses a new tokenization method called FAST, which enables training of an autoregressive vision-language-action policy for high-frequency robotic tasks with improved performance and reduced training time.
+{% elif model_name == "sac" %}
+[Soft Actor-Critic (SAC)](https://huggingface.co/papers/1801.01290) is an entropy-regularised actor-critic algorithm offering stable, sample-efficient learning in continuous-control environments.
+{% elif model_name == "reward_classifier" %}
+A reward classifier is a lightweight neural network that scores observations or trajectories for task success, providing a learned reward signal or offline evaluation when explicit rewards are unavailable.
+{% else %}
+_Model type not recognized — please update this template._
+{% endif %}
+
+This policy has been trained and pushed to the Hub using [LeRobot](https://github.com/huggingface/lerobot).
+See the full documentation at [LeRobot Docs](https://huggingface.co/docs/lerobot/index).
+
+---
+
+## How to Get Started with the Model
+
+For a complete walkthrough, see the [training guide](https://huggingface.co/docs/lerobot/il_robots#train-a-policy).
+Below is the short version on how to train and run inference/eval:
+
+### Train from scratch
+
+```bash
+python -m lerobot.scripts.train \
+  --dataset.repo_id=${HF_USER}/<dataset> \
+  --policy.type=act \
+  --output_dir=outputs/train/<desired_policy_repo_id> \
+  --job_name=lerobot_training \
+  --policy.device=cuda \
+  --policy.repo_id=${HF_USER}/<desired_policy_repo_id>
+  --wandb.enable=true
+```
+
+*Writes checkpoints to `outputs/train/<desired_policy_repo_id>/checkpoints/`.*
+
+### Evaluate the policy/run inference
+
+```bash
+python -m lerobot.record \
+  --robot.type=so100_follower \
+  --dataset.repo_id=<hf_user>/eval_<dataset> \
+  --policy.path=<hf_user>/<desired_policy_repo_id> \
+  --episodes=10
+```
+
+Prefix the dataset repo with **eval\_** and supply `--policy.path` pointing to a local or hub checkpoint.
+
+---
+
+## Model Details
+
+* **License:** {{ license | default("\[More Information Needed]", true) }}
diff --git a/lerobot/templates/visualize_dataset_homepage.html b/src/lerobot/templates/visualize_dataset_homepage.html
similarity index 100%
rename from lerobot/templates/visualize_dataset_homepage.html
rename to src/lerobot/templates/visualize_dataset_homepage.html
diff --git a/lerobot/templates/visualize_dataset_template.html b/src/lerobot/templates/visualize_dataset_template.html
similarity index 100%
rename from lerobot/templates/visualize_dataset_template.html
rename to src/lerobot/templates/visualize_dataset_template.html
diff --git a/lerobot/common/transport/services.proto b/src/lerobot/transport/services.proto
similarity index 97%
rename from lerobot/common/transport/services.proto
rename to src/lerobot/transport/services.proto
index 29d00005a..89bfc107a 100644
--- a/lerobot/common/transport/services.proto
+++ b/src/lerobot/transport/services.proto
@@ -15,7 +15,7 @@
 
 // To generate a classes for transport part (services_pb2.py and services_pb2_grpc.py) use the following command:
 //
-// python -m grpc_tools.protoc -I . --python_out=. --grpc_python_out=. lerobot/common/transport/services.proto
+// python -m grpc_tools.protoc -I . --python_out=. --grpc_python_out=. src/lerobot/transport/services.proto
 //
 // The command should be launched from the root of the project.
 
diff --git a/src/lerobot/transport/services_pb2.py b/src/lerobot/transport/services_pb2.py
new file mode 100644
index 000000000..8a2137687
--- /dev/null
+++ b/src/lerobot/transport/services_pb2.py
@@ -0,0 +1,45 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# NO CHECKED-IN PROTOBUF GENCODE
+# source: src/lerobot/transport/services.proto
+# Protobuf Python Version: 5.29.0
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    5,
+    29,
+    0,
+    '',
+    'src/lerobot/transport/services.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n$src/lerobot/transport/services.proto\x12\ttransport\"L\n\nTransition\x12\x30\n\x0etransfer_state\x18\x01 \x01(\x0e\x32\x18.transport.TransferState\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"L\n\nParameters\x12\x30\n\x0etransfer_state\x18\x01 \x01(\x0e\x32\x18.transport.TransferState\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"T\n\x12InteractionMessage\x12\x30\n\x0etransfer_state\x18\x01 \x01(\x0e\x32\x18.transport.TransferState\x12\x0c\n\x04\x64\x61ta\x18\x02 \x01(\x0c\"\x07\n\x05\x45mpty*`\n\rTransferState\x12\x14\n\x10TRANSFER_UNKNOWN\x10\x00\x12\x12\n\x0eTRANSFER_BEGIN\x10\x01\x12\x13\n\x0fTRANSFER_MIDDLE\x10\x02\x12\x10\n\x0cTRANSFER_END\x10\x03\x32\x81\x02\n\x0eLearnerService\x12=\n\x10StreamParameters\x12\x10.transport.Empty\x1a\x15.transport.Parameters0\x01\x12<\n\x0fSendTransitions\x12\x15.transport.Transition\x1a\x10.transport.Empty(\x01\x12\x45\n\x10SendInteractions\x12\x1d.transport.InteractionMessage\x1a\x10.transport.Empty(\x01\x12+\n\x05Ready\x12\x10.transport.Empty\x1a\x10.transport.Emptyb\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'src.lerobot.transport.services_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_TRANSFERSTATE']._serialized_start=302
+  _globals['_TRANSFERSTATE']._serialized_end=398
+  _globals['_TRANSITION']._serialized_start=51
+  _globals['_TRANSITION']._serialized_end=127
+  _globals['_PARAMETERS']._serialized_start=129
+  _globals['_PARAMETERS']._serialized_end=205
+  _globals['_INTERACTIONMESSAGE']._serialized_start=207
+  _globals['_INTERACTIONMESSAGE']._serialized_end=291
+  _globals['_EMPTY']._serialized_start=293
+  _globals['_EMPTY']._serialized_end=300
+  _globals['_LEARNERSERVICE']._serialized_start=401
+  _globals['_LEARNERSERVICE']._serialized_end=658
+# @@protoc_insertion_point(module_scope)
diff --git a/lerobot/common/transport/services_pb2_grpc.py b/src/lerobot/transport/services_pb2_grpc.py
similarity index 72%
rename from lerobot/common/transport/services_pb2_grpc.py
rename to src/lerobot/transport/services_pb2_grpc.py
index 5a7a924fd..a4fe8c576 100644
--- a/lerobot/common/transport/services_pb2_grpc.py
+++ b/src/lerobot/transport/services_pb2_grpc.py
@@ -3,7 +3,7 @@
 import grpc
 import warnings
 
-from lerobot.common.transport import services_pb2 as lerobot_dot_common_dot_transport_dot_services__pb2
+from src.lerobot.transport import services_pb2 as src_dot_lerobot_dot_transport_dot_services__pb2
 
 GRPC_GENERATED_VERSION = '1.71.0'
 GRPC_VERSION = grpc.__version__
@@ -18,7 +18,7 @@ except ImportError:
 if _version_not_supported:
     raise RuntimeError(
         f'The grpc package installed is at version {GRPC_VERSION},'
-        + f' but the generated code in lerobot/common/transport/services_pb2_grpc.py depends on'
+        + f' but the generated code in src/lerobot/transport/services_pb2_grpc.py depends on'
         + f' grpcio>={GRPC_GENERATED_VERSION}.'
         + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
         + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
@@ -38,23 +38,23 @@ class LearnerServiceStub:
         """
         self.StreamParameters = channel.unary_stream(
                 '/transport.LearnerService/StreamParameters',
-                request_serializer=lerobot_dot_common_dot_transport_dot_services__pb2.Empty.SerializeToString,
-                response_deserializer=lerobot_dot_common_dot_transport_dot_services__pb2.Parameters.FromString,
+                request_serializer=src_dot_lerobot_dot_transport_dot_services__pb2.Empty.SerializeToString,
+                response_deserializer=src_dot_lerobot_dot_transport_dot_services__pb2.Parameters.FromString,
                 _registered_method=True)
         self.SendTransitions = channel.stream_unary(
                 '/transport.LearnerService/SendTransitions',
-                request_serializer=lerobot_dot_common_dot_transport_dot_services__pb2.Transition.SerializeToString,
-                response_deserializer=lerobot_dot_common_dot_transport_dot_services__pb2.Empty.FromString,
+                request_serializer=src_dot_lerobot_dot_transport_dot_services__pb2.Transition.SerializeToString,
+                response_deserializer=src_dot_lerobot_dot_transport_dot_services__pb2.Empty.FromString,
                 _registered_method=True)
         self.SendInteractions = channel.stream_unary(
                 '/transport.LearnerService/SendInteractions',
-                request_serializer=lerobot_dot_common_dot_transport_dot_services__pb2.InteractionMessage.SerializeToString,
-                response_deserializer=lerobot_dot_common_dot_transport_dot_services__pb2.Empty.FromString,
+                request_serializer=src_dot_lerobot_dot_transport_dot_services__pb2.InteractionMessage.SerializeToString,
+                response_deserializer=src_dot_lerobot_dot_transport_dot_services__pb2.Empty.FromString,
                 _registered_method=True)
         self.Ready = channel.unary_unary(
                 '/transport.LearnerService/Ready',
-                request_serializer=lerobot_dot_common_dot_transport_dot_services__pb2.Empty.SerializeToString,
-                response_deserializer=lerobot_dot_common_dot_transport_dot_services__pb2.Empty.FromString,
+                request_serializer=src_dot_lerobot_dot_transport_dot_services__pb2.Empty.SerializeToString,
+                response_deserializer=src_dot_lerobot_dot_transport_dot_services__pb2.Empty.FromString,
                 _registered_method=True)
 
 
@@ -93,23 +93,23 @@ def add_LearnerServiceServicer_to_server(servicer, server):
     rpc_method_handlers = {
             'StreamParameters': grpc.unary_stream_rpc_method_handler(
                     servicer.StreamParameters,
-                    request_deserializer=lerobot_dot_common_dot_transport_dot_services__pb2.Empty.FromString,
-                    response_serializer=lerobot_dot_common_dot_transport_dot_services__pb2.Parameters.SerializeToString,
+                    request_deserializer=src_dot_lerobot_dot_transport_dot_services__pb2.Empty.FromString,
+                    response_serializer=src_dot_lerobot_dot_transport_dot_services__pb2.Parameters.SerializeToString,
             ),
             'SendTransitions': grpc.stream_unary_rpc_method_handler(
                     servicer.SendTransitions,
-                    request_deserializer=lerobot_dot_common_dot_transport_dot_services__pb2.Transition.FromString,
-                    response_serializer=lerobot_dot_common_dot_transport_dot_services__pb2.Empty.SerializeToString,
+                    request_deserializer=src_dot_lerobot_dot_transport_dot_services__pb2.Transition.FromString,
+                    response_serializer=src_dot_lerobot_dot_transport_dot_services__pb2.Empty.SerializeToString,
             ),
             'SendInteractions': grpc.stream_unary_rpc_method_handler(
                     servicer.SendInteractions,
-                    request_deserializer=lerobot_dot_common_dot_transport_dot_services__pb2.InteractionMessage.FromString,
-                    response_serializer=lerobot_dot_common_dot_transport_dot_services__pb2.Empty.SerializeToString,
+                    request_deserializer=src_dot_lerobot_dot_transport_dot_services__pb2.InteractionMessage.FromString,
+                    response_serializer=src_dot_lerobot_dot_transport_dot_services__pb2.Empty.SerializeToString,
             ),
             'Ready': grpc.unary_unary_rpc_method_handler(
                     servicer.Ready,
-                    request_deserializer=lerobot_dot_common_dot_transport_dot_services__pb2.Empty.FromString,
-                    response_serializer=lerobot_dot_common_dot_transport_dot_services__pb2.Empty.SerializeToString,
+                    request_deserializer=src_dot_lerobot_dot_transport_dot_services__pb2.Empty.FromString,
+                    response_serializer=src_dot_lerobot_dot_transport_dot_services__pb2.Empty.SerializeToString,
             ),
     }
     generic_handler = grpc.method_handlers_generic_handler(
@@ -139,8 +139,8 @@ class LearnerService:
             request,
             target,
             '/transport.LearnerService/StreamParameters',
-            lerobot_dot_common_dot_transport_dot_services__pb2.Empty.SerializeToString,
-            lerobot_dot_common_dot_transport_dot_services__pb2.Parameters.FromString,
+            src_dot_lerobot_dot_transport_dot_services__pb2.Empty.SerializeToString,
+            src_dot_lerobot_dot_transport_dot_services__pb2.Parameters.FromString,
             options,
             channel_credentials,
             insecure,
@@ -166,8 +166,8 @@ class LearnerService:
             request_iterator,
             target,
             '/transport.LearnerService/SendTransitions',
-            lerobot_dot_common_dot_transport_dot_services__pb2.Transition.SerializeToString,
-            lerobot_dot_common_dot_transport_dot_services__pb2.Empty.FromString,
+            src_dot_lerobot_dot_transport_dot_services__pb2.Transition.SerializeToString,
+            src_dot_lerobot_dot_transport_dot_services__pb2.Empty.FromString,
             options,
             channel_credentials,
             insecure,
@@ -193,8 +193,8 @@ class LearnerService:
             request_iterator,
             target,
             '/transport.LearnerService/SendInteractions',
-            lerobot_dot_common_dot_transport_dot_services__pb2.InteractionMessage.SerializeToString,
-            lerobot_dot_common_dot_transport_dot_services__pb2.Empty.FromString,
+            src_dot_lerobot_dot_transport_dot_services__pb2.InteractionMessage.SerializeToString,
+            src_dot_lerobot_dot_transport_dot_services__pb2.Empty.FromString,
             options,
             channel_credentials,
             insecure,
@@ -220,8 +220,8 @@ class LearnerService:
             request,
             target,
             '/transport.LearnerService/Ready',
-            lerobot_dot_common_dot_transport_dot_services__pb2.Empty.SerializeToString,
-            lerobot_dot_common_dot_transport_dot_services__pb2.Empty.FromString,
+            src_dot_lerobot_dot_transport_dot_services__pb2.Empty.SerializeToString,
+            src_dot_lerobot_dot_transport_dot_services__pb2.Empty.FromString,
             options,
             channel_credentials,
             insecure,
diff --git a/lerobot/common/transport/utils.py b/src/lerobot/transport/utils.py
similarity index 97%
rename from lerobot/common/transport/utils.py
rename to src/lerobot/transport/utils.py
index 774721fc6..1c6683262 100644
--- a/lerobot/common/transport/utils.py
+++ b/src/lerobot/transport/utils.py
@@ -23,8 +23,8 @@ from typing import Any
 
 import torch
 
-from lerobot.common.transport import services_pb2
-from lerobot.common.utils.transition import Transition
+from lerobot.transport import services_pb2
+from lerobot.utils.transition import Transition
 
 CHUNK_SIZE = 2 * 1024 * 1024  # 2 MB
 
diff --git a/lerobot/common/utils/benchmark.py b/src/lerobot/utils/benchmark.py
similarity index 100%
rename from lerobot/common/utils/benchmark.py
rename to src/lerobot/utils/benchmark.py
diff --git a/lerobot/common/utils/buffer.py b/src/lerobot/utils/buffer.py
similarity index 99%
rename from lerobot/common/utils/buffer.py
rename to src/lerobot/utils/buffer.py
index 9ae231ad9..7f8d989dd 100644
--- a/lerobot/common/utils/buffer.py
+++ b/src/lerobot/utils/buffer.py
@@ -22,8 +22,8 @@ import torch
 import torch.nn.functional as F  # noqa: N812
 from tqdm import tqdm
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.utils.transition import Transition
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.utils.transition import Transition
 
 
 class BatchTransition(TypedDict):
diff --git a/lerobot/common/utils/control_utils.py b/src/lerobot/utils/control_utils.py
similarity index 97%
rename from lerobot/common/utils/control_utils.py
rename to src/lerobot/utils/control_utils.py
index b66977a72..4bcc241da 100644
--- a/lerobot/common/utils/control_utils.py
+++ b/src/lerobot/utils/control_utils.py
@@ -28,10 +28,10 @@ import torch
 from deepdiff import DeepDiff
 from termcolor import colored
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.datasets.utils import DEFAULT_FEATURES
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.robots import Robot
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.utils import DEFAULT_FEATURES
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.robots import Robot
 
 
 def log_control_info(robot: Robot, dt_s, episode_index=None, frame_index=None, fps=None):
diff --git a/lerobot/common/utils/encoding_utils.py b/src/lerobot/utils/encoding_utils.py
similarity index 100%
rename from lerobot/common/utils/encoding_utils.py
rename to src/lerobot/utils/encoding_utils.py
diff --git a/lerobot/common/utils/hub.py b/src/lerobot/utils/hub.py
similarity index 100%
rename from lerobot/common/utils/hub.py
rename to src/lerobot/utils/hub.py
diff --git a/lerobot/common/utils/import_utils.py b/src/lerobot/utils/import_utils.py
similarity index 100%
rename from lerobot/common/utils/import_utils.py
rename to src/lerobot/utils/import_utils.py
diff --git a/lerobot/common/utils/io_utils.py b/src/lerobot/utils/io_utils.py
similarity index 100%
rename from lerobot/common/utils/io_utils.py
rename to src/lerobot/utils/io_utils.py
diff --git a/lerobot/common/utils/logging_utils.py b/src/lerobot/utils/logging_utils.py
similarity index 98%
rename from lerobot/common/utils/logging_utils.py
rename to src/lerobot/utils/logging_utils.py
index 56c9abb23..b6404e66d 100644
--- a/lerobot/common/utils/logging_utils.py
+++ b/src/lerobot/utils/logging_utils.py
@@ -15,7 +15,7 @@
 # limitations under the License.
 from typing import Any
 
-from lerobot.common.utils.utils import format_big_number
+from lerobot.utils.utils import format_big_number
 
 
 class AverageMeter:
diff --git a/lerobot/common/utils/process.py b/src/lerobot/utils/process.py
similarity index 100%
rename from lerobot/common/utils/process.py
rename to src/lerobot/utils/process.py
diff --git a/lerobot/common/utils/queue.py b/src/lerobot/utils/queue.py
similarity index 100%
rename from lerobot/common/utils/queue.py
rename to src/lerobot/utils/queue.py
diff --git a/lerobot/common/utils/random_utils.py b/src/lerobot/utils/random_utils.py
similarity index 98%
rename from lerobot/common/utils/random_utils.py
rename to src/lerobot/utils/random_utils.py
index 3d9bf4dd8..31fed1da6 100644
--- a/lerobot/common/utils/random_utils.py
+++ b/src/lerobot/utils/random_utils.py
@@ -22,8 +22,8 @@ import numpy as np
 import torch
 from safetensors.torch import load_file, save_file
 
-from lerobot.common.constants import RNG_STATE
-from lerobot.common.datasets.utils import flatten_dict, unflatten_dict
+from lerobot.constants import RNG_STATE
+from lerobot.datasets.utils import flatten_dict, unflatten_dict
 
 
 def serialize_python_rng_state() -> dict[str, torch.Tensor]:
diff --git a/lerobot/common/utils/robot_utils.py b/src/lerobot/utils/robot_utils.py
similarity index 100%
rename from lerobot/common/utils/robot_utils.py
rename to src/lerobot/utils/robot_utils.py
diff --git a/lerobot/common/utils/train_utils.py b/src/lerobot/utils/train_utils.py
similarity index 93%
rename from lerobot/common/utils/train_utils.py
rename to src/lerobot/utils/train_utils.py
index a79983128..2859fe057 100644
--- a/lerobot/common/utils/train_utils.py
+++ b/src/lerobot/utils/train_utils.py
@@ -20,19 +20,19 @@ from termcolor import colored
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 
-from lerobot.common.constants import (
+from lerobot.configs.train import TrainPipelineConfig
+from lerobot.constants import (
     CHECKPOINTS_DIR,
     LAST_CHECKPOINT_LINK,
     PRETRAINED_MODEL_DIR,
     TRAINING_STATE_DIR,
     TRAINING_STEP,
 )
-from lerobot.common.datasets.utils import load_json, write_json
-from lerobot.common.optim.optimizers import load_optimizer_state, save_optimizer_state
-from lerobot.common.optim.schedulers import load_scheduler_state, save_scheduler_state
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.utils.random_utils import load_rng_state, save_rng_state
-from lerobot.configs.train import TrainPipelineConfig
+from lerobot.datasets.utils import load_json, write_json
+from lerobot.optim.optimizers import load_optimizer_state, save_optimizer_state
+from lerobot.optim.schedulers import load_scheduler_state, save_scheduler_state
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.utils.random_utils import load_rng_state, save_rng_state
 
 
 def log_output_dir(out_dir):
diff --git a/lerobot/common/utils/transition.py b/src/lerobot/utils/transition.py
similarity index 100%
rename from lerobot/common/utils/transition.py
rename to src/lerobot/utils/transition.py
diff --git a/lerobot/common/utils/utils.py b/src/lerobot/utils/utils.py
similarity index 98%
rename from lerobot/common/utils/utils.py
rename to src/lerobot/utils/utils.py
index cba65ba45..6d9c0338b 100644
--- a/lerobot/common/utils/utils.py
+++ b/src/lerobot/utils/utils.py
@@ -184,7 +184,7 @@ def capture_timestamp_utc():
     return datetime.now(timezone.utc)
 
 
-def say(text, blocking=False):
+def say(text: str, blocking: bool = False):
     system = platform.system()
 
     if system == "Darwin":
@@ -212,7 +212,7 @@ def say(text, blocking=False):
         subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
 
 
-def log_say(text, play_sounds, blocking=False):
+def log_say(text: str, play_sounds: bool = True, blocking: bool = False):
     logging.info(text)
 
     if play_sounds:
diff --git a/lerobot/common/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py
similarity index 56%
rename from lerobot/common/utils/visualization_utils.py
rename to src/lerobot/utils/visualization_utils.py
index dfffece5f..f0f9aebb7 100644
--- a/lerobot/common/utils/visualization_utils.py
+++ b/src/lerobot/utils/visualization_utils.py
@@ -13,7 +13,9 @@
 # limitations under the License.
 
 import os
+from typing import Any
 
+import numpy as np
 import rerun as rr
 
 
@@ -24,3 +26,21 @@ def _init_rerun(session_name: str = "lerobot_control_loop") -> None:
     rr.init(session_name)
     memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%")
     rr.spawn(memory_limit=memory_limit)
+
+
+def log_rerun_data(observation: dict[str | Any], action: dict[str | Any]):
+    for obs, val in observation.items():
+        if isinstance(val, float):
+            rr.log(f"observation.{obs}", rr.Scalar(val))
+        elif isinstance(val, np.ndarray):
+            if val.ndim == 1:
+                for i, v in enumerate(val):
+                    rr.log(f"observation.{obs}_{i}", rr.Scalar(float(v)))
+            else:
+                rr.log(f"observation.{obs}", rr.Image(val), static=True)
+    for act, val in action.items():
+        if isinstance(val, float):
+            rr.log(f"action.{act}", rr.Scalar(val))
+        elif isinstance(val, np.ndarray):
+            for i, v in enumerate(val):
+                rr.log(f"action.{act}_{i}", rr.Scalar(float(v)))
diff --git a/lerobot/common/utils/wandb_utils.py b/src/lerobot/utils/wandb_utils.py
similarity index 99%
rename from lerobot/common/utils/wandb_utils.py
rename to src/lerobot/utils/wandb_utils.py
index ac4d22343..91b4ec95c 100644
--- a/lerobot/common/utils/wandb_utils.py
+++ b/src/lerobot/utils/wandb_utils.py
@@ -22,8 +22,8 @@ from pathlib import Path
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from termcolor import colored
 
-from lerobot.common.constants import PRETRAINED_MODEL_DIR
 from lerobot.configs.train import TrainPipelineConfig
+from lerobot.constants import PRETRAINED_MODEL_DIR
 
 
 def cfg_to_group(cfg: TrainPipelineConfig, return_list: bool = False) -> list[str] | str:
diff --git a/tests/artifacts/datasets/save_dataset_to_safetensors.py b/tests/artifacts/datasets/save_dataset_to_safetensors.py
index 74d42a3d5..419961b20 100644
--- a/tests/artifacts/datasets/save_dataset_to_safetensors.py
+++ b/tests/artifacts/datasets/save_dataset_to_safetensors.py
@@ -31,7 +31,7 @@ from pathlib import Path
 
 from safetensors.torch import save_file
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
 
 
 def save_dataset_to_safetensors(output_dir, repo_id="lerobot/pusht"):
diff --git a/tests/artifacts/image_transforms/save_image_transforms_to_safetensors.py b/tests/artifacts/image_transforms/save_image_transforms_to_safetensors.py
index be47d9a47..ce15d16fd 100644
--- a/tests/artifacts/image_transforms/save_image_transforms_to_safetensors.py
+++ b/tests/artifacts/image_transforms/save_image_transforms_to_safetensors.py
@@ -18,14 +18,14 @@ from pathlib import Path
 import torch
 from safetensors.torch import save_file
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.datasets.transforms import (
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.datasets.transforms import (
     ImageTransformConfig,
     ImageTransforms,
     ImageTransformsConfig,
     make_transform_from_config,
 )
-from lerobot.common.utils.random_utils import seeded_context
+from lerobot.utils.random_utils import seeded_context
 
 ARTIFACT_DIR = Path("tests/artifacts/image_transforms")
 DATASET_REPO_ID = "lerobot/aloha_static_cups_open"
diff --git a/tests/artifacts/policies/save_policy_to_safetensors.py b/tests/artifacts/policies/save_policy_to_safetensors.py
index 106f0dc04..6ccb47c3e 100644
--- a/tests/artifacts/policies/save_policy_to_safetensors.py
+++ b/tests/artifacts/policies/save_policy_to_safetensors.py
@@ -19,12 +19,12 @@ from pathlib import Path
 import torch
 from safetensors.torch import save_file
 
-from lerobot.common.datasets.factory import make_dataset
-from lerobot.common.optim.factory import make_optimizer_and_scheduler
-from lerobot.common.policies.factory import make_policy, make_policy_config
-from lerobot.common.utils.random_utils import set_seed
 from lerobot.configs.default import DatasetConfig
 from lerobot.configs.train import TrainPipelineConfig
+from lerobot.datasets.factory import make_dataset
+from lerobot.optim.factory import make_optimizer_and_scheduler
+from lerobot.policies.factory import make_policy, make_policy_config
+from lerobot.utils.random_utils import set_seed
 
 
 def get_policy_stats(ds_repo_id: str, policy_name: str, policy_kwargs: dict):
@@ -32,7 +32,7 @@ def get_policy_stats(ds_repo_id: str, policy_name: str, policy_kwargs: dict):
     train_cfg = TrainPipelineConfig(
         # TODO(rcadene, aliberts): remove dataset download
         dataset=DatasetConfig(repo_id=ds_repo_id, episodes=[0]),
-        policy=make_policy_config(policy_name, **policy_kwargs),
+        policy=make_policy_config(policy_name, push_to_hub=False, **policy_kwargs),
     )
     train_cfg.validate()  # Needed for auto-setting some parameters
 
diff --git a/tests/cameras/test_opencv.py b/tests/cameras/test_opencv.py
index 7ba04b261..a9c060c4f 100644
--- a/tests/cameras/test_opencv.py
+++ b/tests/cameras/test_opencv.py
@@ -24,9 +24,9 @@ from pathlib import Path
 import numpy as np
 import pytest
 
-from lerobot.common.cameras.configs import Cv2Rotation
-from lerobot.common.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.cameras.configs import Cv2Rotation
+from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 # NOTE(Steven): more tests + assertions?
 TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
diff --git a/tests/cameras/test_realsense.py b/tests/cameras/test_realsense.py
index 5fb1767fe..3957baf2d 100644
--- a/tests/cameras/test_realsense.py
+++ b/tests/cameras/test_realsense.py
@@ -25,12 +25,12 @@ from unittest.mock import patch
 import numpy as np
 import pytest
 
-from lerobot.common.cameras.configs import Cv2Rotation
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.cameras.configs import Cv2Rotation
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
 
 pytest.importorskip("pyrealsense2")
 
-from lerobot.common.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
+from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
 
 TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
 BAG_FILE_PATH = TEST_ARTIFACTS_DIR / "test_rs.bag"
diff --git a/tests/configs/test_plugin_loading.py b/tests/configs/test_plugin_loading.py
index 1a8cceed7..957574eb4 100644
--- a/tests/configs/test_plugin_loading.py
+++ b/tests/configs/test_plugin_loading.py
@@ -5,15 +5,15 @@ from typing import Generator
 
 import pytest
 
-from lerobot.common.envs.configs import EnvConfig
 from lerobot.configs.parser import PluginLoadError, load_plugin, parse_plugin_args, wrap
+from lerobot.envs.configs import EnvConfig
 
 
 def create_plugin_code(*, base_class: str = "EnvConfig", plugin_name: str = "test_env") -> str:
     """Creates a dummy plugin module that implements its own EnvConfig subclass."""
     return f"""
 from dataclasses import dataclass
-from lerobot.common.envs.configs import {base_class}
+from lerobot.envs.configs import {base_class}
 
 @{base_class}.register_subclass("{plugin_name}")
 @dataclass
diff --git a/tests/datasets/test_compute_stats.py b/tests/datasets/test_compute_stats.py
index d9032c8a3..8f8179c29 100644
--- a/tests/datasets/test_compute_stats.py
+++ b/tests/datasets/test_compute_stats.py
@@ -18,7 +18,7 @@ from unittest.mock import patch
 import numpy as np
 import pytest
 
-from lerobot.common.datasets.compute_stats import (
+from lerobot.datasets.compute_stats import (
     _assert_type_and_shape,
     aggregate_feature_stats,
     aggregate_stats,
@@ -61,7 +61,7 @@ def test_sample_indices():
     assert len(indices) == estimate_num_samples(10)
 
 
-@patch("lerobot.common.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy)
+@patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy)
 def test_sample_images(mock_load):
     image_paths = [f"image_{i}.jpg" for i in range(100)]
     images = sample_images(image_paths)
@@ -144,9 +144,7 @@ def test_compute_episode_stats():
         "observation.state": {"dtype": "numeric"},
     }
 
-    with patch(
-        "lerobot.common.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy
-    ):
+    with patch("lerobot.datasets.compute_stats.load_image_as_numpy", side_effect=mock_load_image_as_numpy):
         stats = compute_episode_stats(episode_data, features)
 
     assert "observation.image" in stats and "observation.state" in stats
diff --git a/tests/datasets/test_datasets.py b/tests/datasets/test_datasets.py
index 55a417c30..d3b78ddcc 100644
--- a/tests/datasets/test_datasets.py
+++ b/tests/datasets/test_datasets.py
@@ -28,21 +28,21 @@ from PIL import Image
 from safetensors.torch import load_file
 
 import lerobot
-from lerobot.common.datasets.factory import make_dataset
-from lerobot.common.datasets.image_writer import image_array_to_pil_image
-from lerobot.common.datasets.lerobot_dataset import (
+from lerobot.configs.default import DatasetConfig
+from lerobot.configs.train import TrainPipelineConfig
+from lerobot.datasets.factory import make_dataset
+from lerobot.datasets.image_writer import image_array_to_pil_image
+from lerobot.datasets.lerobot_dataset import (
     LeRobotDataset,
     MultiLeRobotDataset,
 )
-from lerobot.common.datasets.utils import (
+from lerobot.datasets.utils import (
     create_branch,
     flatten_dict,
     unflatten_dict,
 )
-from lerobot.common.envs.factory import make_env_config
-from lerobot.common.policies.factory import make_policy_config
-from lerobot.configs.default import DatasetConfig
-from lerobot.configs.train import TrainPipelineConfig
+from lerobot.envs.factory import make_env_config
+from lerobot.policies.factory import make_policy_config
 from tests.fixtures.constants import DUMMY_CHW, DUMMY_HWC, DUMMY_REPO_ID
 from tests.utils import require_x86_64_kernel
 
@@ -338,8 +338,9 @@ def test_factory(env_name, repo_id, policy_name):
         # TODO(rcadene, aliberts): remove dataset download
         dataset=DatasetConfig(repo_id=repo_id, episodes=[0]),
         env=make_env_config(env_name),
-        policy=make_policy_config(policy_name),
+        policy=make_policy_config(policy_name, push_to_hub=False),
     )
+    cfg.validate()
 
     dataset = make_dataset(cfg)
     delta_timestamps = dataset.delta_timestamps
@@ -557,7 +558,7 @@ def test_create_branch():
 
 def test_dataset_feature_with_forward_slash_raises_error():
     # make sure dir does not exist
-    from lerobot.common.constants import HF_LEROBOT_HOME
+    from lerobot.constants import HF_LEROBOT_HOME
 
     dataset_dir = HF_LEROBOT_HOME / "lerobot/test/with/slash"
     # make sure does not exist
diff --git a/tests/datasets/test_delta_timestamps.py b/tests/datasets/test_delta_timestamps.py
index 350146420..786b90ce2 100644
--- a/tests/datasets/test_delta_timestamps.py
+++ b/tests/datasets/test_delta_timestamps.py
@@ -19,7 +19,7 @@ import pyarrow.compute as pc
 import pytest
 import torch
 
-from lerobot.common.datasets.utils import (
+from lerobot.datasets.utils import (
     check_delta_timestamps,
     check_timestamps_sync,
     get_delta_indices,
diff --git a/tests/datasets/test_image_transforms.py b/tests/datasets/test_image_transforms.py
index 146a4dcd4..3ab93cb2c 100644
--- a/tests/datasets/test_image_transforms.py
+++ b/tests/datasets/test_image_transforms.py
@@ -21,7 +21,7 @@ from safetensors.torch import load_file
 from torchvision.transforms import v2
 from torchvision.transforms.v2 import functional as F  # noqa: N812
 
-from lerobot.common.datasets.transforms import (
+from lerobot.datasets.transforms import (
     ImageTransformConfig,
     ImageTransforms,
     ImageTransformsConfig,
@@ -29,11 +29,11 @@ from lerobot.common.datasets.transforms import (
     SharpnessJitter,
     make_transform_from_config,
 )
-from lerobot.common.utils.random_utils import seeded_context
 from lerobot.scripts.visualize_image_transforms import (
     save_all_transforms,
     save_each_transform,
 )
+from lerobot.utils.random_utils import seeded_context
 from tests.artifacts.image_transforms.save_image_transforms_to_safetensors import ARTIFACT_DIR
 from tests.utils import require_x86_64_kernel
 
diff --git a/tests/datasets/test_image_writer.py b/tests/datasets/test_image_writer.py
index 802fe0d36..99c8b24fc 100644
--- a/tests/datasets/test_image_writer.py
+++ b/tests/datasets/test_image_writer.py
@@ -20,7 +20,7 @@ import numpy as np
 import pytest
 from PIL import Image
 
-from lerobot.common.datasets.image_writer import (
+from lerobot.datasets.image_writer import (
     AsyncImageWriter,
     image_array_to_pil_image,
     safe_stop_image_writer,
diff --git a/tests/datasets/test_online_buffer.py b/tests/datasets/test_online_buffer.py
index 339f68481..887da6041 100644
--- a/tests/datasets/test_online_buffer.py
+++ b/tests/datasets/test_online_buffer.py
@@ -20,7 +20,7 @@ import numpy as np
 import pytest
 import torch
 
-from lerobot.common.datasets.online_buffer import OnlineBuffer, compute_sampler_weights
+from lerobot.datasets.online_buffer import OnlineBuffer, compute_sampler_weights
 
 # Some constants for OnlineBuffer tests.
 data_key = "data"
diff --git a/tests/datasets/test_sampler.py b/tests/datasets/test_sampler.py
index ee143f376..94576a3e2 100644
--- a/tests/datasets/test_sampler.py
+++ b/tests/datasets/test_sampler.py
@@ -15,9 +15,9 @@
 # limitations under the License.
 from datasets import Dataset
 
-from lerobot.common.datasets.push_dataset_to_hub.utils import calculate_episode_data_index
-from lerobot.common.datasets.sampler import EpisodeAwareSampler
-from lerobot.common.datasets.utils import (
+from lerobot.datasets.push_dataset_to_hub.utils import calculate_episode_data_index
+from lerobot.datasets.sampler import EpisodeAwareSampler
+from lerobot.datasets.utils import (
     hf_transform_to_torch,
 )
 
diff --git a/tests/datasets/test_utils.py b/tests/datasets/test_utils.py
index 0d02218aa..ba16874d0 100644
--- a/tests/datasets/test_utils.py
+++ b/tests/datasets/test_utils.py
@@ -18,8 +18,8 @@ import torch
 from datasets import Dataset
 from huggingface_hub import DatasetCard
 
-from lerobot.common.datasets.push_dataset_to_hub.utils import calculate_episode_data_index
-from lerobot.common.datasets.utils import create_lerobot_dataset_card, hf_transform_to_torch
+from lerobot.datasets.push_dataset_to_hub.utils import calculate_episode_data_index
+from lerobot.datasets.utils import create_lerobot_dataset_card, hf_transform_to_torch
 
 
 def test_default_parameters():
diff --git a/tests/envs/test_envs.py b/tests/envs/test_envs.py
index b318abb4a..140e9dfb9 100644
--- a/tests/envs/test_envs.py
+++ b/tests/envs/test_envs.py
@@ -21,8 +21,8 @@ import torch
 from gymnasium.utils.env_checker import check_env
 
 import lerobot
-from lerobot.common.envs.factory import make_env, make_env_config
-from lerobot.common.envs.utils import preprocess_observation
+from lerobot.envs.factory import make_env, make_env_config
+from lerobot.envs.utils import preprocess_observation
 from tests.utils import require_env
 
 OBS_TYPES = ["state", "pixels", "pixels_agent_pos"]
diff --git a/tests/fixtures/constants.py b/tests/fixtures/constants.py
index 5e5c762c8..d69a4634f 100644
--- a/tests/fixtures/constants.py
+++ b/tests/fixtures/constants.py
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from lerobot.common.constants import HF_LEROBOT_HOME
+from lerobot.constants import HF_LEROBOT_HOME
 
 LEROBOT_TEST_DIR = HF_LEROBOT_HOME / "_testing"
 DUMMY_REPO_ID = "dummy/repo"
diff --git a/tests/fixtures/dataset_factories.py b/tests/fixtures/dataset_factories.py
index 531977dab..047db3393 100644
--- a/tests/fixtures/dataset_factories.py
+++ b/tests/fixtures/dataset_factories.py
@@ -23,8 +23,8 @@ import PIL.Image
 import pytest
 import torch
 
-from lerobot.common.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset, LeRobotDatasetMetadata
-from lerobot.common.datasets.utils import (
+from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset, LeRobotDatasetMetadata
+from lerobot.datasets.utils import (
     DEFAULT_CHUNK_SIZE,
     DEFAULT_FEATURES,
     DEFAULT_PARQUET_PATH,
@@ -351,10 +351,8 @@ def lerobot_dataset_metadata_factory(
             episodes=episodes,
         )
         with (
-            patch("lerobot.common.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version_patch,
-            patch(
-                "lerobot.common.datasets.lerobot_dataset.snapshot_download"
-            ) as mock_snapshot_download_patch,
+            patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version_patch,
+            patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download_patch,
         ):
             mock_get_safe_version_patch.side_effect = lambda repo_id, version: version
             mock_snapshot_download_patch.side_effect = mock_snapshot_download
@@ -428,11 +426,9 @@ def lerobot_dataset_factory(
             episodes=episode_dicts,
         )
         with (
-            patch("lerobot.common.datasets.lerobot_dataset.LeRobotDatasetMetadata") as mock_metadata_patch,
-            patch("lerobot.common.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version_patch,
-            patch(
-                "lerobot.common.datasets.lerobot_dataset.snapshot_download"
-            ) as mock_snapshot_download_patch,
+            patch("lerobot.datasets.lerobot_dataset.LeRobotDatasetMetadata") as mock_metadata_patch,
+            patch("lerobot.datasets.lerobot_dataset.get_safe_version") as mock_get_safe_version_patch,
+            patch("lerobot.datasets.lerobot_dataset.snapshot_download") as mock_snapshot_download_patch,
         ):
             mock_metadata_patch.return_value = mock_metadata
             mock_get_safe_version_patch.side_effect = lambda repo_id, version: version
diff --git a/tests/fixtures/files.py b/tests/fixtures/files.py
index 678d1f382..e0553f77e 100644
--- a/tests/fixtures/files.py
+++ b/tests/fixtures/files.py
@@ -20,7 +20,7 @@ import pyarrow.compute as pc
 import pyarrow.parquet as pq
 import pytest
 
-from lerobot.common.datasets.utils import (
+from lerobot.datasets.utils import (
     EPISODES_PATH,
     EPISODES_STATS_PATH,
     INFO_PATH,
diff --git a/tests/fixtures/hub.py b/tests/fixtures/hub.py
index aa2768e4c..f7c5f5b04 100644
--- a/tests/fixtures/hub.py
+++ b/tests/fixtures/hub.py
@@ -17,7 +17,7 @@ import datasets
 import pytest
 from huggingface_hub.utils import filter_repo_objects
 
-from lerobot.common.datasets.utils import (
+from lerobot.datasets.utils import (
     EPISODES_PATH,
     EPISODES_STATS_PATH,
     INFO_PATH,
diff --git a/tests/fixtures/optimizers.py b/tests/fixtures/optimizers.py
index 65488566e..a1b4a9da0 100644
--- a/tests/fixtures/optimizers.py
+++ b/tests/fixtures/optimizers.py
@@ -14,8 +14,8 @@
 import pytest
 import torch
 
-from lerobot.common.optim.optimizers import AdamConfig
-from lerobot.common.optim.schedulers import VQBeTSchedulerConfig
+from lerobot.optim.optimizers import AdamConfig
+from lerobot.optim.schedulers import VQBeTSchedulerConfig
 
 
 @pytest.fixture
diff --git a/tests/mocks/mock_dynamixel.py b/tests/mocks/mock_dynamixel.py
index d446bf272..64592439f 100644
--- a/tests/mocks/mock_dynamixel.py
+++ b/tests/mocks/mock_dynamixel.py
@@ -5,7 +5,7 @@ import dynamixel_sdk as dxl
 import serial
 from mock_serial.mock_serial import MockSerial
 
-from lerobot.common.motors.dynamixel.dynamixel import _split_into_byte_chunks
+from lerobot.motors.dynamixel.dynamixel import _split_into_byte_chunks
 
 from .mock_serial_patch import WaitableStub
 
diff --git a/tests/mocks/mock_feetech.py b/tests/mocks/mock_feetech.py
index 5279b1dc8..e0b677d57 100644
--- a/tests/mocks/mock_feetech.py
+++ b/tests/mocks/mock_feetech.py
@@ -5,7 +5,7 @@ import scservo_sdk as scs
 import serial
 from mock_serial import MockSerial
 
-from lerobot.common.motors.feetech.feetech import _split_into_byte_chunks, patch_setPacketTimeout
+from lerobot.motors.feetech.feetech import _split_into_byte_chunks, patch_setPacketTimeout
 
 from .mock_serial_patch import WaitableStub
 
diff --git a/tests/mocks/mock_motors_bus.py b/tests/mocks/mock_motors_bus.py
index e322eae8a..91e33473d 100644
--- a/tests/mocks/mock_motors_bus.py
+++ b/tests/mocks/mock_motors_bus.py
@@ -1,6 +1,6 @@
 # ruff: noqa: N802
 
-from lerobot.common.motors.motors_bus import (
+from lerobot.motors.motors_bus import (
     Motor,
     MotorsBus,
 )
diff --git a/tests/mocks/mock_robot.py b/tests/mocks/mock_robot.py
index 40d8fbde6..971fc00ad 100644
--- a/tests/mocks/mock_robot.py
+++ b/tests/mocks/mock_robot.py
@@ -3,9 +3,9 @@ from dataclasses import dataclass, field
 from functools import cached_property
 from typing import Any
 
-from lerobot.common.cameras import CameraConfig, make_cameras_from_configs
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
-from lerobot.common.robots import Robot, RobotConfig
+from lerobot.cameras import CameraConfig, make_cameras_from_configs
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.robots import Robot, RobotConfig
 
 
 @RobotConfig.register_subclass("mock_robot")
diff --git a/tests/mocks/mock_teleop.py b/tests/mocks/mock_teleop.py
index a7f5cad35..c29cc9219 100644
--- a/tests/mocks/mock_teleop.py
+++ b/tests/mocks/mock_teleop.py
@@ -3,8 +3,8 @@ from dataclasses import dataclass
 from functools import cached_property
 from typing import Any
 
-from lerobot.common.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
-from lerobot.common.teleoperators import Teleoperator, TeleoperatorConfig
+from lerobot.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError
+from lerobot.teleoperators import Teleoperator, TeleoperatorConfig
 
 
 @TeleoperatorConfig.register_subclass("mock_teleop")
diff --git a/tests/motors/test_dynamixel.py b/tests/motors/test_dynamixel.py
index dcce8f691..a54e49056 100644
--- a/tests/motors/test_dynamixel.py
+++ b/tests/motors/test_dynamixel.py
@@ -5,10 +5,10 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
-from lerobot.common.motors import Motor, MotorCalibration, MotorNormMode
-from lerobot.common.motors.dynamixel import MODEL_NUMBER_TABLE, DynamixelMotorsBus
-from lerobot.common.motors.dynamixel.tables import X_SERIES_CONTROL_TABLE
-from lerobot.common.utils.encoding_utils import encode_twos_complement
+from lerobot.motors import Motor, MotorCalibration, MotorNormMode
+from lerobot.motors.dynamixel import MODEL_NUMBER_TABLE, DynamixelMotorsBus
+from lerobot.motors.dynamixel.tables import X_SERIES_CONTROL_TABLE
+from lerobot.utils.encoding_utils import encode_twos_complement
 
 try:
     import dynamixel_sdk as dxl
@@ -389,7 +389,7 @@ def test_record_ranges_of_motion(mock_motors, dummy_motors):
     read_pos_stub = mock_motors.build_sequential_sync_read_stub(
         *X_SERIES_CONTROL_TABLE["Present_Position"], positions
     )
-    with patch("lerobot.common.motors.motors_bus.enter_pressed", side_effect=[False, True]):
+    with patch("lerobot.motors.motors_bus.enter_pressed", side_effect=[False, True]):
         bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
         bus.connect(handshake=False)
 
diff --git a/tests/motors/test_feetech.py b/tests/motors/test_feetech.py
index 2e7b2ff77..7f9e5dd7b 100644
--- a/tests/motors/test_feetech.py
+++ b/tests/motors/test_feetech.py
@@ -5,10 +5,10 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
-from lerobot.common.motors import Motor, MotorCalibration, MotorNormMode
-from lerobot.common.motors.feetech import MODEL_NUMBER, MODEL_NUMBER_TABLE, FeetechMotorsBus
-from lerobot.common.motors.feetech.tables import STS_SMS_SERIES_CONTROL_TABLE
-from lerobot.common.utils.encoding_utils import encode_sign_magnitude
+from lerobot.motors import Motor, MotorCalibration, MotorNormMode
+from lerobot.motors.feetech import MODEL_NUMBER, MODEL_NUMBER_TABLE, FeetechMotorsBus
+from lerobot.motors.feetech.tables import STS_SMS_SERIES_CONTROL_TABLE
+from lerobot.utils.encoding_utils import encode_sign_magnitude
 
 try:
     import scservo_sdk as scs
@@ -432,7 +432,7 @@ def test_record_ranges_of_motion(mock_motors, dummy_motors):
     stub = mock_motors.build_sequential_sync_read_stub(
         *STS_SMS_SERIES_CONTROL_TABLE["Present_Position"], positions
     )
-    with patch("lerobot.common.motors.motors_bus.enter_pressed", side_effect=[False, True]):
+    with patch("lerobot.motors.motors_bus.enter_pressed", side_effect=[False, True]):
         bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
         bus.connect(handshake=False)
 
diff --git a/tests/motors/test_motors_bus.py b/tests/motors/test_motors_bus.py
index 78b7a47da..966af3fb0 100644
--- a/tests/motors/test_motors_bus.py
+++ b/tests/motors/test_motors_bus.py
@@ -3,7 +3,7 @@ from unittest.mock import patch
 
 import pytest
 
-from lerobot.common.motors.motors_bus import (
+from lerobot.motors.motors_bus import (
     Motor,
     MotorNormMode,
     assert_same_address,
diff --git a/tests/optim/test_optimizers.py b/tests/optim/test_optimizers.py
index 630353fca..4152c7f8d 100644
--- a/tests/optim/test_optimizers.py
+++ b/tests/optim/test_optimizers.py
@@ -14,11 +14,11 @@
 import pytest
 import torch
 
-from lerobot.common.constants import (
+from lerobot.constants import (
     OPTIMIZER_PARAM_GROUPS,
     OPTIMIZER_STATE,
 )
-from lerobot.common.optim.optimizers import (
+from lerobot.optim.optimizers import (
     AdamConfig,
     AdamWConfig,
     MultiAdamConfig,
diff --git a/tests/optim/test_schedulers.py b/tests/optim/test_schedulers.py
index d6191dcea..43851c458 100644
--- a/tests/optim/test_schedulers.py
+++ b/tests/optim/test_schedulers.py
@@ -13,8 +13,8 @@
 # limitations under the License.
 from torch.optim.lr_scheduler import LambdaLR
 
-from lerobot.common.constants import SCHEDULER_STATE
-from lerobot.common.optim.schedulers import (
+from lerobot.constants import SCHEDULER_STATE
+from lerobot.optim.schedulers import (
     CosineDecayWithWarmupSchedulerConfig,
     DiffuserSchedulerConfig,
     VQBeTSchedulerConfig,
diff --git a/tests/policies/hilserl/test_modeling_classifier.py b/tests/policies/hilserl/test_modeling_classifier.py
index 526e1f17d..0be1b9c7c 100644
--- a/tests/policies/hilserl/test_modeling_classifier.py
+++ b/tests/policies/hilserl/test_modeling_classifier.py
@@ -16,9 +16,9 @@
 
 import torch
 
-from lerobot.common.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
-from lerobot.common.policies.sac.reward_model.modeling_classifier import ClassifierOutput
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
+from lerobot.policies.sac.reward_model.modeling_classifier import ClassifierOutput
 from tests.utils import require_package
 
 
@@ -37,7 +37,7 @@ def test_classifier_output():
 
 @require_package("transformers")
 def test_binary_classifier_with_default_params():
-    from lerobot.common.policies.sac.reward_model.modeling_classifier import Classifier
+    from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
 
     config = RewardClassifierConfig()
     config.input_features = {
@@ -78,7 +78,7 @@ def test_binary_classifier_with_default_params():
 
 @require_package("transformers")
 def test_multiclass_classifier():
-    from lerobot.common.policies.sac.reward_model.modeling_classifier import Classifier
+    from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
 
     num_classes = 5
     config = RewardClassifierConfig()
@@ -117,7 +117,7 @@ def test_multiclass_classifier():
 
 @require_package("transformers")
 def test_default_device():
-    from lerobot.common.policies.sac.reward_model.modeling_classifier import Classifier
+    from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
 
     config = RewardClassifierConfig()
     assert config.device == "cpu"
@@ -129,7 +129,7 @@ def test_default_device():
 
 @require_package("transformers")
 def test_explicit_device_setup():
-    from lerobot.common.policies.sac.reward_model.modeling_classifier import Classifier
+    from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
 
     config = RewardClassifierConfig(device="cpu")
     assert config.device == "cpu"
diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py
index dff5975ae..ed37fedd6 100644
--- a/tests/policies/test_policies.py
+++ b/tests/policies/test_policies.py
@@ -24,23 +24,23 @@ from packaging import version
 from safetensors.torch import load_file
 
 from lerobot import available_policies
-from lerobot.common.datasets.factory import make_dataset
-from lerobot.common.datasets.utils import cycle, dataset_to_policy_features
-from lerobot.common.envs.factory import make_env, make_env_config
-from lerobot.common.envs.utils import preprocess_observation
-from lerobot.common.optim.factory import make_optimizer_and_scheduler
-from lerobot.common.policies.act.modeling_act import ACTTemporalEnsembler
-from lerobot.common.policies.factory import (
+from lerobot.configs.default import DatasetConfig
+from lerobot.configs.train import TrainPipelineConfig
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.datasets.factory import make_dataset
+from lerobot.datasets.utils import cycle, dataset_to_policy_features
+from lerobot.envs.factory import make_env, make_env_config
+from lerobot.envs.utils import preprocess_observation
+from lerobot.optim.factory import make_optimizer_and_scheduler
+from lerobot.policies.act.modeling_act import ACTTemporalEnsembler
+from lerobot.policies.factory import (
     get_policy_class,
     make_policy,
     make_policy_config,
 )
-from lerobot.common.policies.normalize import Normalize, Unnormalize
-from lerobot.common.policies.pretrained import PreTrainedPolicy
-from lerobot.common.utils.random_utils import seeded_context
-from lerobot.configs.default import DatasetConfig
-from lerobot.configs.train import TrainPipelineConfig
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.policies.normalize import Normalize, Unnormalize
+from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.utils.random_utils import seeded_context
 from tests.artifacts.policies.save_policy_to_safetensors import get_policy_stats
 from tests.utils import DEVICE, require_cpu, require_env, require_x86_64_kernel
 
@@ -142,9 +142,10 @@ def test_policy(ds_repo_id, env_name, env_kwargs, policy_name, policy_kwargs):
     train_cfg = TrainPipelineConfig(
         # TODO(rcadene, aliberts): remove dataset download
         dataset=DatasetConfig(repo_id=ds_repo_id, episodes=[0]),
-        policy=make_policy_config(policy_name, **policy_kwargs),
+        policy=make_policy_config(policy_name, push_to_hub=False, **policy_kwargs),
         env=make_env_config(env_name, **env_kwargs),
     )
+    train_cfg.validate()
 
     # Check that we can make the policy object.
     dataset = make_dataset(train_cfg)
@@ -213,7 +214,7 @@ def test_act_backbone_lr():
     cfg = TrainPipelineConfig(
         # TODO(rcadene, aliberts): remove dataset download
         dataset=DatasetConfig(repo_id="lerobot/aloha_sim_insertion_scripted", episodes=[0]),
-        policy=make_policy_config("act", optimizer_lr=0.01, optimizer_lr_backbone=0.001),
+        policy=make_policy_config("act", optimizer_lr=0.01, optimizer_lr_backbone=0.001, push_to_hub=False),
     )
     cfg.validate()  # Needed for auto-setting some parameters
 
@@ -415,6 +416,7 @@ def test_backward_compatibility(ds_repo_id: str, policy_name: str, policy_kwargs
     https://github.com/huggingface/lerobot/pull/1127.
 
     """
+
     # NOTE: ACT policy has different randomness, after PyTorch 2.7.0
     if policy_name == "act" and version.parse(torch.__version__) < version.parse("2.7.0"):
         pytest.skip(f"Skipping act policy test with PyTorch {torch.__version__}. Requires PyTorch >= 2.7.0")
diff --git a/tests/policies/test_sac_config.py b/tests/policies/test_sac_config.py
index d94ee41e0..a67815eed 100644
--- a/tests/policies/test_sac_config.py
+++ b/tests/policies/test_sac_config.py
@@ -16,7 +16,8 @@
 
 import pytest
 
-from lerobot.common.policies.sac.configuration_sac import (
+from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.policies.sac.configuration_sac import (
     ActorLearnerConfig,
     ActorNetworkConfig,
     ConcurrencyConfig,
@@ -24,7 +25,6 @@ from lerobot.common.policies.sac.configuration_sac import (
     PolicyConfig,
     SACConfig,
 )
-from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 
 
 def test_sac_config_default_initialization():
diff --git a/tests/policies/test_sac_policy.py b/tests/policies/test_sac_policy.py
index e4e2dd8a9..7891c2e52 100644
--- a/tests/policies/test_sac_policy.py
+++ b/tests/policies/test_sac_policy.py
@@ -20,10 +20,10 @@ import pytest
 import torch
 from torch import Tensor, nn
 
-from lerobot.common.policies.sac.configuration_sac import SACConfig
-from lerobot.common.policies.sac.modeling_sac import MLP, SACPolicy
-from lerobot.common.utils.random_utils import seeded_context, set_seed
 from lerobot.configs.types import FeatureType, PolicyFeature
+from lerobot.policies.sac.configuration_sac import SACConfig
+from lerobot.policies.sac.modeling_sac import MLP, SACPolicy
+from lerobot.utils.random_utils import seeded_context, set_seed
 
 try:
     import transformers  # noqa: F401
diff --git a/tests/rl/test_actor.py b/tests/rl/test_actor.py
index 0cf6a8f64..f078b4602 100644
--- a/tests/rl/test_actor.py
+++ b/tests/rl/test_actor.py
@@ -21,14 +21,14 @@ import pytest
 import torch
 from torch.multiprocessing import Event, Queue
 
-from lerobot.common.utils.transition import Transition
+from lerobot.utils.transition import Transition
 from tests.utils import require_package
 
 
 def create_learner_service_stub():
     import grpc
 
-    from lerobot.common.transport import services_pb2, services_pb2_grpc
+    from lerobot.transport import services_pb2, services_pb2_grpc
 
     class MockLearnerService(services_pb2_grpc.LearnerServiceServicer):
         def __init__(self):
@@ -101,8 +101,8 @@ def test_establish_learner_connection_failure():
 
 @require_package("grpc")
 def test_push_transitions_to_transport_queue():
-    from lerobot.common.transport.utils import bytes_to_transitions
     from lerobot.scripts.rl.actor import push_transitions_to_transport_queue
+    from lerobot.transport.utils import bytes_to_transitions
     from tests.transport.test_transport_utils import assert_transitions_equal
 
     """Test pushing transitions to transport queue."""
@@ -169,8 +169,8 @@ def test_transitions_stream():
 @require_package("grpc")
 @pytest.mark.timeout(3)  # force cross-platform watchdog
 def test_interactions_stream():
-    from lerobot.common.transport.utils import bytes_to_python_object, python_object_to_bytes
     from lerobot.scripts.rl.actor import interactions_stream
+    from lerobot.transport.utils import bytes_to_python_object, python_object_to_bytes
 
     """Test interactions stream functionality."""
     shutdown_event = Event()
diff --git a/tests/rl/test_actor_learner.py b/tests/rl/test_actor_learner.py
index cb72da7e4..b2a7a5d5f 100644
--- a/tests/rl/test_actor_learner.py
+++ b/tests/rl/test_actor_learner.py
@@ -22,9 +22,9 @@ import pytest
 import torch
 from torch.multiprocessing import Event, Queue
 
-from lerobot.common.policies.sac.configuration_sac import SACConfig
-from lerobot.common.utils.transition import Transition
 from lerobot.configs.train import TrainRLServerPipelineConfig
+from lerobot.policies.sac.configuration_sac import SACConfig
+from lerobot.utils.transition import Transition
 from tests.utils import require_package
 
 
@@ -90,7 +90,6 @@ def cfg():
 @require_package("grpc")
 @pytest.mark.timeout(10)  # force cross-platform watchdog
 def test_end_to_end_transitions_flow(cfg):
-    from lerobot.common.transport.utils import bytes_to_transitions
     from lerobot.scripts.rl.actor import (
         establish_learner_connection,
         learner_service_client,
@@ -98,6 +97,7 @@ def test_end_to_end_transitions_flow(cfg):
         send_transitions,
     )
     from lerobot.scripts.rl.learner import start_learner
+    from lerobot.transport.utils import bytes_to_transitions
     from tests.transport.test_transport_utils import assert_transitions_equal
 
     """Test complete transitions flow from actor to learner."""
@@ -152,13 +152,13 @@ def test_end_to_end_transitions_flow(cfg):
 @require_package("grpc")
 @pytest.mark.timeout(10)
 def test_end_to_end_interactions_flow(cfg):
-    from lerobot.common.transport.utils import bytes_to_python_object, python_object_to_bytes
     from lerobot.scripts.rl.actor import (
         establish_learner_connection,
         learner_service_client,
         send_interactions,
     )
     from lerobot.scripts.rl.learner import start_learner
+    from lerobot.transport.utils import bytes_to_python_object, python_object_to_bytes
 
     """Test complete interactions flow from actor to learner."""
     # Queues for actor-learner communication
@@ -226,9 +226,9 @@ def test_end_to_end_interactions_flow(cfg):
 @pytest.mark.parametrize("data_size", ["small", "large"])
 @pytest.mark.timeout(10)
 def test_end_to_end_parameters_flow(cfg, data_size):
-    from lerobot.common.transport.utils import bytes_to_state_dict, state_to_bytes
     from lerobot.scripts.rl.actor import establish_learner_connection, learner_service_client, receive_policy
     from lerobot.scripts.rl.learner import start_learner
+    from lerobot.transport.utils import bytes_to_state_dict, state_to_bytes
 
     """Test complete parameter flow from learner to actor, with small and large data."""
     # Actor's local queue to receive params
diff --git a/tests/rl/test_learner_service.py b/tests/rl/test_learner_service.py
index ee9d06e91..f5e1e8d48 100644
--- a/tests/rl/test_learner_service.py
+++ b/tests/rl/test_learner_service.py
@@ -50,8 +50,8 @@ def create_learner_service_stub(
 ):
     import grpc
 
-    from lerobot.common.transport import services_pb2_grpc  # generated from .proto
     from lerobot.scripts.rl.learner_service import LearnerService
+    from lerobot.transport import services_pb2_grpc  # generated from .proto
 
     """Fixture to start a LearnerService gRPC server and provide a connected stub."""
 
@@ -83,7 +83,7 @@ def close_learner_service_stub(channel, server):
 
 @pytest.mark.timeout(3)  # force cross-platform watchdog
 def test_ready_method(learner_service_stub):
-    from lerobot.common.transport import services_pb2
+    from lerobot.transport import services_pb2
 
     """Test the ready method of the UserService."""
     request = services_pb2.Empty()
@@ -94,7 +94,7 @@ def test_ready_method(learner_service_stub):
 @require_package("grpc")
 @pytest.mark.timeout(3)  # force cross-platform watchdog
 def test_send_interactions():
-    from lerobot.common.transport import services_pb2
+    from lerobot.transport import services_pb2
 
     shutdown_event = Event()
 
@@ -138,7 +138,7 @@ def test_send_interactions():
 @require_package("grpc")
 @pytest.mark.timeout(3)  # force cross-platform watchdog
 def test_send_transitions():
-    from lerobot.common.transport import services_pb2
+    from lerobot.transport import services_pb2
 
     """Test the SendTransitions method with various transition data."""
     shutdown_event = Event()
@@ -184,7 +184,7 @@ def test_send_transitions():
 @require_package("grpc")
 @pytest.mark.timeout(3)  # force cross-platform watchdog
 def test_send_transitions_empty_stream():
-    from lerobot.common.transport import services_pb2
+    from lerobot.transport import services_pb2
 
     """Test SendTransitions with empty stream."""
     shutdown_event = Event()
@@ -214,7 +214,7 @@ def test_send_transitions_empty_stream():
 def test_stream_parameters():
     import time
 
-    from lerobot.common.transport import services_pb2
+    from lerobot.transport import services_pb2
 
     """Test the StreamParameters method."""
     shutdown_event = Event()
@@ -270,7 +270,7 @@ def test_stream_parameters():
 @require_package("grpc")
 @pytest.mark.timeout(3)  # force cross-platform watchdog
 def test_stream_parameters_with_shutdown():
-    from lerobot.common.transport import services_pb2
+    from lerobot.transport import services_pb2
 
     """Test StreamParameters handles shutdown gracefully."""
     shutdown_event = Event()
@@ -325,7 +325,7 @@ def test_stream_parameters_waits_and_retries_on_empty_queue():
     import threading
     import time
 
-    from lerobot.common.transport import services_pb2
+    from lerobot.transport import services_pb2
 
     """Test that StreamParameters waits and retries when the queue is empty."""
     shutdown_event = Event()
diff --git a/tests/robots/test_so100_follower.py b/tests/robots/test_so100_follower.py
index 81d9d6a91..498eec94b 100644
--- a/tests/robots/test_so100_follower.py
+++ b/tests/robots/test_so100_follower.py
@@ -3,7 +3,7 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
-from lerobot.common.robots.so100_follower import (
+from lerobot.robots.so100_follower import (
     SO100Follower,
     SO100FollowerConfig,
 )
@@ -50,7 +50,7 @@ def follower():
 
     with (
         patch(
-            "lerobot.common.robots.so100_follower.so100_follower.FeetechMotorsBus",
+            "lerobot.robots.so100_follower.so100_follower.FeetechMotorsBus",
             side_effect=_bus_side_effect,
         ),
         patch.object(SO100Follower, "configure", lambda self: None),
diff --git a/tests/test_available.py b/tests/test_available.py
index a18b95ffa..19e39b2b6 100644
--- a/tests/test_available.py
+++ b/tests/test_available.py
@@ -19,10 +19,10 @@ import gymnasium as gym
 import pytest
 
 import lerobot
-from lerobot.common.policies.act.modeling_act import ACTPolicy
-from lerobot.common.policies.diffusion.modeling_diffusion import DiffusionPolicy
-from lerobot.common.policies.tdmpc.modeling_tdmpc import TDMPCPolicy
-from lerobot.common.policies.vqbet.modeling_vqbet import VQBeTPolicy
+from lerobot.policies.act.modeling_act import ACTPolicy
+from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
+from lerobot.policies.tdmpc.modeling_tdmpc import TDMPCPolicy
+from lerobot.policies.vqbet.modeling_vqbet import VQBeTPolicy
 from tests.utils import require_env
 
 
diff --git a/tests/transport/test_transport_utils.py b/tests/transport/test_transport_utils.py
index cf33f52c0..79edad4e4 100644
--- a/tests/transport/test_transport_utils.py
+++ b/tests/transport/test_transport_utils.py
@@ -21,13 +21,13 @@ from pickle import UnpicklingError
 import pytest
 import torch
 
-from lerobot.common.utils.transition import Transition
+from lerobot.utils.transition import Transition
 from tests.utils import require_cuda, require_package
 
 
 @require_package("grpc")
 def test_bytes_buffer_size_empty_buffer():
-    from lerobot.common.transport.utils import bytes_buffer_size
+    from lerobot.transport.utils import bytes_buffer_size
 
     """Test with an empty buffer."""
     buffer = io.BytesIO()
@@ -38,7 +38,7 @@ def test_bytes_buffer_size_empty_buffer():
 
 @require_package("grpc")
 def test_bytes_buffer_size_small_buffer():
-    from lerobot.common.transport.utils import bytes_buffer_size
+    from lerobot.transport.utils import bytes_buffer_size
 
     """Test with a small buffer."""
     buffer = io.BytesIO(b"Hello, World!")
@@ -48,7 +48,7 @@ def test_bytes_buffer_size_small_buffer():
 
 @require_package("grpc")
 def test_bytes_buffer_size_large_buffer():
-    from lerobot.common.transport.utils import CHUNK_SIZE, bytes_buffer_size
+    from lerobot.transport.utils import CHUNK_SIZE, bytes_buffer_size
 
     """Test with a large buffer."""
     data = b"x" * (CHUNK_SIZE * 2 + 1000)
@@ -59,7 +59,7 @@ def test_bytes_buffer_size_large_buffer():
 
 @require_package("grpc")
 def test_send_bytes_in_chunks_empty_data():
-    from lerobot.common.transport.utils import send_bytes_in_chunks, services_pb2
+    from lerobot.transport.utils import send_bytes_in_chunks, services_pb2
 
     """Test sending empty data."""
     message_class = services_pb2.InteractionMessage
@@ -69,7 +69,7 @@ def test_send_bytes_in_chunks_empty_data():
 
 @require_package("grpc")
 def test_single_chunk_small_data():
-    from lerobot.common.transport.utils import send_bytes_in_chunks, services_pb2
+    from lerobot.transport.utils import send_bytes_in_chunks, services_pb2
 
     """Test data that fits in a single chunk."""
     data = b"Some data"
@@ -83,7 +83,7 @@ def test_single_chunk_small_data():
 
 @require_package("grpc")
 def test_not_silent_mode():
-    from lerobot.common.transport.utils import send_bytes_in_chunks, services_pb2
+    from lerobot.transport.utils import send_bytes_in_chunks, services_pb2
 
     """Test not silent mode."""
     data = b"Some data"
@@ -95,7 +95,7 @@ def test_not_silent_mode():
 
 @require_package("grpc")
 def test_send_bytes_in_chunks_large_data():
-    from lerobot.common.transport.utils import CHUNK_SIZE, send_bytes_in_chunks, services_pb2
+    from lerobot.transport.utils import CHUNK_SIZE, send_bytes_in_chunks, services_pb2
 
     """Test sending large data."""
     data = b"x" * (CHUNK_SIZE * 2 + 1000)
@@ -112,7 +112,7 @@ def test_send_bytes_in_chunks_large_data():
 
 @require_package("grpc")
 def test_send_bytes_in_chunks_large_data_with_exact_chunk_size():
-    from lerobot.common.transport.utils import CHUNK_SIZE, send_bytes_in_chunks, services_pb2
+    from lerobot.transport.utils import CHUNK_SIZE, send_bytes_in_chunks, services_pb2
 
     """Test sending large data with exact chunk size."""
     data = b"x" * CHUNK_SIZE
@@ -125,7 +125,7 @@ def test_send_bytes_in_chunks_large_data_with_exact_chunk_size():
 
 @require_package("grpc")
 def test_receive_bytes_in_chunks_empty_data():
-    from lerobot.common.transport.utils import receive_bytes_in_chunks
+    from lerobot.transport.utils import receive_bytes_in_chunks
 
     """Test receiving empty data."""
     queue = Queue()
@@ -139,7 +139,7 @@ def test_receive_bytes_in_chunks_empty_data():
 
 @require_package("grpc")
 def test_receive_bytes_in_chunks_single_chunk():
-    from lerobot.common.transport.utils import receive_bytes_in_chunks, services_pb2
+    from lerobot.transport.utils import receive_bytes_in_chunks, services_pb2
 
     """Test receiving a single chunk message."""
     queue = Queue()
@@ -158,7 +158,7 @@ def test_receive_bytes_in_chunks_single_chunk():
 
 @require_package("grpc")
 def test_receive_bytes_in_chunks_single_not_end_chunk():
-    from lerobot.common.transport.utils import receive_bytes_in_chunks, services_pb2
+    from lerobot.transport.utils import receive_bytes_in_chunks, services_pb2
 
     """Test receiving a single chunk message."""
     queue = Queue()
@@ -176,7 +176,7 @@ def test_receive_bytes_in_chunks_single_not_end_chunk():
 
 @require_package("grpc")
 def test_receive_bytes_in_chunks_multiple_chunks():
-    from lerobot.common.transport.utils import receive_bytes_in_chunks, services_pb2
+    from lerobot.transport.utils import receive_bytes_in_chunks, services_pb2
 
     """Test receiving a multi-chunk message."""
     queue = Queue()
@@ -200,7 +200,7 @@ def test_receive_bytes_in_chunks_multiple_chunks():
 
 @require_package("grpc")
 def test_receive_bytes_in_chunks_multiple_messages():
-    from lerobot.common.transport.utils import receive_bytes_in_chunks, services_pb2
+    from lerobot.transport.utils import receive_bytes_in_chunks, services_pb2
 
     """Test receiving multiple complete messages in sequence."""
     queue = Queue()
@@ -236,7 +236,7 @@ def test_receive_bytes_in_chunks_multiple_messages():
 
 @require_package("grpc")
 def test_receive_bytes_in_chunks_shutdown_during_receive():
-    from lerobot.common.transport.utils import receive_bytes_in_chunks, services_pb2
+    from lerobot.transport.utils import receive_bytes_in_chunks, services_pb2
 
     """Test that shutdown event stops receiving mid-stream."""
     queue = Queue()
@@ -260,7 +260,7 @@ def test_receive_bytes_in_chunks_shutdown_during_receive():
 
 @require_package("grpc")
 def test_receive_bytes_in_chunks_only_begin_chunk():
-    from lerobot.common.transport.utils import receive_bytes_in_chunks, services_pb2
+    from lerobot.transport.utils import receive_bytes_in_chunks, services_pb2
 
     """Test receiving only a BEGIN chunk without END."""
     queue = Queue()
@@ -280,7 +280,7 @@ def test_receive_bytes_in_chunks_only_begin_chunk():
 
 @require_package("grpc")
 def test_receive_bytes_in_chunks_missing_begin():
-    from lerobot.common.transport.utils import receive_bytes_in_chunks, services_pb2
+    from lerobot.transport.utils import receive_bytes_in_chunks, services_pb2
 
     """Test receiving chunks starting with MIDDLE instead of BEGIN."""
     queue = Queue()
@@ -304,7 +304,7 @@ def test_receive_bytes_in_chunks_missing_begin():
 # Tests for state_to_bytes and bytes_to_state_dict
 @require_package("grpc")
 def test_state_to_bytes_empty_dict():
-    from lerobot.common.transport.utils import bytes_to_state_dict, state_to_bytes
+    from lerobot.transport.utils import bytes_to_state_dict, state_to_bytes
 
     """Test converting empty state dict to bytes."""
     state_dict = {}
@@ -315,7 +315,7 @@ def test_state_to_bytes_empty_dict():
 
 @require_package("grpc")
 def test_bytes_to_state_dict_empty_data():
-    from lerobot.common.transport.utils import bytes_to_state_dict
+    from lerobot.transport.utils import bytes_to_state_dict
 
     """Test converting empty data to state dict."""
     with pytest.raises(EOFError):
@@ -324,7 +324,7 @@ def test_bytes_to_state_dict_empty_data():
 
 @require_package("grpc")
 def test_state_to_bytes_simple_dict():
-    from lerobot.common.transport.utils import bytes_to_state_dict, state_to_bytes
+    from lerobot.transport.utils import bytes_to_state_dict, state_to_bytes
 
     """Test converting simple state dict to bytes."""
     state_dict = {
@@ -348,7 +348,7 @@ def test_state_to_bytes_simple_dict():
 
 @require_package("grpc")
 def test_state_to_bytes_various_dtypes():
-    from lerobot.common.transport.utils import bytes_to_state_dict, state_to_bytes
+    from lerobot.transport.utils import bytes_to_state_dict, state_to_bytes
 
     """Test converting state dict with various tensor dtypes."""
     state_dict = {
@@ -373,7 +373,7 @@ def test_state_to_bytes_various_dtypes():
 
 @require_package("grpc")
 def test_bytes_to_state_dict_invalid_data():
-    from lerobot.common.transport.utils import bytes_to_state_dict
+    from lerobot.transport.utils import bytes_to_state_dict
 
     """Test bytes_to_state_dict with invalid data."""
     with pytest.raises(UnpicklingError):
@@ -383,7 +383,7 @@ def test_bytes_to_state_dict_invalid_data():
 @require_cuda
 @require_package("grpc")
 def test_state_to_bytes_various_dtypes_cuda():
-    from lerobot.common.transport.utils import bytes_to_state_dict, state_to_bytes
+    from lerobot.transport.utils import bytes_to_state_dict, state_to_bytes
 
     """Test converting state dict with various tensor dtypes."""
     state_dict = {
@@ -408,7 +408,7 @@ def test_state_to_bytes_various_dtypes_cuda():
 
 @require_package("grpc")
 def test_python_object_to_bytes_none():
-    from lerobot.common.transport.utils import bytes_to_python_object, python_object_to_bytes
+    from lerobot.transport.utils import bytes_to_python_object, python_object_to_bytes
 
     """Test converting None to bytes."""
     obj = None
@@ -440,7 +440,7 @@ def test_python_object_to_bytes_none():
 )
 @require_package("grpc")
 def test_python_object_to_bytes_simple_types(obj):
-    from lerobot.common.transport.utils import bytes_to_python_object, python_object_to_bytes
+    from lerobot.transport.utils import bytes_to_python_object, python_object_to_bytes
 
     """Test converting simple Python types."""
     data = python_object_to_bytes(obj)
@@ -451,7 +451,7 @@ def test_python_object_to_bytes_simple_types(obj):
 
 @require_package("grpc")
 def test_python_object_to_bytes_with_tensors():
-    from lerobot.common.transport.utils import bytes_to_python_object, python_object_to_bytes
+    from lerobot.transport.utils import bytes_to_python_object, python_object_to_bytes
 
     """Test converting objects containing PyTorch tensors."""
     obj = {
@@ -476,7 +476,7 @@ def test_python_object_to_bytes_with_tensors():
 
 @require_package("grpc")
 def test_transitions_to_bytes_empty_list():
-    from lerobot.common.transport.utils import bytes_to_transitions, transitions_to_bytes
+    from lerobot.transport.utils import bytes_to_transitions, transitions_to_bytes
 
     """Test converting empty transitions list."""
     transitions = []
@@ -488,7 +488,7 @@ def test_transitions_to_bytes_empty_list():
 
 @require_package("grpc")
 def test_transitions_to_bytes_single_transition():
-    from lerobot.common.transport.utils import bytes_to_transitions, transitions_to_bytes
+    from lerobot.transport.utils import bytes_to_transitions, transitions_to_bytes
 
     """Test converting a single transition."""
     transition = Transition(
@@ -528,7 +528,7 @@ def assert_observation_equal(o1: dict, o2: dict):
 
 @require_package("grpc")
 def test_transitions_to_bytes_multiple_transitions():
-    from lerobot.common.transport.utils import bytes_to_transitions, transitions_to_bytes
+    from lerobot.transport.utils import bytes_to_transitions, transitions_to_bytes
 
     """Test converting multiple transitions."""
     transitions = []
@@ -552,7 +552,7 @@ def test_transitions_to_bytes_multiple_transitions():
 
 @require_package("grpc")
 def test_receive_bytes_in_chunks_unknown_state():
-    from lerobot.common.transport.utils import receive_bytes_in_chunks
+    from lerobot.transport.utils import receive_bytes_in_chunks
 
     """Test receive_bytes_in_chunks with an unknown transfer state."""
 
diff --git a/tests/utils.py b/tests/utils.py
index 103b973fb..800b7d4b3 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -21,7 +21,7 @@ import pytest
 import torch
 
 from lerobot import available_cameras, available_motors, available_robots
-from lerobot.common.utils.import_utils import is_package_available
+from lerobot.utils.import_utils import is_package_available
 
 DEVICE = os.environ.get("LEROBOT_TEST_DEVICE", "cuda") if torch.cuda.is_available() else "cpu"
 
diff --git a/tests/utils/test_encoding_utils.py b/tests/utils/test_encoding_utils.py
index c21c8081e..9c9796762 100644
--- a/tests/utils/test_encoding_utils.py
+++ b/tests/utils/test_encoding_utils.py
@@ -1,6 +1,6 @@
 import pytest
 
-from lerobot.common.utils.encoding_utils import (
+from lerobot.utils.encoding_utils import (
     decode_sign_magnitude,
     decode_twos_complement,
     encode_sign_magnitude,
diff --git a/tests/utils/test_io_utils.py b/tests/utils/test_io_utils.py
index c1b776db6..9768a5ef9 100644
--- a/tests/utils/test_io_utils.py
+++ b/tests/utils/test_io_utils.py
@@ -17,7 +17,7 @@ from typing import Any
 
 import pytest
 
-from lerobot.common.utils.io_utils import deserialize_json_into_object
+from lerobot.utils.io_utils import deserialize_json_into_object
 
 
 @pytest.fixture
diff --git a/tests/utils/test_logging_utils.py b/tests/utils/test_logging_utils.py
index 1ba1829e8..927fdc14d 100644
--- a/tests/utils/test_logging_utils.py
+++ b/tests/utils/test_logging_utils.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 import pytest
 
-from lerobot.common.utils.logging_utils import AverageMeter, MetricsTracker
+from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
 
 
 @pytest.fixture
diff --git a/tests/utils/test_process.py b/tests/utils/test_process.py
index 054a8593a..61e6e2c73 100644
--- a/tests/utils/test_process.py
+++ b/tests/utils/test_process.py
@@ -22,7 +22,7 @@ from unittest.mock import patch
 
 import pytest
 
-from lerobot.common.utils.process import ProcessSignalHandler
+from lerobot.utils.process import ProcessSignalHandler
 
 
 # Fixture to reset shutdown_event_counter and original signal handlers before and after each test
diff --git a/tests/utils/test_queue.py b/tests/utils/test_queue.py
index 863231e82..0a0d21770 100644
--- a/tests/utils/test_queue.py
+++ b/tests/utils/test_queue.py
@@ -18,7 +18,7 @@ import threading
 import time
 from queue import Queue
 
-from lerobot.common.utils.queue import get_last_item_from_queue
+from lerobot.utils.queue import get_last_item_from_queue
 
 
 def test_get_last_item_single_item():
diff --git a/tests/utils/test_random_utils.py b/tests/utils/test_random_utils.py
index daf08a89f..5865361d0 100644
--- a/tests/utils/test_random_utils.py
+++ b/tests/utils/test_random_utils.py
@@ -17,7 +17,7 @@ import numpy as np
 import pytest
 import torch
 
-from lerobot.common.utils.random_utils import (
+from lerobot.utils.random_utils import (
     deserialize_numpy_rng_state,
     deserialize_python_rng_state,
     deserialize_rng_state,
diff --git a/tests/utils/test_replay_buffer.py b/tests/utils/test_replay_buffer.py
index f7a055b20..260276032 100644
--- a/tests/utils/test_replay_buffer.py
+++ b/tests/utils/test_replay_buffer.py
@@ -20,8 +20,8 @@ from typing import Callable
 import pytest
 import torch
 
-from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.common.utils.buffer import BatchTransition, ReplayBuffer, random_crop_vectorized
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.utils.buffer import BatchTransition, ReplayBuffer, random_crop_vectorized
 from tests.fixtures.constants import DUMMY_REPO_ID
 
 
diff --git a/tests/utils/test_train_utils.py b/tests/utils/test_train_utils.py
index b78f6e497..2d963d7ae 100644
--- a/tests/utils/test_train_utils.py
+++ b/tests/utils/test_train_utils.py
@@ -14,7 +14,7 @@
 from pathlib import Path
 from unittest.mock import Mock, patch
 
-from lerobot.common.constants import (
+from lerobot.constants import (
     CHECKPOINTS_DIR,
     LAST_CHECKPOINT_LINK,
     OPTIMIZER_PARAM_GROUPS,
@@ -24,7 +24,7 @@ from lerobot.common.constants import (
     TRAINING_STATE_DIR,
     TRAINING_STEP,
 )
-from lerobot.common.utils.train_utils import (
+from lerobot.utils.train_utils import (
     get_step_checkpoint_dir,
     get_step_identifier,
     load_training_state,
@@ -69,7 +69,7 @@ def test_update_last_checkpoint(tmp_path):
     assert last_checkpoint.resolve() == checkpoint
 
 
-@patch("lerobot.common.utils.train_utils.save_training_state")
+@patch("lerobot.utils.train_utils.save_training_state")
 def test_save_checkpoint(mock_save_training_state, tmp_path, optimizer):
     policy = Mock()
     cfg = Mock()