feat(pi05): compose relative poses in SE3

examples/dataset/slurm_recompute_stats.py

@@ -1,489 +0,0 @@
-#!/usr/bin/env python
-
-# 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.
-
-"""
-SLURM-distributed recomputation of a LeRobotDataset's ``meta/stats.json``.
-
-Modified copy of lerobot's examples/dataset/slurm_recompute_stats.py
-(feat/recompute-stats-readonly-and-visual branch) with cluster-friendly additions:
-
-  1. --qos              : pass a SLURM QoS through to every worker's sbatch.
-  2. --venv-path        : activate a venv on each worker before the python step.
-  3. --env-command      : raw shell snippet injected before the python step (e.g. to
-                          export HF_LEROBOT_HOME). Runs in addition to --venv-path.
-  4. --chain-aggregate  : submit ``aggregate`` with an afterok dependency on
-                          ``compute`` so it only runs once all shards exist
-                          (no manual squeue-wait, no gap/overlap race).
-  5. --update-episode-stats : in ``aggregate``, also rewrite the per-episode stats in the
-                          episodes parquet so they stay consistent with meta/stats.json
-                          (default: only stats.json is written).
-
-Data access: no filesystem mount. Point HF_LEROBOT_HOME at a node-visible shared
-cache (e.g. /fsx/$USER/.cache) so the dataset downloads once and all workers read
-it. This is the download route; the source dataset is fetched from the Hub on the
-CPU workers.
-
-IMPORTANT — how to run (do NOT sbatch this file):
-  Run it as a normal python process on the LOGIN node. datatrove submits the
-  workers for you. The reference copy (--new-root) is built on the login node and
-  references the shared HF cache, so /fsx must be visible there (it is).
-
-Requires: pip install 'lerobot[dataset]' datatrove
-
-Example (single command, compute then dependent aggregate):
-
-    export HF_LEROBOT_HOME=/fsx/$USER/.cache
-
-    python slurm_recompute_stats_patched.py compute \
-        --repo-id behavior-1k/2026-challenge-demos \
-        --new-root /fsx/$USER/behavior-1k_recomputed \
-        --shard-dir /fsx/$USER/behavior-1k_recomputed/stats_shards \
-        --logs-dir /fsx/$USER/logs/recompute \
-        --skip-image-video 0 \
-        --workers 250 \
-        --partition hopper-cpu \
-        --qos normal \
-        --cpus-per-task 8 --mem-per-cpu 4G \
-        --venv-path /fsx/$USER/venvs/lerobot/bin/activate \
-        --env-command 'export HF_LEROBOT_HOME=/fsx/'"$USER"'/.cache' \
-        --chain-aggregate
-
-REHEARSE FIRST with --workers 2 --skip-image-video 1 and inspect one worker's log
-under --logs-dir to confirm QoS was accepted and a numeric stats.json is written.
-"""
-
-import argparse
-from pathlib import Path
-
-from datatrove.executor import LocalPipelineExecutor
-from datatrove.executor.slurm import SlurmPipelineExecutor
-from datatrove.pipeline.base import PipelineStep
-
-class ComputeEpisodeStatsShards(PipelineStep):
-    """Each worker computes per-episode stats for its ``episodes[rank::world_size]`` shard."""
-
-    def __init__(self, repo_id, root, new_root, skip_image_video, shard_dir, video_backend=None):
-        super().__init__()
-        self.repo_id = repo_id
-        self.root = root
-        self.new_root = new_root
-        self.skip_image_video = skip_image_video
-        self.shard_dir = shard_dir
-        self.video_backend = video_backend
-
-    def run(self, data=None, rank: int = 0, world_size: int = 1):
-        # NOTE: this method is pickled and executed on a worker, where this script's module
-        # globals are NOT available. Keep it self-contained: import locally and don't reference
-        # module-level helpers/constants.
-        import logging
-        import pickle
-        from pathlib import Path
-
-        from lerobot.datasets import LeRobotDataset, compute_dataset_episode_stats
-        from lerobot.utils.utils import init_logging
-
-        init_logging()
-        load_kwargs = {"video_backend": self.video_backend} if self.video_backend else {}
-        root = self.new_root if self.new_root and Path(self.new_root).exists() else self.root
-        dataset = LeRobotDataset(self.repo_id, root=root, **load_kwargs)
-
-        my_episodes = list(range(dataset.meta.total_episodes))[rank::world_size]
-        if not my_episodes:
-            logging.info(f"Rank {rank}: no episodes assigned")
-            return
-        logging.info(f"Rank {rank}: {len(my_episodes)} / {dataset.meta.total_episodes} episodes")
-
-        episode_stats = compute_dataset_episode_stats(
-            dataset,
-            episode_indices=my_episodes,
-            skip_image_video=self.skip_image_video,
-        )
-
-        shard_dir = Path(self.shard_dir)
-        shard_dir.mkdir(parents=True, exist_ok=True)
-        out = shard_dir / f"episode_stats_{rank:05d}.pkl"
-        with open(out, "wb") as f:
-            pickle.dump(episode_stats, f)
-        logging.info(f"Rank {rank}: saved {len(episode_stats)} episode stats to {out}")
-
-
-class AggregateEpisodeStats(PipelineStep):
-    """Merge all per-episode stat shards into meta/stats.json."""
-
-    def __init__(
-        self,
-        repo_id,
-        root,
-        new_root,
-        shard_dir,
-        push_to_hub=False,
-        video_backend=None,
-        update_episode_stats=False,
-    ):
-        super().__init__()
-        self.repo_id = repo_id
-        self.root = root
-        self.new_root = new_root
-        self.shard_dir = shard_dir
-        self.push_to_hub = push_to_hub
-        self.video_backend = video_backend
-        self.update_episode_stats = update_episode_stats
-
-    def run(self, data=None, rank: int = 0, world_size: int = 1):
-        # NOTE: pickled and executed on a worker; keep self-contained (see ComputeEpisodeStatsShards.run).
-        import logging
-        import pickle
-        from pathlib import Path
-
-        from lerobot.datasets import LeRobotDataset, aggregate_episode_stats
-        from lerobot.utils.utils import init_logging
-
-        init_logging()
-        if rank != 0:
-            return
-
-        shard_dir = Path(self.shard_dir)
-        shards = sorted(shard_dir.glob("episode_stats_*.pkl"))
-        if not shards:
-            raise FileNotFoundError(f"No episode stat shards found in {shard_dir}")
-
-        # Shards map episode_index -> stats; merging by key makes a dropped shard show up as a
-        # missing episode and a re-run shard overwrite rather than double-count.
-        all_episode_stats = {}
-        for shard in shards:
-            with open(shard, "rb") as f:
-                all_episode_stats.update(pickle.load(f))
-        logging.info(f"Aggregating {len(all_episode_stats)} episode stats from {len(shards)} shards")
-
-        load_kwargs = {"video_backend": self.video_backend} if self.video_backend else {}
-        root = self.new_root if self.new_root and Path(self.new_root).exists() else self.root
-        dataset = LeRobotDataset(self.repo_id, root=root, **load_kwargs)
-
-        # Aggregation is order-independent, so the only way sharding changes the result is a
-        # gap (dropped shard) or an overlap (episode counted twice). Verify the shards cover
-        # every episode exactly once before writing stats.json.
-        expected_episodes = dataset.meta.total_episodes
-        if len(all_episode_stats) != expected_episodes:
-            raise ValueError(
-                f"Expected {expected_episodes} per-episode stats (one per episode) but got "
-                f"{len(all_episode_stats)} across {len(shards)} shards. A compute shard is likely "
-                "missing or was written more than once; re-run the failed shards before aggregating."
-            )
-
-        # Frame-count check catches the case where a duplicate and a gap cancel out in the
-        # episode count: summed per-episode frame counts must equal the dataset's total frames.
-        stats_values = list(all_episode_stats.values())
-        numeric_key = next(
-            (
-                k
-                for k, v in dataset.meta.features.items()
-                if v["dtype"] not in ("image", "video", "string") and stats_values and k in stats_values[0]
-            ),
-            None,
-        )
-        if numeric_key is not None:
-            total_frames = sum(int(s[numeric_key]["count"][0]) for s in stats_values)
-            if total_frames != dataset.meta.total_frames:
-                raise ValueError(
-                    f"Summed frame count from shards ({total_frames}) != dataset total_frames "
-                    f"({dataset.meta.total_frames}); episodes are double-counted or missing."
-                )
-
-        new_stats = aggregate_episode_stats(
-            dataset, all_episode_stats, update_episode_stats=self.update_episode_stats
-        )
-        if new_stats is None:
-            raise RuntimeError("Aggregation produced no stats")
-        logging.info(f"Wrote stats for features: {list(new_stats.keys())} to {dataset.root}")
-
-        if self.push_to_hub:
-            logging.info(f"Pushing {self.repo_id} to hub")
-            dataset.push_to_hub()
-
-
-def _mem_gb(mem: str) -> int:
-    """Parse '4G' / '4GB' / '4' into an int number of GB for datatrove's mem_per_cpu_gb."""
-    s = str(mem).strip().lower().rstrip("b").rstrip("g")
-    return int(float(s))
-
-
-def _make_executor(
-    pipeline,
-    logs_dir,
-    job_name,
-    slurm,
-    workers,
-    tasks,
-    time,
-    partition,
-    cpus,
-    mem,
-    qos=None,
-    env_command=None,
-    venv_path=None,
-    depends=None,
-):
-    kwargs = {"pipeline": pipeline, "logging_dir": str(Path(logs_dir) / job_name)}
-    if slurm:
-        kwargs.update(
-            {
-                "job_name": job_name,
-                "tasks": tasks,
-                "workers": workers,
-                "time": time,
-                "partition": partition,
-                "cpus_per_task": cpus,
-                "mem_per_cpu_gb": _mem_gb(mem),  # datatrove's native field (int GB)
-                "sbatch_args": {},
-            }
-        )
-        if qos:
-            kwargs["qos"] = qos  # -> "#SBATCH --qos=<qos>" on every worker
-        if venv_path:
-            kwargs["venv_path"] = venv_path  # datatrove sources this before the python step
-        if env_command:
-            kwargs["env_command"] = env_command  # extra raw snippet before python (composes with venv_path)
-        if depends is not None:
-            kwargs["depends"] = depends  # chains --dependency=afterok:<compute jobid>
-        return SlurmPipelineExecutor(**kwargs)
-    kwargs.update({"tasks": tasks, "workers": 1})
-    return LocalPipelineExecutor(**kwargs)
-
-
-def _maybe_reference_copy(repo_id, root, new_root, download_videos):
-    """Create the read-only-safe reference copy once, before submitting workers.
-
-    Loads metadata only (to resolve the source root and revision) instead of a full
-    ``LeRobotDataset``, which would also memory-map the entire frame index just to read a
-    path. Fetches the source into the shared cache so the copy's symlinks point at real
-    files and workers don't each re-download, pulling videos only when the run needs them
-    (i.e. when image/video stats are being recomputed).
-    """
-    if not new_root:
-        return
-    from huggingface_hub import snapshot_download
-
-    from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
-    from lerobot.scripts.lerobot_edit_dataset import _reference_copy_dataset
-    from lerobot.utils.constants import HF_LEROBOT_HUB_CACHE
-
-    new_root_path = Path(new_root)
-    if new_root_path.exists():
-        return
-
-    meta = LeRobotDatasetMetadata(repo_id, root=Path(root) if root else None)
-    ignore_patterns = None if download_videos else "videos/"
-    if root:
-        snapshot_download(
-            repo_id,
-            repo_type="dataset",
-            revision=meta.revision,
-            local_dir=meta.root,
-            ignore_patterns=ignore_patterns,
-        )
-        src_root = Path(meta.root)
-    else:
-        src_root = Path(
-            snapshot_download(
-                repo_id,
-                repo_type="dataset",
-                revision=meta.revision,
-                cache_dir=HF_LEROBOT_HUB_CACHE,
-                ignore_patterns=ignore_patterns,
-            )
-        )
-    _reference_copy_dataset(src_root, new_root_path)
-
-
-def _add_shared_args(p):
-    p.add_argument("--repo-id", type=str, required=True, help="Dataset identifier, e.g. 'user/dataset'.")
-    p.add_argument("--root", type=str, default=None, help="Source dataset root (defaults to the Hub cache).")
-    p.add_argument(
-        "--new-root",
-        type=str,
-        default=None,
-        help="Writable output root; a read-only-safe reference copy of --root. If omitted, stats "
-        "are written in place at --root.",
-    )
-    p.add_argument("--shard-dir", type=Path, default=Path("stats_shards"), help="Per-rank shard dir.")
-    p.add_argument("--logs-dir", type=Path, default=Path("logs"), help="datatrove logs dir.")
-    p.add_argument("--job-name", type=str, default=None, help="SLURM job name.")
-    p.add_argument("--slurm", type=int, default=1, help="1 = submit via SLURM; 0 = run locally.")
-    p.add_argument("--partition", type=str, default=None, help="SLURM partition, e.g. 'hopper-cpu'.")
-    p.add_argument("--qos", type=str, default=None, help="SLURM QoS, e.g. 'normal'. Passed to every worker.")
-    p.add_argument("--cpus-per-task", type=int, default=4, help="CPUs per SLURM task.")
-    p.add_argument("--mem-per-cpu", type=str, default="4G", help="Memory per CPU, e.g. '4G'.")
-    p.add_argument(
-        "--video-backend",
-        type=str,
-        default=None,
-        help="Video decoding backend (e.g. 'pyav', 'torchcodec'). Defaults to the dataset's default; "
-        "use 'pyav' if torchcodec fails to load locally.",
-    )
-    p.add_argument("--venv-path", type=str, default=None, help="venv activate script sourced on each worker.")
-    p.add_argument(
-        "--env-command",
-        type=str,
-        default=None,
-        help="Raw shell snippet injected into each worker's sbatch before the python step "
-        "(e.g. to export HF_LEROBOT_HOME). Runs in addition to --venv-path.",
-    )
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="PATCHED SLURM-distributed LeRobotDataset stats recomputation",
-        formatter_class=argparse.RawDescriptionHelpFormatter,
-    )
-    sub = parser.add_subparsers(dest="command", required=True)
-
-    cp = sub.add_parser("compute", help="Distribute per-episode stats across SLURM workers.")
-    _add_shared_args(cp)
-    cp.add_argument("--workers", type=int, default=50, help="Number of parallel SLURM tasks.")
-    cp.add_argument(
-        "--skip-image-video",
-        type=int,
-        default=1,
-        help="1 = numeric features only (fast); 0 = also recompute image/video stats (decodes frames).",
-    )
-    cp.add_argument(
-        "--chain-aggregate",
-        action="store_true",
-        help="After building compute, submit aggregate with an afterok dependency (single command).",
-    )
-    cp.add_argument("--push-to-hub", action="store_true", help="For the chained aggregate: push after done.")
-    cp.add_argument(
-        "--update-episode-stats",
-        action="store_true",
-        help="For the chained aggregate: also rewrite per-episode stats in the episodes parquet.",
-    )
-
-    ap = sub.add_parser("aggregate", help="Merge shards into meta/stats.json.")
-    _add_shared_args(ap)
-    ap.add_argument("--push-to-hub", action="store_true", help="Push the dataset after aggregation.")
-    ap.add_argument(
-        "--update-episode-stats",
-        action="store_true",
-        help="Also rewrite per-episode stats in the episodes parquet to match stats.json.",
-    )
-    ap.add_argument(
-        "--depends-job-id",
-        type=str,
-        default=None,
-        help="Optional SLURM job id; aggregate waits for it (afterok) before running.",
-    )
-
-    args = parser.parse_args()
-    slurm = args.slurm == 1
-
-    if args.command == "compute":
-        # The reference copy (if any) is created once on the submitting node so workers
-        # can all load --new-root without racing to build it. Videos are only fetched when
-        # image/video stats are being recomputed.
-        _maybe_reference_copy(
-            args.repo_id, args.root, args.new_root, download_videos=not bool(args.skip_image_video)
-        )
-
-        compute_exec = _make_executor(
-            pipeline=[
-                ComputeEpisodeStatsShards(
-                    args.repo_id,
-                    args.root,
-                    args.new_root,
-                    bool(args.skip_image_video),
-                    str(args.shard_dir),
-                    args.video_backend,
-                )
-            ],
-            logs_dir=args.logs_dir,
-            job_name=args.job_name or "recompute_stats_compute",
-            slurm=slurm,
-            workers=args.workers,
-            tasks=args.workers,
-            time="24:00:00",
-            partition=args.partition,
-            cpus=args.cpus_per_task,
-            mem=args.mem_per_cpu,
-            qos=args.qos,
-            env_command=args.env_command,
-            venv_path=args.venv_path,
-        )
-
-        if args.chain_aggregate and slurm:
-            # Build aggregate depending on compute. datatrove launches the dependency
-            # (compute) first, then submits aggregate with --dependency=afterok:<jobid>.
-            aggregate_exec = _make_executor(
-                pipeline=[
-                    AggregateEpisodeStats(
-                        args.repo_id,
-                        args.root,
-                        args.new_root,
-                        str(args.shard_dir),
-                        args.push_to_hub,
-                        args.video_backend,
-                        args.update_episode_stats,
-                    )
-                ],
-                logs_dir=args.logs_dir,
-                job_name="recompute_stats_aggregate",
-                slurm=slurm,
-                workers=1,
-                tasks=1,
-                time="02:00:00",
-                partition=args.partition,
-                cpus=args.cpus_per_task,
-                mem=args.mem_per_cpu,
-                qos=args.qos,
-                env_command=args.env_command,
-                venv_path=args.venv_path,
-                depends=compute_exec,
-            )
-            aggregate_exec.run()
-        else:
-            compute_exec.run()
-    else:
-        aggregate_exec = _make_executor(
-            pipeline=[
-                AggregateEpisodeStats(
-                    args.repo_id,
-                    args.root,
-                    args.new_root,
-                    str(args.shard_dir),
-                    args.push_to_hub,
-                    args.video_backend,
-                    args.update_episode_stats,
-                )
-            ],
-            logs_dir=args.logs_dir,
-            job_name=args.job_name or "recompute_stats_aggregate",
-            slurm=slurm,
-            workers=1,
-            tasks=1,
-            time="02:00:00",
-            partition=args.partition,
-            cpus=args.cpus_per_task,
-            mem=args.mem_per_cpu,
-            qos=args.qos,
-            env_command=args.env_command,
-            venv_path=args.venv_path,
-        )
-        if args.depends_job_id is not None:
-            aggregate_exec.depends_job_id = args.depends_job_id
-        aggregate_exec.run()
-
-
-if __name__ == "__main__":
-    main()

src/lerobot/datasets/__init__.py

@@ -25,8 +25,6 @@ from .compute_stats import DEFAULT_QUANTILES, aggregate_stats, get_feature_stats
 from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
 from .dataset_tools import (
     add_features,
-    aggregate_episode_stats,
-    compute_dataset_episode_stats,
     convert_image_to_video_dataset,
     delete_episodes,
     merge_datasets,
@@ -36,7 +34,6 @@ from .dataset_tools import (
     reencode_dataset,
     remove_feature,
     split_dataset,
-    write_episode_stats,
 )
 from .factory import make_dataset, make_train_eval_datasets, resolve_delta_timestamps
 from .image_writer import safe_stop_image_writer
@@ -81,10 +78,8 @@ __all__ = [
     "detect_available_encoders_pyav",
     "add_features",
     "aggregate_datasets",
-    "aggregate_episode_stats",
     "aggregate_pipeline_dataset_features",
     "aggregate_stats",
-    "compute_dataset_episode_stats",
     "convert_image_to_video_dataset",
     "create_initial_features",
     "compute_sampler_state",
@@ -104,6 +99,5 @@ __all__ = [
     "resolve_delta_timestamps",
     "safe_stop_image_writer",
     "split_dataset",
-    "write_episode_stats",
     "write_stats",
 ]

src/lerobot/datasets/compute_stats.py

@@ -18,8 +18,9 @@ from __future__ import annotations
 import logging
 
 import numpy as np
+import torch
 
-from lerobot.processor import RelativeActionsProcessorStep
+from lerobot.processor import RelativeActionsProcessorStep, to_relative_actions
 from lerobot.utils.constants import ACTION, OBS_STATE
 
 from .io_utils import load_image_as_numpy
@@ -660,6 +661,8 @@ def _compute_relative_chunk_batch(
     all_states: np.ndarray,
     chunk_size: int,
     relative_mask: np.ndarray,
+    pose_representation: str = "componentwise",
+    se3_pose_groups: list[list[int]] | None = None,
 ) -> np.ndarray:
     """Vectorised relative-action computation for a batch of start indices.
 
@@ -671,6 +674,18 @@ def _compute_relative_chunk_batch(
     frame_idx = start_indices[:, None] + offsets[None, :]
     chunks = all_actions[frame_idx].copy()
     states = all_states[start_indices]
+    if pose_representation == "se3":
+        return (
+            to_relative_actions(
+                torch.from_numpy(chunks),
+                torch.from_numpy(states),
+                relative_mask.astype(bool).tolist(),
+                pose_representation=pose_representation,
+                se3_pose_groups=se3_pose_groups,
+            )
+            .numpy()
+            .reshape(-1, all_actions.shape[1])
+        )
     mask_dim = len(relative_mask)
     chunks[:, :, :mask_dim] -= states[:, None, :mask_dim] * relative_mask[None, None, :]
     return chunks.reshape(-1, all_actions.shape[1])
@@ -682,6 +697,9 @@ def compute_relative_action_stats(
     chunk_size: int,
     exclude_joints: list[str] | None = None,
     num_workers: int = 0,
+    state_from_action: bool = False,
+    pose_representation: str = "componentwise",
+    se3_pose_groups: list[list[int]] | None = None,
 ) -> dict[str, np.ndarray]:
     """Compute normalization statistics for relative actions over the full dataset.
 
@@ -700,6 +718,9 @@ def compute_relative_action_stats(
         num_workers: Number of parallel threads for computation. Values ≤1
             mean single-threaded. Numpy releases the GIL so threads give
             real parallelism here.
+        state_from_action: Use the current absolute action as state. This is
+            intended for state-less pose datasets where each action row is the
+            synchronized measured robot pose.
 
     Returns:
         Statistics dict with keys "mean", "std", "min", "max", "q01", …, "q99".
@@ -722,7 +743,7 @@ def compute_relative_action_stats(
 
     logging.info("Loading action/state data for relative action stats...")
     all_actions = np.array(hf_dataset[ACTION], dtype=np.float32)
-    all_states = np.array(hf_dataset[OBS_STATE], dtype=np.float32)
+    all_states = all_actions if state_from_action else np.array(hf_dataset[OBS_STATE], dtype=np.float32)
     episode_indices = np.array(hf_dataset["episode_index"])
 
     valid_starts = _get_valid_chunk_starts(episode_indices, chunk_size)
@@ -754,6 +775,8 @@ def compute_relative_action_stats(
                     all_states,
                     chunk_size,
                     relative_mask,
+                    pose_representation,
+                    se3_pose_groups,
                 )
                 for batch in batches
             ]
@@ -762,7 +785,15 @@ def compute_relative_action_stats(
     else:
         for batch in batches:
             running_stats.update(
-                _compute_relative_chunk_batch(batch, all_actions, all_states, chunk_size, relative_mask)
+                _compute_relative_chunk_batch(
+                    batch,
+                    all_actions,
+                    all_states,
+                    chunk_size,
+                    relative_mask,
+                    pose_representation,
+                    se3_pose_groups,
+                )
             )
 
     stats = running_stats.get_statistics()
@@ -777,3 +808,58 @@ def compute_relative_action_stats(
     )
 
     return stats
+
+
+def compute_state_history_stats(
+    hf_dataset,
+    features: dict,
+    history_steps: int,
+    exclude_joints: list[str] | None = None,
+    relative: bool = False,
+    pose_representation: str = "componentwise",
+    se3_pose_groups: list[list[int]] | None = None,
+) -> dict[str, np.ndarray]:
+    """Compute stats for flattened state history synthesized from absolute actions.
+
+    History is left-padded with the first action of each episode, matching dataset
+    boundary padding. When ``relative`` is enabled, every history pose is expressed
+    relative to its newest pose while excluded dimensions remain absolute.
+    """
+    if history_steps < 1:
+        raise ValueError("history_steps must be at least 1")
+    if exclude_joints is None:
+        exclude_joints = []
+
+    actions = np.asarray(hf_dataset[ACTION], dtype=np.float32)
+    episode_indices = np.asarray(hf_dataset["episode_index"])
+    sample_indices = np.arange(len(actions))
+    episode_starts = np.maximum.accumulate(
+        np.where(
+            np.concatenate(([True], episode_indices[1:] != episode_indices[:-1])),
+            sample_indices,
+            0,
+        )
+    )
+    offsets = np.arange(-(history_steps - 1), 1)
+    history_indices = np.maximum(sample_indices[:, None] + offsets[None, :], episode_starts[:, None])
+    history = actions[history_indices].copy()
+
+    if relative:
+        state_dim = actions.shape[-1]
+        names = features.get(ACTION, {}).get("names")
+        mask_step = RelativeActionsProcessorStep(
+            enabled=True,
+            exclude_joints=exclude_joints,
+            action_names=names,
+        )
+        mask = mask_step._build_mask(state_dim)
+        history = to_relative_actions(
+            torch.from_numpy(history),
+            torch.from_numpy(history[:, -1].copy()),
+            mask,
+            pose_representation=pose_representation,
+            se3_pose_groups=se3_pose_groups,
+        ).numpy()
+
+    flattened = history.reshape(len(history), -1)
+    return get_feature_stats(flattened, axis=0, keepdims=False)

src/lerobot/datasets/dataset_tools.py

@@ -33,13 +33,11 @@ from pathlib import Path
 import datasets
 import numpy as np
 import pandas as pd
-import pyarrow as pa
 import pyarrow.parquet as pq
 import torch
 from tqdm import tqdm
 
 from lerobot.configs import (
-    DEFAULT_DEPTH_UNIT,
     DepthEncoderConfig,
     RGBEncoderConfig,
     VideoEncoderConfig,
@@ -53,15 +51,12 @@ from lerobot.utils.utils import flatten_dict
 
 from .aggregate import aggregate_datasets
 from .compute_stats import (
-    RunningQuantileStats,
     aggregate_stats,
-    auto_downsample_height_width,
     compute_episode_stats,
     compute_relative_action_stats,
-    sample_indices,
+    compute_state_history_stats,
 )
 from .dataset_metadata import LeRobotDatasetMetadata
-from .depth_utils import dequantize_depth
 from .image_writer import write_image
 from .io_utils import (
     get_parquet_file_size_in_mb,
@@ -83,7 +78,6 @@ from .utils import (
     update_chunk_file_indices,
 )
 from .video_utils import (
-    decode_video_frames,
     encode_video_frames,
     reencode_video,
 )
@@ -1566,191 +1560,6 @@ def modify_tasks(
     return dataset
 
 
-def _load_episode_image_frames(
-    dataset: LeRobotDataset,
-    key: str,
-    ep_idx: int,
-    frame_offsets: list[int],
-    is_depth: bool,
-) -> np.ndarray:
-    """Load sampled frames of an image feature for one episode as a (N, C, H, W) array."""
-    ep = dataset.meta.episodes[ep_idx]
-    from_idx = ep["dataset_from_index"]
-    column = dataset.hf_dataset.with_format(None).select_columns(key)
-
-    frames = []
-    for offset in frame_offsets:
-        img = column[from_idx + offset][key]
-        if is_depth:
-            arr = np.array(img)
-            if arr.ndim == 2:
-                arr = arr[np.newaxis, ...]
-        else:
-            arr = np.transpose(np.array(img.convert("RGB"), dtype=np.uint8), (2, 0, 1))
-        frames.append(auto_downsample_height_width(arr))
-    return np.stack(frames)
-
-
-def _load_episode_video_frames(
-    dataset: LeRobotDataset,
-    key: str,
-    ep_idx: int,
-    frame_offsets: list[int],
-    is_depth: bool,
-) -> np.ndarray:
-    """Load sampled frames of a video feature for one episode as a (N, C, H, W) array."""
-    ep = dataset.meta.episodes[ep_idx]
-    video_path = dataset.root / dataset.meta.get_video_file_path(ep_idx, key)
-    from_timestamp = ep[f"videos/{key}/from_timestamp"]
-    timestamps = [from_timestamp + offset / dataset.meta.fps for offset in frame_offsets]
-
-    frames = decode_video_frames(
-        video_path,
-        timestamps,
-        dataset.tolerance_s,
-        backend=dataset._video_backend,
-        return_uint8=not is_depth,
-        is_depth=is_depth,
-    )
-    if is_depth:
-        # ``decode_video_frames`` returns raw 12-bit codec values; dequantize back to
-        # the recorded depth unit so stats match record-time stats (which are stored in
-        # ``info.depth_unit`` and only rescaled to the output unit on read).
-        info = dataset.meta.features[key].get("info") or {}
-        depth_encoder = DepthEncoderConfig.from_video_info(info)
-        frames = dequantize_depth(
-            frames,
-            depth_min=depth_encoder.depth_min,
-            depth_max=depth_encoder.depth_max,
-            shift=depth_encoder.shift,
-            use_log=depth_encoder.use_log,
-            output_unit=info.get("depth_unit") or DEFAULT_DEPTH_UNIT,
-        )
-    return np.stack([auto_downsample_height_width(frame) for frame in frames.numpy()])
-
-
-def _compute_visual_episode_stats(
-    dataset: LeRobotDataset,
-    ep_idx: int,
-    visual_keys: list[str],
-    frame_batch_size: int = 32,
-) -> dict:
-    """Compute per-episode statistics for image/video features by sampling frames.
-
-    Mirrors the image/video branch of :func:`compute_episode_stats`: per-channel stats
-    are computed on downsampled sampled frames, then RGB stats are rescaled to [0, 1]
-    (depth maps keep their native units).
-
-    Frames are decoded and accumulated into a :class:`RunningQuantileStats` in batches of
-    ``frame_batch_size`` rather than materialising every sampled frame at once. Peak memory
-    is bounded by one batch (``frame_batch_size x C x H x W``) regardless of episode length,
-    which keeps long, high-resolution episodes from exhausting memory.
-    """
-    ep_length = dataset.meta.episodes[ep_idx]["length"]
-    frame_offsets = sample_indices(ep_length)
-
-    ep_stats = {}
-    for key in visual_keys:
-        is_depth = key in dataset.meta.depth_keys
-        is_video = dataset.meta.features[key]["dtype"] == "video"
-
-        running = RunningQuantileStats()
-        for start in range(0, len(frame_offsets), frame_batch_size):
-            batch_offsets = frame_offsets[start : start + frame_batch_size]
-            if is_video:
-                frames = _load_episode_video_frames(dataset, key, ep_idx, batch_offsets, is_depth)
-            else:
-                frames = _load_episode_image_frames(dataset, key, ep_idx, batch_offsets, is_depth)
-            # (N, C, H, W) -> (N * H * W, C) so stats are accumulated per channel.
-            running.update(np.moveaxis(frames, 1, -1).reshape(-1, frames.shape[1]))
-
-        stats = running.get_statistics()
-        normalization_factor = 1.0 if is_depth else 255.0
-        num_channels = stats["mean"].shape[0]
-        # ``count`` follows the per-frame convention of ``get_feature_stats`` (number of
-        # sampled frames), not the per-pixel count tracked internally by RunningQuantileStats.
-        ep_stats[key] = {
-            k: np.array([len(frame_offsets)])
-            if k == "count"
-            else v.reshape(num_channels, 1, 1) / normalization_factor
-            for k, v in stats.items()
-        }
-
-    return ep_stats
-
-
-def compute_dataset_episode_stats(
-    dataset: LeRobotDataset,
-    episode_indices: list[int] | None = None,
-    skip_image_video: bool = True,
-    drop_keys: list[str] | None = None,
-) -> dict[int, dict]:
-    """Compute per-episode statistics for a subset of episodes.
-
-    This is the shardable unit of work behind :func:`recompute_stats`: distribute
-    ``episode_indices`` across workers (e.g. ``list(range(n))[rank::world_size]``),
-    then combine the results with :func:`aggregate_episode_stats`.
-
-    Args:
-        dataset: The LeRobotDataset to compute stats for.
-        episode_indices: Episodes to process. When ``None``, all episodes are processed.
-        skip_image_video: If True (default), only numeric features are computed. If False,
-            image/video stats are also computed by sampling and decoding frames.
-        drop_keys: Feature keys to exclude (e.g. ``action`` when it is computed separately
-            in relative-action space).
-
-    Returns:
-        A mapping of episode index to its per-episode stat dict. Keeping the episode index
-        (rather than a bare list) lets callers write the stats back to the correct episode
-        row, and survives sharding since shards can be merged by key.
-    """
-    features = dataset.meta.features
-    meta_keys = {"index", "episode_index", "task_index", "frame_index", "timestamp"}
-    drop = set(drop_keys or [])
-    features_to_compute = {
-        k: v
-        for k, v in features.items()
-        if v["dtype"] != "string"
-        and k not in meta_keys
-        and k not in drop
-        and (not skip_image_video or v["dtype"] not in ["image", "video"])
-    }
-    numeric_keys = [k for k, v in features_to_compute.items() if v["dtype"] not in ["image", "video"]]
-    visual_keys = [k for k, v in features_to_compute.items() if v["dtype"] in ["image", "video"]]
-
-    if dataset.meta.episodes is None:
-        dataset.meta.episodes = load_episodes(dataset.meta.root)
-
-    if episode_indices is None:
-        episode_indices = list(range(dataset.meta.total_episodes))
-
-    # Group requested episodes by their data parquet file so each file is read once.
-    file_to_episodes: dict[Path, list[int]] = {}
-    for ep_idx in episode_indices:
-        file_to_episodes.setdefault(dataset.meta.get_data_file_path(ep_idx), []).append(ep_idx)
-
-    all_episode_stats = {}
-    for src_path, eps in tqdm(sorted(file_to_episodes.items()), desc="Computing stats from data files"):
-        df = pd.read_parquet(dataset.root / src_path) if numeric_keys else None
-        for ep_idx in sorted(eps):
-            episode_data = {}
-            if numeric_keys:
-                ep_df = df[df["episode_index"] == ep_idx]
-                for key in numeric_keys:
-                    if key in ep_df.columns:
-                        values = ep_df[key].values
-                        episode_data[key] = (
-                            np.stack(values) if hasattr(values[0], "__len__") else np.array(values)
-                        )
-
-            ep_stats = compute_episode_stats(episode_data, features_to_compute)
-            if visual_keys:
-                ep_stats.update(_compute_visual_episode_stats(dataset, int(ep_idx), visual_keys))
-            all_episode_stats[int(ep_idx)] = ep_stats
-
-    return all_episode_stats
-
-
 def recompute_stats(
     dataset: LeRobotDataset,
     skip_image_video: bool = True,
@@ -1758,21 +1567,19 @@ def recompute_stats(
     relative_exclude_joints: list[str] | None = None,
     chunk_size: int = 50,
     num_workers: int = 0,
-    update_episode_stats: bool = False,
+    state_from_action: bool = False,
+    state_history_steps: int = 1,
+    relative_state_history: bool = False,
+    relative_state_exclude_joints: list[str] | None = None,
+    relative_pose_representation: str = "componentwise",
+    relative_se3_pose_groups: list[list[int]] | None = None,
 ) -> LeRobotDataset:
     """Recompute stats.json from scratch by iterating all episodes.
 
     Args:
         dataset: The LeRobotDataset to recompute stats for.
         skip_image_video: If True (default), only recompute stats for numeric features
-            (action, state, etc.) and keep existing image/video stats unchanged. If False,
-            image/video stats are also recomputed by sampling and decoding frames from each
-            episode (this reads the image/video files, unlike the numeric-only path).
-        update_episode_stats: If True, also rewrite the per-episode ``stats/*`` columns in the
-            episodes parquet files so they stay consistent with the aggregated ``stats.json``.
-            Defaults to False (only ``stats.json`` is rewritten). Requires a writable
-            ``dataset.root``. Note that relative-action stats are aggregate-only and are not
-            written per-episode.
+            (action, state, etc.) and keep existing image/video stats unchanged.
         relative_action: If True, compute action stats in relative space by
             iterating all valid action chunks and subtracting the current state.
             This matches the normalization distribution the model sees during
@@ -1783,18 +1590,53 @@ def recompute_stats(
             ``policy.chunk_size``. Only used when ``relative_action=True``.
         num_workers: Number of parallel threads for relative action stats computation.
             Values ≤1 mean single-threaded. Only used when ``relative_action=True``.
+        state_from_action: Use absolute action rows as synthetic state while
+            computing relative-action stats, and write their absolute statistics
+            under ``observation.state``.
+        state_history_steps: Number of consecutive synthesized state samples.
+        relative_state_history: Express state history relative to its newest pose.
+        relative_state_exclude_joints: State dimensions to retain as absolute.
+        relative_pose_representation: ``componentwise`` for legacy subtraction or
+            ``se3`` for ``inv(T_current) @ T_target`` pose composition.
+        relative_se3_pose_groups: Six-index xyz+rotation-vector pose groups.
 
     Returns:
         The same dataset with updated stats.
     """
     features = dataset.meta.features
+    meta_keys = {"index", "episode_index", "task_index", "frame_index", "timestamp"}
+    numeric_features = {
+        k: v
+        for k, v in features.items()
+        if v["dtype"] not in ["image", "video", "string"] and k not in meta_keys
+    }
+
+    if skip_image_video:
+        features_to_compute = numeric_features
+    else:
+        features_to_compute = {
+            k: v for k, v in features.items() if v["dtype"] != "string" and k not in meta_keys
+        }
 
     # When relative_action is enabled, compute action stats via chunk-based sampling
     # (matching what the model sees during training) and skip action in the
     # per-episode pass below.
     relative_action_stats = None
-    drop_keys = None
-    if relative_action and ACTION in features and OBS_STATE in features:
+    synthetic_state_stats = None
+    if state_from_action:
+        if ACTION not in features:
+            raise ValueError("state_from_action requires an action feature")
+        synthetic_state_stats = compute_state_history_stats(
+            dataset.hf_dataset,
+            features,
+            history_steps=state_history_steps,
+            exclude_joints=relative_state_exclude_joints,
+            relative=relative_state_history,
+            pose_representation=relative_pose_representation,
+            se3_pose_groups=relative_se3_pose_groups,
+        )
+
+    if relative_action and ACTION in features and (OBS_STATE in features or state_from_action):
         if relative_exclude_joints is None:
             relative_exclude_joints = ["gripper"]
         relative_action_stats = compute_relative_action_stats(
@@ -1803,106 +1645,62 @@ def recompute_stats(
             chunk_size=chunk_size,
             exclude_joints=relative_exclude_joints,
             num_workers=num_workers,
+            state_from_action=state_from_action,
+            pose_representation=relative_pose_representation,
+            se3_pose_groups=relative_se3_pose_groups,
         )
-        drop_keys = [ACTION]
+        features_to_compute.pop(ACTION, None)
 
-    all_episode_stats = compute_dataset_episode_stats(
-        dataset, skip_image_video=skip_image_video, drop_keys=drop_keys
-    )
+    logging.info(f"Recomputing stats for features: {list(features_to_compute.keys())}")
 
-    new_stats = aggregate_episode_stats(
-        dataset,
-        all_episode_stats,
-        extra_stats={ACTION: relative_action_stats} if relative_action_stats else None,
-        update_episode_stats=update_episode_stats,
-    )
-    if new_stats is None:
+    data_dir = dataset.root / DATA_DIR
+    parquet_files = sorted(data_dir.glob("*/*.parquet"))
+    if not parquet_files:
+        raise ValueError(f"No parquet files found in {data_dir}")
+
+    all_episode_stats = []
+    # TODO: enable image and video stats re-computation
+    numeric_keys = [k for k, v in features_to_compute.items() if v["dtype"] not in ["image", "video"]]
+
+    for parquet_path in tqdm(parquet_files, desc="Computing stats from data files"):
+        df = pd.read_parquet(parquet_path)
+
+        for ep_idx in sorted(df["episode_index"].unique()):
+            ep_df = df[df["episode_index"] == ep_idx]
+            episode_data = {}
+            for key in numeric_keys:
+                if key in ep_df.columns:
+                    values = ep_df[key].values
+                    if hasattr(values[0], "__len__"):
+                        episode_data[key] = np.stack(values)
+                    else:
+                        episode_data[key] = np.array(values)
+
+            ep_stats = compute_episode_stats(episode_data, features_to_compute)
+            all_episode_stats.append(ep_stats)
+
+    if features_to_compute and not all_episode_stats:
         logging.warning("No episode stats computed")
-    else:
-        logging.info("Stats recomputed successfully")
-    return dataset
+        return dataset
 
+    new_stats = aggregate_stats(all_episode_stats) if all_episode_stats else {}
 
-def write_episode_stats(dataset: LeRobotDataset, episode_stats: dict[int, dict]) -> None:
-    """Overwrite the per-episode ``stats/*`` columns in the episodes parquet files in place.
+    if relative_action_stats is not None:
+        new_stats[ACTION] = relative_action_stats
+    if synthetic_state_stats is not None:
+        new_stats[OBS_STATE] = synthetic_state_stats
 
-    Only the features present in ``episode_stats[ep_idx]`` are rewritten; stats columns for
-    features that were not recomputed are left untouched. Every other episode column (tasks,
-    length, chunk/file indices, frame ranges, …) is preserved. ``dataset.root`` must be
-    writable (e.g. the reference copy created for read-only sources).
-    """
-    if not episode_stats:
-        return
-
-    meta = dataset.meta
-    if meta.episodes is None:
-        meta.episodes = load_episodes(meta.root)
-
-    # Group episodes by the parquet file that holds them so each file is rewritten once.
-    file_to_episodes: dict[tuple[int, int], list[int]] = {}
-    for ep_idx in episode_stats:
-        ep = meta.episodes[ep_idx]
-        key = (ep["meta/episodes/chunk_index"], ep["meta/episodes/file_index"])
-        file_to_episodes.setdefault(key, []).append(ep_idx)
-
-    for (chunk_idx, file_idx), eps in file_to_episodes.items():
-        path = meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
-        table = pq.read_table(path)
-        rows = table.to_pylist()
-        row_by_ep = {row["episode_index"]: row for row in rows}
-        for ep_idx in eps:
-            row = row_by_ep[ep_idx]
-            for feature, feature_stats in episode_stats[ep_idx].items():
-                for stat_name, value in feature_stats.items():
-                    col = f"stats/{feature}/{stat_name}"
-                    if col in row:
-                        row[col] = np.asarray(value).tolist()
-        # Reuse the source schema so the rewritten stats keep the exact on-disk types.
-        new_table = pa.Table.from_pylist(rows, schema=table.schema)
-        pq.write_table(new_table, path, compression="snappy", use_dictionary=True)
-
-
-def aggregate_episode_stats(
-    dataset: LeRobotDataset,
-    episode_stats: dict[int, dict],
-    extra_stats: dict | None = None,
-    update_episode_stats: bool = False,
-) -> dict | None:
-    """Aggregate per-episode stats, merge with existing stats, and write ``stats.json``.
-
-    Companion to :func:`compute_dataset_episode_stats` for the distributed workflow: pass the
-    merged ``{episode_index: stats}`` mapping of every worker's per-episode stats. ``extra_stats``
-    lets callers inject feature stats computed outside the per-episode pass (e.g. relative-action
-    stats).
-
-    Args:
-        dataset: The dataset whose ``meta/stats.json`` (and optionally episode stats) is updated.
-        episode_stats: Mapping of episode index to its per-episode stat dict.
-        extra_stats: Feature stats to inject into the aggregate (not written per-episode).
-        update_episode_stats: If True, also rewrite the per-episode ``stats/*`` columns in the
-            episodes parquet files via :func:`write_episode_stats`.
-
-    Returns the written stats dict, or ``None`` if there was nothing to aggregate.
-    """
-    if not episode_stats and not extra_stats:
-        return None
-
-    new_stats = aggregate_stats(list(episode_stats.values())) if episode_stats else {}
-    if extra_stats:
-        new_stats.update(extra_stats)
-
-    # Merge: keep existing stats for features we didn't recompute.
+    # Merge: keep existing stats for features we didn't recompute
     if dataset.meta.stats:
         for key, value in dataset.meta.stats.items():
-            new_stats.setdefault(key, value)
+            if key not in new_stats:
+                new_stats[key] = value
 
     write_stats(new_stats, dataset.root)
     dataset.meta.stats = new_stats
 
-    if update_episode_stats:
-        write_episode_stats(dataset, episode_stats)
-
-    return new_stats
+    logging.info("Stats recomputed successfully")
+    return dataset
 
 
 def convert_image_to_video_dataset(

src/lerobot/policies/pi05/configuration_pi05.py

@@ -55,6 +55,18 @@ class PI05Config(PreTrainedConfig):
     relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
     # Populated at runtime from dataset metadata by make_policy.
     action_feature_names: list[str] | None = None
+    # ``se3`` uses inv(T_current) @ T_target for each xyz+rotation-vector pose group.
+    # ``componentwise`` preserves the legacy action - state behavior.
+    relative_pose_representation: str = "componentwise"
+    relative_se3_pose_groups: list[list[int]] = field(default_factory=lambda: [list(range(6))])
+
+    # Build proprioception from absolute action samples when the dataset has no
+    # observation.state. With history_steps=2, training samples request t-1 as
+    # well as the normal t..t+chunk_size-1 action targets.
+    state_from_action: bool = False
+    proprioception_history_steps: int = 1
+    use_relative_state_history: bool = False
+    relative_state_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
 
     # Real-Time Chunking (RTC) configuration
     rtc_config: RTCConfig | None = None
@@ -121,6 +133,20 @@ class PI05Config(PreTrainedConfig):
         if self.dtype not in ["bfloat16", "float32"]:
             raise ValueError(f"Invalid dtype: {self.dtype}")
 
+        if self.proprioception_history_steps < 1:
+            raise ValueError("proprioception_history_steps must be at least 1")
+
+        if self.relative_pose_representation not in {"componentwise", "se3"}:
+            raise ValueError(
+                "relative_pose_representation must be either 'componentwise' or 'se3', got "
+                f"{self.relative_pose_representation!r}"
+            )
+        for group in self.relative_se3_pose_groups:
+            if len(group) != 6 or len(set(group)) != 6 or any(index < 0 for index in group):
+                raise ValueError(f"Invalid six-index SE(3) pose group: {group}")
+        if self.relative_pose_representation == "se3" and not self.relative_se3_pose_groups:
+            raise ValueError("relative_pose_representation='se3' requires relative_se3_pose_groups")
+
     def validate_features(self) -> None:
         """Validate and set up input/output features."""
         for i in range(self.empty_cameras):
@@ -131,13 +157,6 @@ class PI05Config(PreTrainedConfig):
             )
             self.input_features[key] = empty_camera
 
-        if OBS_STATE not in self.input_features:
-            state_feature = PolicyFeature(
-                type=FeatureType.STATE,
-                shape=(self.max_state_dim,),  # Padded to max_state_dim
-            )
-            self.input_features[OBS_STATE] = state_feature
-
         if ACTION not in self.output_features:
             action_feature = PolicyFeature(
                 type=FeatureType.ACTION,
@@ -145,6 +164,25 @@ class PI05Config(PreTrainedConfig):
             )
             self.output_features[ACTION] = action_feature
 
+        if OBS_STATE not in self.input_features:
+            state_shape = (self.max_state_dim,)
+            if self.state_from_action and ACTION in self.output_features:
+                state_shape = self.output_features[ACTION].shape
+            state_feature = PolicyFeature(
+                type=FeatureType.STATE,
+                shape=state_shape,
+            )
+            self.input_features[OBS_STATE] = state_feature
+
+        state_dim = self.input_features[OBS_STATE].shape[-1]
+        history_state_dim = state_dim * self.proprioception_history_steps
+        if history_state_dim > self.max_state_dim:
+            raise ValueError(
+                "Flattened proprioception history exceeds max_state_dim: "
+                f"{state_dim} * {self.proprioception_history_steps} = {history_state_dim} > "
+                f"{self.max_state_dim}"
+            )
+
     def get_optimizer_preset(self) -> AdamWConfig:
         return AdamWConfig(
             lr=self.optimizer_lr,
@@ -168,7 +206,8 @@ class PI05Config(PreTrainedConfig):
 
     @property
     def action_delta_indices(self) -> list:
-        return list(range(self.chunk_size))
+        history_prefix = self.proprioception_history_steps - 1 if self.state_from_action else 0
+        return list(range(-history_prefix, self.chunk_size))
 
     @property
     def reward_delta_indices(self) -> None:

src/lerobot/policies/pi05/processor_pi05.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 
 from copy import deepcopy
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 import numpy as np
@@ -36,6 +36,7 @@ from lerobot.processor import (
     TokenizerProcessorStep,
     UnnormalizerProcessorStep,
     policy_action_to_transition,
+    to_relative_actions,
     transition_to_policy_action,
 )
 from lerobot.types import EnvTransition, TransitionKey
@@ -48,6 +49,150 @@ from lerobot.utils.constants import (
 from .configuration_pi05 import PI05Config
 
 
+@ProcessorStepRegistry.register(name="pi05_state_from_action_processor_step")
+@dataclass
+class Pi05StateFromActionProcessorStep(ProcessorStep):
+    """Synthesize proprioception from absolute actions in state-less datasets.
+
+    The dataset loader supplies ``history_steps - 1`` actions before the normal
+    target chunk. Those leading samples and action(t) become state history; only
+    the leading samples are then removed from the action targets.
+    """
+
+    enabled: bool = False
+    history_steps: int = 1
+    _inference_history: torch.Tensor | None = field(default=None, init=False, repr=False)
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        if not self.enabled:
+            return transition
+
+        observation = transition.get(TransitionKey.OBSERVATION, {})
+        observed_state = observation.get(OBS_STATE)
+        if observed_state is not None:
+            # At inference the robot normally provides only the current state and
+            # there is no action target. Build a rolling history in the processor.
+            if transition.get(TransitionKey.ACTION) is None and observed_state.ndim == 2:
+                if self._inference_history is None:
+                    self._inference_history = observed_state.unsqueeze(1).repeat(1, self.history_steps, 1)
+                else:
+                    self._inference_history = torch.cat(
+                        [self._inference_history[:, 1:], observed_state.unsqueeze(1)], dim=1
+                    )
+                new_transition = transition.copy()
+                new_observation = dict(observation)
+                new_observation[OBS_STATE] = self._inference_history.clone()
+                new_transition[TransitionKey.OBSERVATION] = new_observation
+                return new_transition
+            return transition
+
+        action = transition.get(TransitionKey.ACTION)
+        if action is None:
+            raise ValueError("Cannot synthesize PI0.5 state without action")
+        if action.ndim != 3:
+            raise ValueError(f"Expected batched action chunks with shape (B, T, D), got {action.shape}")
+        if action.shape[1] < self.history_steps:
+            raise ValueError(
+                f"Action chunk has {action.shape[1]} steps, fewer than history_steps={self.history_steps}"
+            )
+
+        new_transition = transition.copy()
+        new_observation = dict(observation)
+        state = action[:, : self.history_steps].clone()
+        if self.history_steps == 1:
+            state = state[:, 0]
+        new_observation[OBS_STATE] = state
+        new_transition[TransitionKey.OBSERVATION] = new_observation
+        new_transition[TransitionKey.ACTION] = action[:, self.history_steps - 1 :]
+        return new_transition
+
+    def get_config(self) -> dict[str, Any]:
+        return {"enabled": self.enabled, "history_steps": self.history_steps}
+
+    def reset(self) -> None:
+        self._inference_history = None
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        return features
+
+
+@ProcessorStepRegistry.register(name="pi05_flatten_state_history_processor_step")
+@dataclass
+class Pi05FlattenStateHistoryProcessorStep(ProcessorStep):
+    """Optionally relativize raw state history, then flatten it for PI0.5."""
+
+    history_steps: int = 1
+    max_state_dim: int = 32
+    relative: bool = False
+    exclude_joints: list[str] = field(default_factory=list)
+    state_names: list[str] | None = None
+    pose_representation: str = "componentwise"
+    se3_pose_groups: list[list[int]] = field(default_factory=list)
+
+    def __call__(self, transition: EnvTransition) -> EnvTransition:
+        observation = transition.get(TransitionKey.OBSERVATION, {})
+        state = observation.get(OBS_STATE)
+        if state is None:
+            raise ValueError("State is required for PI05")
+        if self.history_steps == 1 and state.ndim == 2:
+            state = state.unsqueeze(1)
+        if state.ndim != 3 or state.shape[1] != self.history_steps:
+            raise ValueError(
+                f"Expected state history with shape (B, {self.history_steps}, D), got {state.shape}"
+            )
+
+        flattened_dim = state.shape[1] * state.shape[2]
+        if flattened_dim > self.max_state_dim:
+            raise ValueError(
+                f"Flattened state history has {flattened_dim} dimensions, above max_state_dim={self.max_state_dim}"
+            )
+
+        processed_state = state.clone()
+        if self.relative:
+            mask_step = RelativeActionsProcessorStep(
+                enabled=True,
+                exclude_joints=self.exclude_joints,
+                action_names=self.state_names,
+            )
+            processed_state = to_relative_actions(
+                state,
+                state[:, -1],
+                mask_step._build_mask(state.shape[-1]),
+                pose_representation=self.pose_representation,
+                se3_pose_groups=self.se3_pose_groups,
+            )
+
+        new_transition = transition.copy()
+        new_observation = dict(observation)
+        new_observation[OBS_STATE] = processed_state.flatten(start_dim=1)
+        new_transition[TransitionKey.OBSERVATION] = new_observation
+        return new_transition
+
+    def get_config(self) -> dict[str, Any]:
+        return {
+            "history_steps": self.history_steps,
+            "max_state_dim": self.max_state_dim,
+            "relative": self.relative,
+            "exclude_joints": self.exclude_joints,
+            "state_names": self.state_names,
+            "pose_representation": self.pose_representation,
+            "se3_pose_groups": self.se3_pose_groups,
+        }
+
+    def transform_features(
+        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
+    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
+        transformed = deepcopy(features)
+        for feature_group in transformed.values():
+            state_feature = feature_group.get(OBS_STATE)
+            if state_feature is not None and self.history_steps > 1:
+                state_dim = state_feature.shape[-1] * self.history_steps
+                feature_group[OBS_STATE] = PolicyFeature(type=state_feature.type, shape=(state_dim,))
+        return transformed
+
+
 @ProcessorStepRegistry.register(name="pi05_prepare_state_tokenizer_processor_step")
 @dataclass
 class Pi05PrepareStateTokenizerProcessorStep(ProcessorStep):
@@ -133,13 +278,28 @@ def make_pi05_pre_post_processors(
         enabled=config.use_relative_actions,
         exclude_joints=getattr(config, "relative_exclude_joints", []),
         action_names=getattr(config, "action_feature_names", None),
+        pose_representation=config.relative_pose_representation,
+        se3_pose_groups=config.relative_se3_pose_groups,
     )
 
     # OpenPI order: raw → relative → normalize → model → unnormalize → absolute
     input_steps: list[ProcessorStep] = [
         RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
         AddBatchDimensionProcessorStep(),
+        Pi05StateFromActionProcessorStep(
+            enabled=config.state_from_action,
+            history_steps=config.proprioception_history_steps,
+        ),
         relative_step,
+        Pi05FlattenStateHistoryProcessorStep(
+            history_steps=config.proprioception_history_steps,
+            max_state_dim=config.max_state_dim,
+            relative=config.use_relative_state_history,
+            exclude_joints=config.relative_state_exclude_joints,
+            state_names=config.action_feature_names,
+            pose_representation=config.relative_pose_representation,
+            se3_pose_groups=config.relative_se3_pose_groups,
+        ),
         # NOTE: NormalizerProcessorStep MUST come before Pi05PrepareStateTokenizerProcessorStep
         # because the tokenizer step expects normalized state in [-1, 1] range for discretization
         NormalizerProcessorStep(

src/lerobot/processor/__init__.py

@@ -90,7 +90,9 @@ from .relative_action_processor import (
     AbsoluteActionsProcessorStep,
     RelativeActionsProcessorStep,
     to_absolute_actions,
+    to_absolute_se3_pose,
     to_relative_actions,
+    to_relative_se3_pose,
 )
 from .rename_processor import RenameObservationsProcessorStep, rename_stats
 from .tokenizer_processor import ActionTokenizerProcessorStep, TokenizerProcessorStep
@@ -135,6 +137,10 @@ __all__ = [
     "make_default_robot_observation_processor",
     "AbsoluteActionsProcessorStep",
     "RelativeActionsProcessorStep",
+    "to_absolute_actions",
+    "to_absolute_se3_pose",
+    "to_relative_actions",
+    "to_relative_se3_pose",
     "MapDeltaActionToRobotActionStep",
     "MapTensorToDeltaActionDictStep",
     "NewLineTaskProcessorStep",
@@ -168,8 +174,6 @@ __all__ = [
     "transition_to_batch",
     "TransitionKey",
     "TruncatedProcessorStep",
-    "to_absolute_actions",
-    "to_relative_actions",
     "UnnormalizerProcessorStep",
     "VanillaObservationProcessorStep",
 ]

src/lerobot/processor/relative_action_processor.py

@@ -34,57 +34,206 @@ __all__ = [
     "AbsoluteActionsProcessorStep",
     "to_relative_actions",
     "to_absolute_actions",
+    "to_relative_se3_pose",
+    "to_absolute_se3_pose",
 ]
 
 
-def to_relative_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) -> Tensor:
-    """Convert absolute actions to relative: relative = action - state (for masked dims).
+def _rotvec_to_quaternion(rotvec: Tensor) -> Tensor:
+    angle = torch.linalg.vector_norm(rotvec, dim=-1, keepdim=True)
+    angle_sq = angle.square()
+    small_scale = 0.5 - angle_sq / 48.0 + angle_sq.square() / 3840.0
+    scale = torch.where(angle > 1e-6, torch.sin(angle / 2.0) / angle.clamp_min(1e-12), small_scale)
+    return torch.cat((torch.cos(angle / 2.0), rotvec * scale), dim=-1)
+
+
+def _quaternion_to_rotvec(quaternion: Tensor) -> Tensor:
+    quaternion = quaternion / torch.linalg.vector_norm(quaternion, dim=-1, keepdim=True).clamp_min(1e-12)
+    quaternion = quaternion * torch.where(quaternion[..., :1] < 0, -1.0, 1.0)
+    vector = quaternion[..., 1:]
+    sin_half_angle = torch.linalg.vector_norm(vector, dim=-1, keepdim=True)
+    angle = 2.0 * torch.atan2(sin_half_angle, quaternion[..., :1].clamp_min(0.0))
+    small_scale = 2.0 + sin_half_angle.square() / 3.0
+    scale = torch.where(
+        sin_half_angle > 1e-6,
+        angle / sin_half_angle.clamp_min(1e-12),
+        small_scale,
+    )
+    return vector * scale
+
+
+def _quaternion_multiply(left: Tensor, right: Tensor) -> Tensor:
+    left_w, left_xyz = left[..., :1], left[..., 1:]
+    right_w, right_xyz = right[..., :1], right[..., 1:]
+    return torch.cat(
+        (
+            left_w * right_w - (left_xyz * right_xyz).sum(dim=-1, keepdim=True),
+            left_w * right_xyz + right_w * left_xyz + torch.linalg.cross(left_xyz, right_xyz, dim=-1),
+        ),
+        dim=-1,
+    )
+
+
+def _quaternion_conjugate(quaternion: Tensor) -> Tensor:
+    return torch.cat((quaternion[..., :1], -quaternion[..., 1:]), dim=-1)
+
+
+def _quaternion_rotate(quaternion: Tensor, vector: Tensor) -> Tensor:
+    quaternion_xyz = quaternion[..., 1:]
+    uv = torch.linalg.cross(quaternion_xyz, vector, dim=-1)
+    uuv = torch.linalg.cross(quaternion_xyz, uv, dim=-1)
+    return vector + 2.0 * (quaternion[..., :1] * uv + uuv)
+
+
+def to_relative_se3_pose(target_pose: Tensor, reference_pose: Tensor) -> Tensor:
+    """Encode a pose as ``inv(T_reference) @ T_target``.
+
+    Poses use ``[x, y, z, rx, ry, rz]`` with an axis-angle rotation vector.
+    The relative translation is therefore expressed in the reference EE frame.
+    """
+    if target_pose.shape[-1] != 6 or reference_pose.shape[-1] != 6:
+        raise ValueError("SE(3) poses must have six values: xyz followed by a rotation vector")
+    reference_quaternion = _rotvec_to_quaternion(reference_pose[..., 3:])
+    target_quaternion = _rotvec_to_quaternion(target_pose[..., 3:])
+    inverse_reference_quaternion = _quaternion_conjugate(reference_quaternion)
+    relative_translation = _quaternion_rotate(
+        inverse_reference_quaternion, target_pose[..., :3] - reference_pose[..., :3]
+    )
+    relative_quaternion = _quaternion_multiply(inverse_reference_quaternion, target_quaternion)
+    return torch.cat((relative_translation, _quaternion_to_rotvec(relative_quaternion)), dim=-1)
+
+
+def to_absolute_se3_pose(relative_pose: Tensor, reference_pose: Tensor) -> Tensor:
+    """Decode a pose with ``T_target = T_reference @ T_relative``."""
+    if relative_pose.shape[-1] != 6 or reference_pose.shape[-1] != 6:
+        raise ValueError("SE(3) poses must have six values: xyz followed by a rotation vector")
+    reference_quaternion = _rotvec_to_quaternion(reference_pose[..., 3:])
+    relative_quaternion = _rotvec_to_quaternion(relative_pose[..., 3:])
+    target_translation = reference_pose[..., :3] + _quaternion_rotate(
+        reference_quaternion, relative_pose[..., :3]
+    )
+    target_quaternion = _quaternion_multiply(reference_quaternion, relative_quaternion)
+    return torch.cat((target_translation, _quaternion_to_rotvec(target_quaternion)), dim=-1)
+
+
+def _broadcast_reference(actions: Tensor, state: Tensor) -> Tensor:
+    if state.device != actions.device or state.dtype != actions.dtype:
+        state = state.to(device=actions.device, dtype=actions.dtype)
+    if actions.ndim == state.ndim + 1:
+        state = state.unsqueeze(-2)
+    return state
+
+
+def _validate_se3_pose_groups(
+    pose_representation: str,
+    se3_pose_groups: Sequence[Sequence[int]] | None,
+    mask: Sequence[bool],
+    action_dim: int,
+) -> list[list[int]]:
+    if pose_representation not in {"componentwise", "se3"}:
+        raise ValueError(
+            f"Unsupported pose_representation={pose_representation!r}; expected 'componentwise' or 'se3'"
+        )
+    if pose_representation == "componentwise":
+        return []
+    if not se3_pose_groups:
+        raise ValueError("pose_representation='se3' requires at least one six-index se3_pose_group")
+
+    normalized_groups: list[list[int]] = []
+    used_indices: set[int] = set()
+    for raw_group in se3_pose_groups:
+        group = [int(index) for index in raw_group]
+        if len(group) != 6:
+            raise ValueError(f"Each SE(3) pose group must contain six indices, got {group}")
+        if len(set(group)) != 6 or any(index < 0 or index >= action_dim for index in group):
+            raise ValueError(f"Invalid SE(3) pose group for action_dim={action_dim}: {group}")
+        if any(index >= len(mask) for index in group):
+            raise ValueError(f"SE(3) pose group lies outside the relative mask: {group}")
+        if used_indices.intersection(group):
+            raise ValueError(f"SE(3) pose groups must not overlap: {group}")
+        group_mask = [bool(mask[index]) for index in group]
+        if any(group_mask) and not all(group_mask):
+            raise ValueError(f"An SE(3) pose group must be wholly relative or wholly absolute: {group}")
+        used_indices.update(group)
+        if all(group_mask):
+            normalized_groups.append(group)
+    return normalized_groups
+
+
+def to_relative_actions(
+    actions: Tensor,
+    state: Tensor,
+    mask: Sequence[bool],
+    *,
+    pose_representation: str = "componentwise",
+    se3_pose_groups: Sequence[Sequence[int]] | None = None,
+) -> Tensor:
+    """Convert absolute actions to a configured relative representation.
+
+    Component-wise mode computes ``action - state``. SE(3) mode computes
+    ``inv(T_state) @ T_action`` for each configured pose group.
 
     Args:
         actions: (B, T, action_dim) or (B, action_dim).
         state: (B, state_dim). Broadcast across time dimension.
         mask: Which dims to convert. Can be shorter than action_dim.
     """
+    groups = _validate_se3_pose_groups(pose_representation, se3_pose_groups, mask, actions.shape[-1])
     mask_t = torch.tensor(mask, dtype=actions.dtype, device=actions.device)
     dims = mask_t.shape[0]
     # Align state to the same device/dtype as actions. _last_state is cached before
     # DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA.
-    if state.device != actions.device or state.dtype != actions.dtype:
-        state = state.to(device=actions.device, dtype=actions.dtype)
-    state_offset = state[..., :dims] * mask_t
-    if actions.ndim == 3:
-        state_offset = state_offset.unsqueeze(-2)
+    state = _broadcast_reference(actions, state)
+    component_mask = mask_t.clone()
+    for group in groups:
+        component_mask[group] = 0
+    state_offset = state[..., :dims] * component_mask
     actions = actions.clone()
     actions[..., :dims] -= state_offset
+    for group in groups:
+        actions[..., group] = to_relative_se3_pose(actions[..., group], state[..., group])
     return actions
 
 
-def to_absolute_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) -> Tensor:
-    """Convert relative actions back to absolute: absolute = relative + state (for masked dims).
+def to_absolute_actions(
+    actions: Tensor,
+    state: Tensor,
+    mask: Sequence[bool],
+    *,
+    pose_representation: str = "componentwise",
+    se3_pose_groups: Sequence[Sequence[int]] | None = None,
+) -> Tensor:
+    """Convert relative actions back to absolute actions.
+
+    Component-wise mode computes ``relative + state``. SE(3) mode computes
+    ``T_state @ T_relative`` for each configured pose group.
 
     Args:
         actions: (B, T, action_dim) or (B, action_dim).
         state: (B, state_dim). Broadcast across time dimension.
         mask: Which dims to convert. Can be shorter than action_dim.
     """
+    groups = _validate_se3_pose_groups(pose_representation, se3_pose_groups, mask, actions.shape[-1])
     mask_t = torch.tensor(mask, dtype=actions.dtype, device=actions.device)
     dims = mask_t.shape[0]
     # Align state to the same device/dtype as actions. _last_state is cached before
     # DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA.
-    if state.device != actions.device or state.dtype != actions.dtype:
-        state = state.to(device=actions.device, dtype=actions.dtype)
-    state_offset = state[..., :dims] * mask_t
-    if actions.ndim == 3:
-        state_offset = state_offset.unsqueeze(-2)
+    state = _broadcast_reference(actions, state)
+    component_mask = mask_t.clone()
+    for group in groups:
+        component_mask[group] = 0
+    state_offset = state[..., :dims] * component_mask
     actions = actions.clone()
     actions[..., :dims] += state_offset
+    for group in groups:
+        actions[..., group] = to_absolute_se3_pose(actions[..., group], state[..., group])
     return actions
 
 
 @ProcessorStepRegistry.register("relative_actions_processor")
 @dataclass
 class RelativeActionsProcessorStep(ProcessorStep):
-    """Converts absolute actions to relative actions (action -= state) for masked dimensions.
+    """Converts absolute actions to the configured relative representation.
 
     Mirrors OpenPI's DeltaActions transform. Applied during preprocessing so the model
     trains on relative offsets instead of absolute positions.
@@ -101,6 +250,8 @@ class RelativeActionsProcessorStep(ProcessorStep):
     enabled: bool = False
     exclude_joints: list[str] = field(default_factory=list)
     action_names: list[str] | None = None
+    pose_representation: str = "componentwise"
+    se3_pose_groups: list[list[int]] = field(default_factory=list)
     _last_state: torch.Tensor | None = field(default=None, init=False, repr=False)
 
     def _build_mask(self, action_dim: int) -> list[bool]:
@@ -126,20 +277,30 @@ class RelativeActionsProcessorStep(ProcessorStep):
         observation = transition.get(TransitionKey.OBSERVATION, {})
         state = observation.get(OBS_STATE) if observation else None
 
+        # State history has shape (B, H, D). Relative actions are referenced to
+        # the newest proprioceptive state, not the whole history tensor.
+        reference_state = state[:, -1] if state is not None and state.ndim == 3 else state
+
         # Always cache state for the paired AbsoluteActionsProcessorStep
-        if state is not None:
-            self._last_state = state
+        if reference_state is not None:
+            self._last_state = reference_state
 
         if not self.enabled:
             return transition
 
         new_transition = transition.copy()
         action = new_transition.get(TransitionKey.ACTION)
-        if action is None or state is None:
+        if action is None or reference_state is None:
             return new_transition
 
         mask = self._build_mask(action.shape[-1])
-        new_transition[TransitionKey.ACTION] = to_relative_actions(action, state, mask)
+        new_transition[TransitionKey.ACTION] = to_relative_actions(
+            action,
+            reference_state,
+            mask,
+            pose_representation=self.pose_representation,
+            se3_pose_groups=self.se3_pose_groups,
+        )
         return new_transition
 
     def get_cached_state(self) -> torch.Tensor | None:
@@ -151,6 +312,8 @@ class RelativeActionsProcessorStep(ProcessorStep):
             "enabled": self.enabled,
             "exclude_joints": self.exclude_joints,
             "action_names": self.action_names,
+            "pose_representation": self.pose_representation,
+            "se3_pose_groups": self.se3_pose_groups,
         }
 
     def transform_features(
@@ -199,7 +362,13 @@ class AbsoluteActionsProcessorStep(ProcessorStep):
             return new_transition
 
         mask = self.relative_step._build_mask(action.shape[-1])
-        new_transition[TransitionKey.ACTION] = to_absolute_actions(action, cached_state, mask)
+        new_transition[TransitionKey.ACTION] = to_absolute_actions(
+            action,
+            cached_state,
+            mask,
+            pose_representation=self.relative_step.pose_representation,
+            se3_pose_groups=self.relative_step.se3_pose_groups,
+        )
         return new_transition
 
     def get_config(self) -> dict[str, Any]:

src/lerobot/scripts/lerobot_edit_dataset.py

@@ -167,9 +167,7 @@ Show dataset information without feature details:
         --operation.type info \
         --operation.show_features false
 
-Recompute dataset statistics (saves to lerobot/pusht_recomputed_stats by default). The source
-dataset is never modified: large files are symlinked and only meta/ is copied, so this also works
-on read-only source datasets:
+Recompute dataset statistics (saves to lerobot/pusht_recomputed_stats by default):
     lerobot-edit-dataset \
         --repo_id lerobot/pusht \
         --operation.type recompute_stats
@@ -180,19 +178,6 @@ Recompute stats and save to a specific new repo_id:
         --new_repo_id lerobot/pusht_new_stats \
         --operation.type recompute_stats
 
-Recompute stats including image/video features (samples and decodes frames from each episode):
-    lerobot-edit-dataset \
-        --repo_id lerobot/pusht \
-        --operation.type recompute_stats \
-        --operation.skip_image_video false
-
-Recompute stats and also rewrite the per-episode stats in the episodes parquet (keeps
-meta/stats.json and the per-episode stats consistent):
-    lerobot-edit-dataset \
-        --repo_id lerobot/pusht \
-        --operation.type recompute_stats \
-        --operation.update_episode_stats true
-
 Recompute stats in-place (overwrites original dataset stats):
     lerobot-edit-dataset \
         --repo_id lerobot/pusht \
@@ -340,7 +325,8 @@ class RecomputeStatsConfig(OperationConfig):
     relative_exclude_joints: list[str] | None = None
     chunk_size: int = 50
     num_workers: int = 0
-    update_episode_stats: bool = False
+    relative_pose_representation: str = "componentwise"
+    relative_se3_pose_groups: list[list[int]] | None = None
     overwrite: bool = False
 
 
@@ -393,30 +379,6 @@ def _resolve_io_paths(
     return output_repo_id, input_path, output_path
 
 
-def _reference_copy_dataset(input_root: Path, output_root: Path) -> None:
-    """Create a lightweight copy of a dataset that never modifies the source.
-
-    The directory tree is recreated with real directories, and every file is
-    symlinked to its source counterpart so no data is duplicated and the source is
-    only ever read. Files under ``meta/`` are instead copied as real, writable files
-    so that stats/info can be rewritten without touching the original. Symlinking
-    individual files (rather than whole directories) keeps ``push_to_hub`` working,
-    since ``Path.glob`` follows file symlinks but does not descend into symlinked
-    directories. This makes the operation safe on read-only source datasets.
-    """
-    for src in input_root.rglob("*"):
-        rel = src.relative_to(input_root)
-        dst = output_root / rel
-        if src.is_dir():
-            dst.mkdir(parents=True, exist_ok=True)
-        elif rel.parts[0] == "meta":
-            dst.parent.mkdir(parents=True, exist_ok=True)
-            shutil.copyfile(src, dst)  # copyfile ignores source perms, so dst is writable
-        else:
-            dst.parent.mkdir(parents=True, exist_ok=True)
-            dst.symlink_to(src.resolve())
-
-
 def get_output_path(
     repo_id: str,
     new_repo_id: str | None,
@@ -714,18 +676,14 @@ def handle_recompute_stats(cfg: EditDatasetConfig) -> None:
         )
         dataset = LeRobotDataset(cfg.repo_id, root=input_root)
     else:
-        logging.info(f"Referencing dataset from {input_root} into {output_root} (source is left untouched)")
+        logging.info(f"Copying dataset from {input_root} to {output_root}")
         if output_root.exists():
             backup_path = output_root.with_name(output_root.name + "_old")
             logging.warning(f"Output directory {output_root} already exists. Moving to {backup_path}")
             if backup_path.exists():
                 shutil.rmtree(backup_path)
             shutil.move(output_root, backup_path)
-        # recompute_stats only reads data/ and rewrites files under meta/ (stats.json, and
-        # the episodes parquet when update_episode_stats is set), so symlink the large
-        # immutable files and copy only meta/. This avoids duplicating the dataset and works
-        # even when the source dataset is read-only.
-        _reference_copy_dataset(input_root, output_root)
+        shutil.copytree(input_root, output_root)
         dataset = LeRobotDataset(output_repo_id, root=output_root)
 
     logging.info(f"Recomputing stats for {cfg.repo_id}")
@@ -742,7 +700,8 @@ def handle_recompute_stats(cfg: EditDatasetConfig) -> None:
         relative_exclude_joints=cfg.operation.relative_exclude_joints,
         chunk_size=cfg.operation.chunk_size,
         num_workers=cfg.operation.num_workers,
-        update_episode_stats=cfg.operation.update_episode_stats,
+        relative_pose_representation=cfg.operation.relative_pose_representation,
+        relative_se3_pose_groups=cfg.operation.relative_se3_pose_groups,
     )
 
     logging.info(f"Stats written to {dataset.root}")

tests/processor/test_pi05_state_from_action.py

@@ -0,0 +1,227 @@
+#!/usr/bin/env python
+
+# 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.
+
+from math import pi
+
+import numpy as np
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from lerobot.datasets.compute_stats import (  # noqa: E402
+    compute_relative_action_stats,
+    compute_state_history_stats,
+)
+from lerobot.policies.pi05.configuration_pi05 import PI05Config  # noqa: E402
+from lerobot.policies.pi05.processor_pi05 import (  # noqa: E402
+    Pi05FlattenStateHistoryProcessorStep,
+    Pi05StateFromActionProcessorStep,
+)
+from lerobot.processor.relative_action_processor import (  # noqa: E402
+    AbsoluteActionsProcessorStep,
+    RelativeActionsProcessorStep,
+)
+from lerobot.types import TransitionKey  # noqa: E402
+from lerobot.utils.constants import OBS_STATE  # noqa: E402
+
+
+def _transition(action: torch.Tensor | None, state: torch.Tensor | None = None) -> dict:
+    observation = {} if state is None else {OBS_STATE: state}
+    return {
+        TransitionKey.OBSERVATION: observation,
+        TransitionKey.ACTION: action,
+        TransitionKey.REWARD: None,
+        TransitionKey.DONE: None,
+        TransitionKey.TRUNCATED: None,
+        TransitionKey.COMPLEMENTARY_DATA: {},
+    }
+
+
+def test_pi05_config_requests_action_history_prefix():
+    config = PI05Config(
+        device="cpu",
+        chunk_size=4,
+        n_action_steps=4,
+        state_from_action=True,
+        proprioception_history_steps=2,
+    )
+
+    assert config.action_delta_indices == [-1, 0, 1, 2, 3]
+
+
+def test_state_from_action_extracts_history_and_preserves_target_horizon():
+    action = torch.arange(2 * 5 * 3, dtype=torch.float32).reshape(2, 5, 3)
+    step = Pi05StateFromActionProcessorStep(enabled=True, history_steps=2)
+
+    result = step(_transition(action))
+
+    torch.testing.assert_close(result[TransitionKey.OBSERVATION][OBS_STATE], action[:, :2])
+    torch.testing.assert_close(result[TransitionKey.ACTION], action[:, 1:])
+
+
+def test_relative_actions_use_newest_state_in_history_and_roundtrip():
+    state_history = torch.tensor([[[1.0, 10.0], [2.0, 20.0]]])
+    absolute = torch.tensor([[[3.0, 30.0], [4.0, 40.0]]])
+    relative_step = RelativeActionsProcessorStep(enabled=True)
+    absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
+
+    relative = relative_step(_transition(absolute, state_history))
+    expected = torch.tensor([[[1.0, 10.0], [2.0, 20.0]]])
+    torch.testing.assert_close(relative[TransitionKey.ACTION], expected)
+
+    recovered = absolute_step(_transition(relative[TransitionKey.ACTION]))
+    torch.testing.assert_close(recovered[TransitionKey.ACTION], absolute)
+
+
+def test_flatten_state_history_preserves_chronological_order():
+    state_history = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]])
+    step = Pi05FlattenStateHistoryProcessorStep(history_steps=2, max_state_dim=4)
+
+    result = step(_transition(torch.zeros(1, 2, 2), state_history))
+
+    torch.testing.assert_close(
+        result[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[1.0, 2.0, 3.0, 4.0]])
+    )
+
+
+def test_state_history_can_be_relative_with_absolute_gripper():
+    state_history = torch.tensor([[[1.0, 10.0, 0.2], [3.0, 20.0, 0.4]]])
+    step = Pi05FlattenStateHistoryProcessorStep(
+        history_steps=2,
+        max_state_dim=6,
+        relative=True,
+        exclude_joints=["gripper"],
+        state_names=["x", "y", "gripper_width"],
+    )
+
+    result = step(_transition(torch.zeros(1, 2, 3), state_history))
+
+    torch.testing.assert_close(
+        result[TransitionKey.OBSERVATION][OBS_STATE],
+        torch.tensor([[-2.0, -10.0, 0.2, 0.0, 0.0, 0.4]]),
+    )
+
+
+def test_state_history_can_use_se3_composition_with_absolute_gripper():
+    state_history = torch.tensor(
+        [[[0.0, 1.0, 0.0, 0.0, 0.0, pi / 2, 0.2], [0.0, 0.0, 0.0, 0.0, 0.0, pi / 2, 0.4]]]
+    )
+    step = Pi05FlattenStateHistoryProcessorStep(
+        history_steps=2,
+        max_state_dim=14,
+        relative=True,
+        exclude_joints=["gripper"],
+        state_names=["x", "y", "z", "rx", "ry", "rz", "gripper_width"],
+        pose_representation="se3",
+        se3_pose_groups=[list(range(6))],
+    )
+
+    result = step(_transition(torch.zeros(1, 2, 7), state_history))
+
+    expected = torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.2, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.4]])
+    torch.testing.assert_close(result[TransitionKey.OBSERVATION][OBS_STATE], expected, atol=1e-6, rtol=1e-6)
+
+
+def test_inference_state_history_is_rolled_and_reset():
+    step = Pi05StateFromActionProcessorStep(enabled=True, history_steps=2)
+
+    first = step(_transition(None, torch.tensor([[1.0, 2.0]])))
+    second = step(_transition(None, torch.tensor([[3.0, 4.0]])))
+    torch.testing.assert_close(
+        first[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[[1.0, 2.0], [1.0, 2.0]]])
+    )
+    torch.testing.assert_close(
+        second[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[[1.0, 2.0], [3.0, 4.0]]])
+    )
+
+    step.reset()
+    reset = step(_transition(None, torch.tensor([[5.0, 6.0]])))
+    torch.testing.assert_close(
+        reset[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[[5.0, 6.0], [5.0, 6.0]]])
+    )
+
+
+def test_flatten_state_history_checks_max_state_dim():
+    step = Pi05FlattenStateHistoryProcessorStep(history_steps=2, max_state_dim=3)
+
+    with pytest.raises(ValueError, match="above max_state_dim"):
+        step(_transition(torch.zeros(1, 2, 2), torch.zeros(1, 2, 2)))
+
+
+def test_relative_stats_can_use_absolute_action_as_state():
+    actions = np.asarray([[0.0, 0.0], [1.0, 2.0], [2.0, 4.0], [3.0, 6.0]], dtype=np.float32)
+    dataset = {"action": actions, "episode_index": np.zeros(4, dtype=np.int64)}
+    features = {"action": {"shape": [2], "names": ["x", "y"]}}
+
+    stats = compute_relative_action_stats(
+        dataset,
+        features,
+        chunk_size=2,
+        state_from_action=True,
+    )
+
+    np.testing.assert_allclose(stats["mean"], [0.5, 1.0])
+
+
+def test_relative_state_history_stats_match_processor_representation():
+    actions = np.asarray(
+        [[0.0, 0.1], [1.0, 0.2], [3.0, 0.3]],
+        dtype=np.float32,
+    )
+    dataset = {"action": actions, "episode_index": np.zeros(3, dtype=np.int64)}
+    features = {"action": {"shape": [2], "names": ["x", "gripper_width"]}}
+
+    stats = compute_state_history_stats(
+        dataset,
+        features,
+        history_steps=2,
+        exclude_joints=["gripper"],
+        relative=True,
+    )
+
+    expected = np.asarray([[0.0, 0.1, 0.0, 0.1], [-1.0, 0.1, 0.0, 0.2], [-2.0, 0.2, 0.0, 0.3]])
+    np.testing.assert_allclose(stats["mean"], expected.mean(axis=0))
+
+
+def test_se3_relative_action_stats_use_reference_frame():
+    actions = np.asarray(
+        [
+            [0.0, 0.0, 0.0, 0.0, 0.0, pi / 2, 0.2],
+            [0.0, 1.0, 0.0, 0.0, 0.0, pi / 2, 0.3],
+        ],
+        dtype=np.float32,
+    )
+    dataset = {"action": actions, "episode_index": np.zeros(2, dtype=np.int64)}
+    features = {
+        "action": {
+            "shape": [7],
+            "names": ["x", "y", "z", "rx", "ry", "rz", "gripper_width"],
+        }
+    }
+
+    stats = compute_relative_action_stats(
+        dataset,
+        features,
+        chunk_size=2,
+        exclude_joints=["gripper"],
+        state_from_action=True,
+        pose_representation="se3",
+        se3_pose_groups=[list(range(6))],
+    )
+
+    np.testing.assert_allclose(stats["mean"][:3], [0.5, 0.0, 0.0], atol=1e-6)
+    np.testing.assert_allclose(stats["mean"][6], 0.25, atol=1e-6)

tests/processor/test_se3_relative_actions.py

@@ -0,0 +1,70 @@
+import math
+
+import pytest
+import torch
+
+from lerobot.processor.relative_action_processor import (
+    to_absolute_actions,
+    to_absolute_se3_pose,
+    to_relative_actions,
+    to_relative_se3_pose,
+)
+
+POSE_GROUP = [list(range(6))]
+
+
+def test_se3_translation_is_expressed_in_reference_frame():
+    reference = torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0, math.pi / 2]])
+    target = torch.tensor([[1.0, 3.0, 3.0, 0.0, 0.0, math.pi / 2]])
+
+    relative = to_relative_se3_pose(target, reference)
+
+    torch.testing.assert_close(relative, torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]]), atol=1e-6, rtol=1e-6)
+
+
+def test_se3_pose_roundtrip_for_batched_chunks():
+    torch.manual_seed(0)
+    reference = torch.randn(4, 6)
+    reference[:, 3:] *= 0.8
+    target = torch.randn(4, 11, 6)
+    target[..., 3:] *= 0.8
+
+    relative = to_relative_se3_pose(target, reference.unsqueeze(1))
+    recovered = to_absolute_se3_pose(relative, reference.unsqueeze(1))
+
+    torch.testing.assert_close(recovered, target, atol=2e-5, rtol=2e-5)
+
+
+def test_mixed_se3_pose_and_absolute_gripper_roundtrip():
+    reference = torch.tensor([[0.2, -0.1, 0.4, 0.1, 0.2, -0.3, 0.06]])
+    target = torch.tensor([[[0.3, 0.2, 0.5, -0.2, 0.1, 0.4, 0.03], [0.1, -0.3, 0.2, 0.5, -0.1, 0.2, 0.05]]])
+    mask = [True, True, True, True, True, True, False]
+
+    relative = to_relative_actions(
+        target,
+        reference,
+        mask,
+        pose_representation="se3",
+        se3_pose_groups=POSE_GROUP,
+    )
+    recovered = to_absolute_actions(
+        relative,
+        reference,
+        mask,
+        pose_representation="se3",
+        se3_pose_groups=POSE_GROUP,
+    )
+
+    torch.testing.assert_close(relative[..., 6], target[..., 6])
+    torch.testing.assert_close(recovered, target, atol=2e-5, rtol=2e-5)
+
+
+def test_se3_pose_group_cannot_be_partially_relative():
+    with pytest.raises(ValueError, match="wholly relative or wholly absolute"):
+        to_relative_actions(
+            torch.zeros(1, 7),
+            torch.zeros(1, 7),
+            [True, True, True, False, False, False, False],
+            pose_representation="se3",
+            se3_pose_groups=POSE_GROUP,
+        )