diff --git a/README.md b/README.md
index 64b9a0f..5e6093b 100644
--- a/README.md
+++ b/README.md
@@ -64,22 +64,53 @@ amp_rsl_rl/
 
 ## 📁 Dataset Structure
 
-The AMP-RSL-RL framework expects motion capture datasets in `.npy` format. Each `.npy` file must contain a Python dictionary with the following keys:
+AMP-RSL-RL expects each motion file to be a `.npy` containing a Python dictionary.
+
+### Required keys (always used)
 
 - **`joints_list`**: `List[str]`  
-  A list of joint names. These should correspond to the joint order expected by the agent.
+  Joint names in the same order expected by the robot configuration.
+
+- **`joint_positions`**: array-like, shape `(T, N_joints)`  
+  Joint positions over time.
+
+- **`root_position`**: array-like, shape `(T, 3)`  
+  Base position in world frame.
+
+- **`root_quaternion`**: array-like, shape `(T, 4)`  
+  Base orientation in quaternion **`xyzw`** format (SciPy convention).
+
+- **`fps`**: `float` (or scalar array)  
+  Original dataset framerate. The loader resamples to simulator `dt`.
+
+### Optional keys (needed for body keyframe observations)
+
+These keys are required only if you enable any of:
+`body_pos_b`, `body_ori_b`, `body_lin_vel_b`, `body_ang_vel_b` in `amp_obs_components`.
+
+- **`body_links_list`**: `List[str]`  
+  Names of links for which body keyframes are provided.
 
-- **`joint_positions`**: `List[np.ndarray]`  
-  A list where each element is a NumPy array representing the joint positions at a frame. All arrays should have the same shape `(N,)`, where `N` is the number of joints.
+- **`body_links_pos_w`**: array-like, shape `(T, N_bodies, 3)`  
+  Body positions in world frame.
 
-- **`root_position`**: `List[np.ndarray]`  
-  A list of 3D vectors representing the position of the base (root) of the agent in world coordinates for each frame.
+- **`body_links_quat_w`**: array-like, shape `(T, N_bodies, 4)`  
+  Body orientations in world frame, quaternion **`xyzw`** format.
 
-- **`root_quaternion`**: `List[np.ndarray]`  
-  A list of unit quaternions in **`xyzw`** format (SciPy convention), representing the base orientation of the agent for each frame.
+### Configuration fields for body keyframes
 
-- **`fps`**: `float`  
-  The number of frames per second in the original dataset. This is used to resample the data to match the simulator's timestep.
+When using body keyframe components, set these in `dataset_cfg`:
+
+- **`amp_obs_components`**: ordered list of discriminator components to concatenate.
+- **`body_links_names`**: ordered list of body names to use from the dataset.
+- **`anchor_body_name`**: reference body used to express relative body observations.
+
+Notes:
+
+- `anchor_body_name` must be present in `body_links_names`.
+- The anchor body is excluded from the final body observation vectors.
+- For backward compatibility, if `amp_obs_components` is not provided, the default is
+  `["joint_pos", "joint_vel", "base_lin_vel", "base_ang_vel"]`.
 
 ---
 
@@ -100,19 +131,44 @@ training remain consistent with their symmetric counterparts.
 
 ### Example
 
-Here’s an example of how the structure might look when loaded in Python:
+Example dataset dictionary with optional body keyframes:
 
 ```python
 {
     "joints_list": ["hip", "knee", "ankle"],
-    "joint_positions": [np.array([0.1, -0.2, 0.3]), np.array([0.11, -0.21, 0.31]), ...],
-    "root_position": [np.array([0.0, 0.0, 1.0]), np.array([0.01, 0.0, 1.0]), ...],
-    "root_quaternion": [np.array([0.0, 0.0, 0.0, 1.0]), np.array([0.0, 0.0, 0.1, 0.99]), ...],
-    "fps": 120.0
+    "joint_positions": np.ndarray(shape=(T, 3)),
+    "root_position": np.ndarray(shape=(T, 3)),
+    "root_quaternion": np.ndarray(shape=(T, 4)),  # xyzw
+    "fps": 120.0,
+
+    # Optional: required only for body_* discriminator components
+    "body_links_list": ["root_link", "left_knee", "right_knee"],
+    "body_links_pos_w": np.ndarray(shape=(T, 3, 3)),
+    "body_links_quat_w": np.ndarray(shape=(T, 3, 4)),  # xyzw
 }
 ```
 
-All lists must have the same number of entries (i.e. one per frame). The dataset should represent smooth motion captured over time.
+`T` is the number of frames. All time-dependent fields must share the same `T`.
+
+Example `dataset_cfg` enabling body keyframe observations:
+
+```python
+dataset_cfg = {
+    "amp_data_path": "/path/to/datasets",
+    "datasets": {"walk": 1.0},
+    "slow_down_factor": 1,
+    "amp_obs_components": [
+        "joint_pos",
+        "joint_vel",
+        "base_lin_vel",
+        "base_ang_vel",
+        "body_pos_b",
+        "body_ori_b",
+    ],
+    "body_links_names": ["root_link", "left_knee", "right_knee"],
+    "anchor_body_name": "root_link",
+}
+```
 
 ---
 
diff --git a/amp_rsl_rl/runners/amp_on_policy_runner.py b/amp_rsl_rl/runners/amp_on_policy_runner.py
index b4827c7..fa746d2 100644
--- a/amp_rsl_rl/runners/amp_on_policy_runner.py
+++ b/amp_rsl_rl/runners/amp_on_policy_runner.py
@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Generative Bionics S.R.L.
+# SPDX-License-Identifier: BSD-3-Clause
+
 # Copyright (c) 2025, Istituto Italiano di Tecnologia
 # All rights reserved.
 #
@@ -308,6 +311,9 @@ def __init__(self, env: VecEnv, train_cfg, log_dir=None, device="cpu"):
             expected_joint_names=amp_joint_names,
             velocity_representation=velocity_representation,
             symmetry_cfg=self.alg_cfg.get("symmetry_cfg"),
+            amp_obs_components=self.dataset_cfg.get("amp_obs_components"),
+            body_links_names=self.dataset_cfg.get("body_links_names"),
+            anchor_body_name=self.dataset_cfg.get("anchor_body_name"),
         )
 
         self.discriminator = Discriminator(
diff --git a/amp_rsl_rl/utils/motion_loader.py b/amp_rsl_rl/utils/motion_loader.py
index 8632913..f032f32 100644
--- a/amp_rsl_rl/utils/motion_loader.py
+++ b/amp_rsl_rl/utils/motion_loader.py
@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Generative Bionics S.R.L.
+# SPDX-License-Identifier: BSD-3-Clause
+
 # Copyright (c) 2025, Istituto Italiano di Tecnologia
 # All rights reserved.
 #
@@ -107,6 +110,9 @@ def _call_augmentation_func(
     return result, None
 
 
+_DEFAULT_OBS_COMPONENTS = ["joint_pos", "joint_vel", "base_lin_vel", "base_ang_vel"]
+
+
 @dataclass
 class MotionData:
     """
@@ -124,6 +130,10 @@ class MotionData:
         - base_lin_velocities_local: linear velocity in local (body) frame
         - base_ang_velocities_local: (currently zeros)
         - base_quat: orientation quaternion as torch.Tensor in wxyz order
+        - body_pos_b: (optional) relative body positions w.r.t. anchor, shape (T, (N_bodies-1)*3)
+        - body_ori_b: (optional) relative body orientations w.r.t. anchor (6D repr), shape (T, (N_bodies-1)*6)
+        - body_lin_vel_b: (optional) relative body linear velocities, shape (T, (N_bodies-1)*3)
+        - body_ang_vel_b: (optional) relative body angular velocities, shape (T, (N_bodies-1)*3)
 
     Notes:
         - The quaternion is expected in the dataset as `xyzw` format (SciPy default),
@@ -139,6 +149,10 @@ class MotionData:
     base_ang_velocities_local: Union[torch.Tensor, np.ndarray]
     base_quat: Union[Rotation, torch.Tensor]
     device: torch.device = torch.device("cpu")
+    body_pos_b: Optional[Union[torch.Tensor, np.ndarray]] = None
+    body_ori_b: Optional[Union[torch.Tensor, np.ndarray]] = None
+    body_lin_vel_b: Optional[Union[torch.Tensor, np.ndarray]] = None
+    body_ang_vel_b: Optional[Union[torch.Tensor, np.ndarray]] = None
 
     def __post_init__(self) -> None:
         # Convert numpy arrays (or SciPy Rotations) to torch tensors
@@ -165,6 +179,14 @@ def to_tensor(x):
                 device=self.device,
                 dtype=torch.float32,
             )
+        if isinstance(self.body_pos_b, np.ndarray):
+            self.body_pos_b = to_tensor(self.body_pos_b)
+        if isinstance(self.body_ori_b, np.ndarray):
+            self.body_ori_b = to_tensor(self.body_ori_b)
+        if isinstance(self.body_lin_vel_b, np.ndarray):
+            self.body_lin_vel_b = to_tensor(self.body_lin_vel_b)
+        if isinstance(self.body_ang_vel_b, np.ndarray):
+            self.body_ang_vel_b = to_tensor(self.body_ang_vel_b)
 
     def __len__(self) -> int:
         return self.joint_positions.shape[0]
@@ -173,6 +195,7 @@ def get_amp_dataset_obs(
         self,
         indices: torch.Tensor,
         velocity_representation: VelocityRepresentation = VelocityRepresentation.BODY_FIXED_REPRESENTATION,
+        obs_components: Optional[List[str]] = None,
     ) -> torch.Tensor:
         """
         Returns the AMP observation tensor for given indices.
@@ -181,30 +204,52 @@ def get_amp_dataset_obs(
             indices: indices of samples to retrieve
             velocity_representation: which frame convention to use for the base
                 velocities in the observation. Defaults to BODY_FIXED_REPRESENTATION.
+            obs_components: list of observation component names to include.
+                Valid values: "joint_pos", "joint_vel", "base_lin_vel", "base_ang_vel",
+                "body_pos_b", "body_ori_b", "body_lin_vel_b", "body_ang_vel_b".
+                If None, defaults to ["joint_pos", "joint_vel", "base_lin_vel", "base_ang_vel"].
 
         Returns:
             Concatenated observation tensor
         """
-        if velocity_representation == VelocityRepresentation.MIXED_REPRESENTATION:
-            return torch.cat(
-                (
-                    self.joint_positions[indices],
-                    self.joint_velocities[indices],
-                    self.base_lin_velocities_mixed[indices],
-                    self.base_ang_velocities_mixed[indices],
-                ),
-                dim=1,
-            )
-        return torch.cat(
-            (
-                self.joint_positions[indices],
-                self.joint_velocities[indices],
-                self.base_lin_velocities_local[indices],
-                self.base_ang_velocities_local[indices],
-            ),
-            dim=1,
+        if obs_components is None:
+            obs_components = _DEFAULT_OBS_COMPONENTS
+
+        use_mixed = (
+            velocity_representation == VelocityRepresentation.MIXED_REPRESENTATION
         )
 
+        component_map = {
+            "joint_pos": self.joint_positions,
+            "joint_vel": self.joint_velocities,
+            "base_lin_vel": (
+                self.base_lin_velocities_mixed
+                if use_mixed
+                else self.base_lin_velocities_local
+            ),
+            "base_ang_vel": (
+                self.base_ang_velocities_mixed
+                if use_mixed
+                else self.base_ang_velocities_local
+            ),
+            "body_pos_b": self.body_pos_b,
+            "body_ori_b": self.body_ori_b,
+            "body_lin_vel_b": self.body_lin_vel_b,
+            "body_ang_vel_b": self.body_ang_vel_b,
+        }
+
+        parts = []
+        for comp in obs_components:
+            tensor = component_map.get(comp)
+            if tensor is None:
+                raise ValueError(
+                    f"Observation component '{comp}' is not available in this MotionData. "
+                    f"Available components: {[k for k, v in component_map.items() if v is not None]}"
+                )
+            parts.append(tensor[indices])
+
+        return torch.cat(parts, dim=1)
+
     def get_state_for_reset(
         self,
         indices: torch.Tensor,
@@ -285,12 +330,42 @@ def __init__(
         expected_joint_names: Union[List[str], None] = None,
         symmetry_cfg: Optional[Dict[str, Any]] = None,
         velocity_representation: VelocityRepresentation = VelocityRepresentation.BODY_FIXED_REPRESENTATION,
+        amp_obs_components: Optional[List[str]] = None,
+        body_links_names: Optional[List[str]] = None,
+        anchor_body_name: Optional[str] = None,
     ) -> None:
         self.device = device
         self.velocity_representation = velocity_representation
+        self.obs_components = amp_obs_components
         if isinstance(dataset_path_root, str):
             dataset_path_root = Path(dataset_path_root)
 
+        # ─── Validate body keyframe configuration ───
+        _body_components = {"body_pos_b", "body_ori_b", "body_lin_vel_b", "body_ang_vel_b"}
+        self._needs_body_keyframes = (
+            self.obs_components is not None
+            and bool(_body_components & set(self.obs_components))
+        )
+        if self._needs_body_keyframes:
+            if body_links_names is None:
+                raise ValueError(
+                    "body_links_names must be provided when using body keyframe "
+                    f"observation components: {_body_components & set(self.obs_components)}"
+                )
+            if anchor_body_name is None:
+                raise ValueError(
+                    "anchor_body_name must be provided when using body keyframe "
+                    "observation components."
+                )
+            if anchor_body_name not in body_links_names:
+                raise ValueError(
+                    f"anchor_body_name '{anchor_body_name}' must be present in "
+                    f"body_links_names: {body_links_names}"
+                )
+        self.body_links_names = body_links_names
+        self.anchor_body_name = anchor_body_name
+        # ─────────────────────────────────────────────
+
         # ─── Check symmetry augmentation configuration ───
         if symmetry_cfg is not None:
             aug_fn = symmetry_cfg.get("amp_dataset_augmentation_func", None)
@@ -332,6 +407,8 @@ def __init__(
                 simulation_dt,
                 slow_down_factor,
                 expected_joint_names,
+                body_links_names=self.body_links_names,
+                anchor_body_name=self.anchor_body_name,
             )
             self.motion_data.append(md)
 
@@ -346,9 +423,13 @@ def __init__(
         for data, w in zip(self.motion_data, self.dataset_weights):
             T = len(data)
             idx = torch.arange(T, device=self.device)
-            obs = data.get_amp_dataset_obs(idx, self.velocity_representation)
+            obs = data.get_amp_dataset_obs(
+                idx, self.velocity_representation, self.obs_components
+            )
             next_idx = torch.clamp(idx + 1, max=T - 1)
-            next_obs = data.get_amp_dataset_obs(next_idx, self.velocity_representation)
+            next_obs = data.get_amp_dataset_obs(
+                next_idx, self.velocity_representation, self.obs_components
+            )
 
             if self.symmetry_cfg and self.symmetry_cfg.get(
                 "use_amp_dataset_augmentation", False
@@ -457,10 +538,22 @@ def load_data(
         simulation_dt: float,
         slow_down_factor: int = 1,
         expected_joint_names: Union[List[str], None] = None,
+        body_links_names: Optional[List[str]] = None,
+        anchor_body_name: Optional[str] = None,
     ) -> MotionData:
         """
         Loads and processes one motion dataset.
 
+        Args:
+            dataset_path: Path to the .npy motion file.
+            simulation_dt: Target simulation timestep.
+            slow_down_factor: Integer factor to slow down original data.
+            expected_joint_names: Joint ordering to use.
+            body_links_names: (Optional) List of body link names to extract from
+                the dataset. Required when using body keyframe observations.
+            anchor_body_name: (Optional) Name of the anchor body in body_links_names.
+                Required when using body keyframe observations.
+
         Returns:
             MotionData instance
         """
@@ -526,6 +619,98 @@ def load_data(
             resampled_base_orientations, simulation_dt, local=True
         )
 
+        # ─── Process body keyframes if requested ───
+        body_pos_b = None
+        body_ori_b = None
+        body_lin_vel_b = None
+        body_ang_vel_b = None
+
+        if self._needs_body_keyframes and body_links_names is not None:
+            if "body_links_list" not in data:
+                raise KeyError(
+                    f"Dataset '{dataset_path}' does not contain 'body_links_list'. "
+                    "Body keyframe observation components require datasets with "
+                    "'body_links_list', 'body_links_pos_w', and 'body_links_quat_w' keys."
+                )
+
+            dataset_body_names = list(data["body_links_list"])
+
+            # Build index map for body_links_names
+            body_idx_map: List[int] = []
+            for name in body_links_names:
+                if name not in dataset_body_names:
+                    raise ValueError(
+                        f"Body link '{name}' not found in dataset '{dataset_path}'. "
+                        f"Available body links: {dataset_body_names}"
+                    )
+                body_idx_map.append(dataset_body_names.index(name))
+
+            # Extract body positions and quaternions (xyzw format in dataset)
+            raw_body_pos_w = np.array(
+                data["body_links_pos_w"][:, body_idx_map, :], dtype=np.float64
+            )
+            raw_body_quat_xyzw = np.array(
+                data["body_links_quat_w"][:, body_idx_map, :], dtype=np.float64
+            )
+
+            # Resample body positions: shape (T_new, N_bodies, 3)
+            n_bodies = len(body_links_names)
+            resampled_body_pos_w = np.zeros((T_new, n_bodies, 3), dtype=np.float64)
+            for b in range(n_bodies):
+                resampled_body_pos_w[:, b, :] = self._resample_data_Rn(
+                    raw_body_pos_w[:, b, :], t_orig, t_new
+                )
+
+            # Resample body quaternions via SLERP: shape (T_new, N_bodies, 4) as Rotation
+            resampled_body_rotations = []  # list of Rotation objects, one per body
+            for b in range(n_bodies):
+                rot_b = self._resample_data_SO3(
+                    raw_body_quat_xyzw[:, b, :], t_orig, t_new
+                )
+                resampled_body_rotations.append(rot_b)
+
+            # Find anchor body index
+            anchor_idx = body_links_names.index(anchor_body_name)
+
+            # Get anchor rotation matrices: shape (T_new, 3, 3)
+            anchor_rot_matrices = resampled_body_rotations[anchor_idx].as_matrix()
+            # Anchor positions: shape (T_new, 3)
+            anchor_pos = resampled_body_pos_w[:, anchor_idx, :]
+
+            # Compute relative transforms for non-anchor bodies
+            non_anchor_indices = [i for i in range(n_bodies) if i != anchor_idx]
+            n_non_anchor = len(non_anchor_indices)
+
+            # body_pos_b: relative position in anchor frame, shape (T_new, n_non_anchor, 3)
+            rel_pos_b = np.zeros((T_new, n_non_anchor, 3), dtype=np.float64)
+            # body_ori_b: relative orientation as 6D (first 2 cols of rot matrix),
+            # shape (T_new, n_non_anchor, 6)
+            rel_ori_b = np.zeros((T_new, n_non_anchor, 6), dtype=np.float64)
+
+            for i, body_idx in enumerate(non_anchor_indices):
+                body_pos = resampled_body_pos_w[:, body_idx, :]
+                body_rot_matrices = resampled_body_rotations[body_idx].as_matrix()
+
+                for t in range(T_new):
+                    R_anchor_T = anchor_rot_matrices[t].T
+                    # Relative position: R_anchor^T @ (p_body - p_anchor)
+                    rel_pos_b[t, i, :] = R_anchor_T @ (body_pos[t] - anchor_pos[t])
+                    # Relative orientation: R_anchor^T @ R_body
+                    R_rel = R_anchor_T @ body_rot_matrices[t]
+                    # 6D representation: first 2 columns of rotation matrix
+                    rel_ori_b[t, i, :3] = R_rel[:, 0]
+                    rel_ori_b[t, i, 3:] = R_rel[:, 1]
+
+            # Compute velocities via numerical differentiation
+            rel_pos_b_flat = rel_pos_b.reshape(T_new, n_non_anchor * 3)
+            rel_ori_b_flat = rel_ori_b.reshape(T_new, n_non_anchor * 6)
+
+            body_pos_b = rel_pos_b_flat
+            body_ori_b = rel_ori_b_flat
+            body_lin_vel_b = self._compute_raw_derivative(rel_pos_b_flat, simulation_dt)
+            body_ang_vel_b = self._compute_raw_derivative(rel_ori_b_flat, simulation_dt)
+        # ─────────────────────────────────────────────────
+
         return MotionData(
             joint_positions=resampled_joint_positions,
             joint_velocities=resampled_joint_velocities,
@@ -535,6 +720,10 @@ def load_data(
             base_ang_velocities_local=resampled_base_ang_vel_local,
             base_quat=resampled_base_orientations,
             device=self.device,
+            body_pos_b=body_pos_b,
+            body_ori_b=body_ori_b,
+            body_lin_vel_b=body_lin_vel_b,
+            body_ang_vel_b=body_ang_vel_b,
         )
 
     def feed_forward_generator(
diff --git a/tests/test_body_keyframes.py b/tests/test_body_keyframes.py
new file mode 100644
index 0000000..3410b1e
--- /dev/null
+++ b/tests/test_body_keyframes.py
@@ -0,0 +1,562 @@
+# SPDX-FileCopyrightText: Generative Bionics S.R.L.
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests for body keyframe observations in AMPLoader.
+
+Run with:
+    python -m pytest tests/test_body_keyframes.py -v
+"""
+
+import torch
+import pytest
+import numpy as np
+import tempfile
+from pathlib import Path
+from scipy.spatial.transform import Rotation
+
+from amp_rsl_rl.utils.motion_loader import AMPLoader, MotionData, VelocityRepresentation
+
+
+# ── Helpers ──────────────────────────────────────────────────────────────────
+
+
+def _create_synthetic_dataset(
+    tmpdir: Path,
+    name: str = "test_motion",
+    n_frames: int = 100,
+    fps: float = 50.0,
+    n_joints: int = 6,
+    joint_names: list = None,
+    include_body_links: bool = True,
+    body_links_list: list = None,
+) -> Path:
+    """Create a synthetic .npy motion dataset for testing."""
+    if joint_names is None:
+        joint_names = [f"joint_{i}" for i in range(n_joints)]
+    if body_links_list is None:
+        body_links_list = ["root_link", "left_knee", "right_knee", "left_foot", "right_foot"]
+
+    n_bodies = len(body_links_list)
+
+    # Generate smooth trajectories
+    t = np.linspace(0, 2 * np.pi, n_frames)
+
+    # Joint positions: sinusoidal motion
+    joint_positions = np.zeros((n_frames, n_joints))
+    for i in range(n_joints):
+        joint_positions[:, i] = 0.5 * np.sin(t + i * np.pi / n_joints)
+
+    # Root position: walking forward
+    root_position = np.zeros((n_frames, 3))
+    root_position[:, 0] = np.linspace(0, 2, n_frames)  # x: forward
+    root_position[:, 2] = 0.9  # z: height
+
+    # Root quaternion: slight rotation (xyzw format)
+    root_quaternion = np.zeros((n_frames, 4))
+    for i in range(n_frames):
+        r = Rotation.from_euler("z", t[i] * 0.1)
+        root_quaternion[i] = r.as_quat()  # xyzw
+
+    data = {
+        "joints_list": joint_names,
+        "joint_positions": joint_positions,
+        "root_position": root_position,
+        "root_quaternion": root_quaternion,
+        "fps": np.array([fps]),
+    }
+
+    if include_body_links:
+        # Generate body link positions/quaternions in world frame
+        body_links_pos_w = np.zeros((n_frames, n_bodies, 3))
+        body_links_quat_w = np.zeros((n_frames, n_bodies, 4))
+
+        for b in range(n_bodies):
+            # Each body moves relative to root with some offset
+            offset = np.array([0.0, 0.1 * (b - n_bodies // 2), -0.1 * b])
+            for i in range(n_frames):
+                R_root = Rotation.from_quat(root_quaternion[i])
+                body_links_pos_w[i, b, :] = root_position[i] + R_root.apply(offset)
+                # Body orientation: root orientation + small relative rotation
+                r_rel = Rotation.from_euler("y", 0.1 * np.sin(t[i] + b))
+                body_links_quat_w[i, b, :] = (R_root * r_rel).as_quat()  # xyzw
+
+        data["body_links_list"] = body_links_list
+        data["body_links_pos_w"] = body_links_pos_w
+        data["body_links_quat_w"] = body_links_quat_w
+
+    filepath = tmpdir / f"{name}.npy"
+    np.save(str(filepath), data)
+    return filepath
+
+
+# ── Fixtures ─────────────────────────────────────────────────────────────────
+
+
+@pytest.fixture
+def tmpdir():
+    with tempfile.TemporaryDirectory() as d:
+        yield Path(d)
+
+
+@pytest.fixture
+def dataset_with_bodies(tmpdir):
+    """Create a dataset with body keyframes."""
+    _create_synthetic_dataset(
+        tmpdir,
+        name="motion_with_bodies",
+        include_body_links=True,
+    )
+    return tmpdir
+
+
+@pytest.fixture
+def dataset_without_bodies(tmpdir):
+    """Create a dataset without body keyframes."""
+    _create_synthetic_dataset(
+        tmpdir,
+        name="motion_no_bodies",
+        include_body_links=False,
+    )
+    return tmpdir
+
+
+# ── Tests: Backward Compatibility ───────────────────────────────────────────
+
+
+class TestBackwardCompatibility:
+    """Verify that existing behavior is unchanged when amp_obs_components is None."""
+
+    def test_default_obs_components(self, dataset_with_bodies):
+        """AMPLoader without amp_obs_components should behave as before."""
+        loader = AMPLoader(
+            device="cpu",
+            dataset_path_root=dataset_with_bodies,
+            datasets={"motion_with_bodies": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+        )
+        # Default obs: joint_pos (6) + joint_vel (6) + base_lin_vel (3) + base_ang_vel (3) = 18
+        assert loader.all_obs.shape[1] == 18
+
+    def test_explicit_default_components(self, dataset_with_bodies):
+        """Explicitly passing default components should give same result."""
+        loader_default = AMPLoader(
+            device="cpu",
+            dataset_path_root=dataset_with_bodies,
+            datasets={"motion_with_bodies": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+        )
+        loader_explicit = AMPLoader(
+            device="cpu",
+            dataset_path_root=dataset_with_bodies,
+            datasets={"motion_with_bodies": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+            amp_obs_components=["joint_pos", "joint_vel", "base_lin_vel", "base_ang_vel"],
+        )
+        assert torch.allclose(loader_default.all_obs, loader_explicit.all_obs)
+        assert torch.allclose(loader_default.all_next_obs, loader_explicit.all_next_obs)
+
+    def test_dataset_without_bodies_default(self, dataset_without_bodies):
+        """Dataset without body links should work fine with default obs."""
+        loader = AMPLoader(
+            device="cpu",
+            dataset_path_root=dataset_without_bodies,
+            datasets={"motion_no_bodies": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+        )
+        assert loader.all_obs.shape[1] == 18
+
+    def test_reset_states_unchanged(self, dataset_with_bodies):
+        """Reset states should not be affected by body keyframes."""
+        loader = AMPLoader(
+            device="cpu",
+            dataset_path_root=dataset_with_bodies,
+            datasets={"motion_with_bodies": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+            amp_obs_components=[
+                "joint_pos", "joint_vel", "base_lin_vel", "base_ang_vel",
+                "body_pos_b",
+            ],
+            body_links_names=["root_link", "left_knee", "right_knee", "left_foot", "right_foot"],
+            anchor_body_name="root_link",
+        )
+        # Reset state: quat(4) + joint_pos(6) + joint_vel(6) + base_lin_vel(3) + base_ang_vel(3) = 22
+        assert loader.all_states.shape[1] == 22
+
+
+# ── Tests: Body Keyframe Observations ───────────────────────────────────────
+
+
+class TestBodyKeyframeObs:
+    """Verify body keyframe observation loading and dimensions."""
+
+    def test_body_pos_b_dimensions(self, dataset_with_bodies):
+        """body_pos_b should have (N_bodies-1)*3 dimensions (anchor excluded)."""
+        loader = AMPLoader(
+            device="cpu",
+            dataset_path_root=dataset_with_bodies,
+            datasets={"motion_with_bodies": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+            amp_obs_components=["body_pos_b"],
+            body_links_names=["root_link", "left_knee", "right_knee", "left_foot", "right_foot"],
+            anchor_body_name="root_link",
+        )
+        # 5 bodies - 1 anchor = 4 non-anchor bodies * 3 = 12
+        assert loader.all_obs.shape[1] == 12
+
+    def test_body_ori_b_dimensions(self, dataset_with_bodies):
+        """body_ori_b should have (N_bodies-1)*6 dimensions (6D rotation repr)."""
+        loader = AMPLoader(
+            device="cpu",
+            dataset_path_root=dataset_with_bodies,
+            datasets={"motion_with_bodies": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+            amp_obs_components=["body_ori_b"],
+            body_links_names=["root_link", "left_knee", "right_knee", "left_foot", "right_foot"],
+            anchor_body_name="root_link",
+        )
+        # 4 non-anchor bodies * 6 = 24
+        assert loader.all_obs.shape[1] == 24
+
+    def test_body_lin_vel_b_dimensions(self, dataset_with_bodies):
+        """body_lin_vel_b should have (N_bodies-1)*3 dimensions."""
+        loader = AMPLoader(
+            device="cpu",
+            dataset_path_root=dataset_with_bodies,
+            datasets={"motion_with_bodies": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+            amp_obs_components=["body_lin_vel_b"],
+            body_links_names=["root_link", "left_knee", "right_knee", "left_foot", "right_foot"],
+            anchor_body_name="root_link",
+        )
+        # 4 non-anchor bodies * 3 = 12
+        assert loader.all_obs.shape[1] == 12
+
+    def test_body_ang_vel_b_dimensions(self, dataset_with_bodies):
+        """body_ang_vel_b should have (N_bodies-1)*6 dimensions (derivative of 6D repr)."""
+        loader = AMPLoader(
+            device="cpu",
+            dataset_path_root=dataset_with_bodies,
+            datasets={"motion_with_bodies": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+            amp_obs_components=["body_ang_vel_b"],
+            body_links_names=["root_link", "left_knee", "right_knee", "left_foot", "right_foot"],
+            anchor_body_name="root_link",
+        )
+        # body_ang_vel_b is derivative of body_ori_b (6D repr), so 4*6 = 24
+        assert loader.all_obs.shape[1] == 24
+
+    def test_combined_components(self, dataset_with_bodies):
+        """All components together should sum dimensions correctly."""
+        loader = AMPLoader(
+            device="cpu",
+            dataset_path_root=dataset_with_bodies,
+            datasets={"motion_with_bodies": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+            amp_obs_components=[
+                "joint_pos", "joint_vel", "base_lin_vel", "base_ang_vel",
+                "body_pos_b", "body_ori_b",
+            ],
+            body_links_names=["root_link", "left_knee", "right_knee", "left_foot", "right_foot"],
+            anchor_body_name="root_link",
+        )
+        # joint_pos(6) + joint_vel(6) + base_lin_vel(3) + base_ang_vel(3)
+        # + body_pos_b(4*3=12) + body_ori_b(4*6=24) = 54
+        assert loader.all_obs.shape[1] == 54
+
+    def test_subset_of_bodies(self, dataset_with_bodies):
+        """Only a subset of body links can be specified."""
+        loader = AMPLoader(
+            device="cpu",
+            dataset_path_root=dataset_with_bodies,
+            datasets={"motion_with_bodies": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+            amp_obs_components=["body_pos_b"],
+            body_links_names=["root_link", "left_knee", "right_knee"],
+            anchor_body_name="root_link",
+        )
+        # 3 bodies - 1 anchor = 2 non-anchor * 3 = 6
+        assert loader.all_obs.shape[1] == 6
+
+    def test_obs_and_next_obs_shapes_match(self, dataset_with_bodies):
+        """obs and next_obs should have same shape."""
+        loader = AMPLoader(
+            device="cpu",
+            dataset_path_root=dataset_with_bodies,
+            datasets={"motion_with_bodies": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+            amp_obs_components=["joint_pos", "body_pos_b", "body_ori_b"],
+            body_links_names=["root_link", "left_knee", "right_knee"],
+            anchor_body_name="root_link",
+        )
+        assert loader.all_obs.shape == loader.all_next_obs.shape
+
+    def test_feed_forward_generator(self, dataset_with_bodies):
+        """feed_forward_generator should yield batches of correct shape."""
+        loader = AMPLoader(
+            device="cpu",
+            dataset_path_root=dataset_with_bodies,
+            datasets={"motion_with_bodies": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+            amp_obs_components=["joint_pos", "body_pos_b"],
+            body_links_names=["root_link", "left_knee", "right_knee"],
+            anchor_body_name="root_link",
+        )
+        # joint_pos(6) + body_pos_b(2*3=6) = 12
+        batch_size = 32
+        for obs, next_obs in loader.feed_forward_generator(2, batch_size):
+            assert obs.shape == (batch_size, 12)
+            assert next_obs.shape == (batch_size, 12)
+
+
+# ── Tests: Validation & Error Handling ───────────────────────────────────────
+
+
+class TestValidation:
+    """Test error cases and input validation."""
+
+    def test_missing_body_links_names(self, dataset_with_bodies):
+        """Should raise ValueError when body components requested without body_links_names."""
+        with pytest.raises(ValueError, match="body_links_names must be provided"):
+            AMPLoader(
+                device="cpu",
+                dataset_path_root=dataset_with_bodies,
+                datasets={"motion_with_bodies": 1.0},
+                simulation_dt=0.02,
+                slow_down_factor=1,
+                amp_obs_components=["body_pos_b"],
+            )
+
+    def test_missing_anchor_body_name(self, dataset_with_bodies):
+        """Should raise ValueError when body components requested without anchor_body_name."""
+        with pytest.raises(ValueError, match="anchor_body_name must be provided"):
+            AMPLoader(
+                device="cpu",
+                dataset_path_root=dataset_with_bodies,
+                datasets={"motion_with_bodies": 1.0},
+                simulation_dt=0.02,
+                slow_down_factor=1,
+                amp_obs_components=["body_pos_b"],
+                body_links_names=["root_link", "left_knee"],
+            )
+
+    def test_anchor_not_in_body_links(self, dataset_with_bodies):
+        """Should raise ValueError if anchor_body_name not in body_links_names."""
+        with pytest.raises(ValueError, match="must be present in body_links_names"):
+            AMPLoader(
+                device="cpu",
+                dataset_path_root=dataset_with_bodies,
+                datasets={"motion_with_bodies": 1.0},
+                simulation_dt=0.02,
+                slow_down_factor=1,
+                amp_obs_components=["body_pos_b"],
+                body_links_names=["left_knee", "right_knee"],
+                anchor_body_name="root_link",
+            )
+
+    def test_dataset_missing_body_links(self, dataset_without_bodies):
+        """Should raise KeyError when body components requested but dataset lacks body data."""
+        with pytest.raises(KeyError, match="does not contain 'body_links_list'"):
+            AMPLoader(
+                device="cpu",
+                dataset_path_root=dataset_without_bodies,
+                datasets={"motion_no_bodies": 1.0},
+                simulation_dt=0.02,
+                slow_down_factor=1,
+                amp_obs_components=["body_pos_b"],
+                body_links_names=["root_link", "left_knee"],
+                anchor_body_name="root_link",
+            )
+
+    def test_body_name_not_in_dataset(self, dataset_with_bodies):
+        """Should raise ValueError when requested body name not found in dataset."""
+        with pytest.raises(ValueError, match="not found in dataset"):
+            AMPLoader(
+                device="cpu",
+                dataset_path_root=dataset_with_bodies,
+                datasets={"motion_with_bodies": 1.0},
+                simulation_dt=0.02,
+                slow_down_factor=1,
+                amp_obs_components=["body_pos_b"],
+                body_links_names=["root_link", "nonexistent_body"],
+                anchor_body_name="root_link",
+            )
+
+    def test_invalid_obs_component(self, dataset_with_bodies):
+        """Should raise ValueError when an invalid obs component is specified."""
+        with pytest.raises(ValueError, match="not available"):
+            AMPLoader(
+                device="cpu",
+                dataset_path_root=dataset_with_bodies,
+                datasets={"motion_with_bodies": 1.0},
+                simulation_dt=0.02,
+                slow_down_factor=1,
+                amp_obs_components=["nonexistent_component"],
+            )
+
+
+# ── Tests: Numerical Correctness ────────────────────────────────────────────
+
+
+class TestNumericalCorrectness:
+    """Verify the numerical correctness of relative transforms."""
+
+    def test_body_pos_b_is_relative_to_anchor(self, tmpdir):
+        """When a body is at the same position as anchor, body_pos_b should be zero."""
+        # Create dataset where body 1 is at the same position as root_link
+        n_frames = 50
+        fps = 50.0
+        t = np.linspace(0, 2 * np.pi, n_frames)
+
+        root_position = np.zeros((n_frames, 3))
+        root_position[:, 2] = 0.9
+        root_quaternion = np.tile([0, 0, 0, 1], (n_frames, 1)).astype(np.float64)
+
+        body_links_pos_w = np.zeros((n_frames, 2, 3))
+        body_links_quat_w = np.zeros((n_frames, 2, 4))
+
+        # Body 0 (anchor): same as root
+        body_links_pos_w[:, 0, :] = root_position
+        body_links_quat_w[:, 0, :] = root_quaternion
+
+        # Body 1: same as anchor (so relative position should be ~0)
+        body_links_pos_w[:, 1, :] = root_position
+        body_links_quat_w[:, 1, :] = root_quaternion
+
+        data = {
+            "joints_list": ["j0", "j1"],
+            "joint_positions": np.zeros((n_frames, 2)),
+            "root_position": root_position,
+            "root_quaternion": root_quaternion,
+            "fps": np.array([fps]),
+            "body_links_list": ["anchor", "same_body"],
+            "body_links_pos_w": body_links_pos_w,
+            "body_links_quat_w": body_links_quat_w,
+        }
+        np.save(str(tmpdir / "test.npy"), data)
+
+        loader = AMPLoader(
+            device="cpu",
+            dataset_path_root=tmpdir,
+            datasets={"test": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+            amp_obs_components=["body_pos_b"],
+            body_links_names=["anchor", "same_body"],
+            anchor_body_name="anchor",
+        )
+
+        # body_pos_b for a body coincident with anchor should be ~0
+        assert torch.allclose(loader.all_obs, torch.zeros_like(loader.all_obs), atol=1e-5)
+
+    def test_body_ori_b_identity_when_same_orientation(self, tmpdir):
+        """When body has same orientation as anchor, the 6D repr should be [1,0,0, 0,1,0]."""
+        n_frames = 50
+        fps = 50.0
+
+        root_position = np.zeros((n_frames, 3))
+        root_position[:, 2] = 0.9
+        root_quaternion = np.tile([0, 0, 0, 1], (n_frames, 1)).astype(np.float64)
+
+        body_links_pos_w = np.zeros((n_frames, 2, 3))
+        body_links_quat_w = np.zeros((n_frames, 2, 4))
+
+        # Both bodies have identity orientation
+        body_links_pos_w[:, 0, :] = root_position
+        body_links_quat_w[:, 0, :] = root_quaternion
+        body_links_pos_w[:, 1, :] = root_position + np.array([1, 0, 0])
+        body_links_quat_w[:, 1, :] = root_quaternion
+
+        data = {
+            "joints_list": ["j0"],
+            "joint_positions": np.zeros((n_frames, 1)),
+            "root_position": root_position,
+            "root_quaternion": root_quaternion,
+            "fps": np.array([fps]),
+            "body_links_list": ["anchor", "body1"],
+            "body_links_pos_w": body_links_pos_w,
+            "body_links_quat_w": body_links_quat_w,
+        }
+        np.save(str(tmpdir / "test.npy"), data)
+
+        loader = AMPLoader(
+            device="cpu",
+            dataset_path_root=tmpdir,
+            datasets={"test": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+            amp_obs_components=["body_ori_b"],
+            body_links_names=["anchor", "body1"],
+            anchor_body_name="anchor",
+        )
+
+        # Identity rotation → first 2 cols of I = [1,0,0, 0,1,0]
+        expected = torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float32)
+        for i in range(loader.all_obs.shape[0]):
+            assert torch.allclose(loader.all_obs[i], expected, atol=1e-5), (
+                f"Frame {i}: {loader.all_obs[i]} != {expected}"
+            )
+
+    def test_non_anchor_body_at_known_offset(self, tmpdir):
+        """Body at a known offset from anchor should produce correct relative position."""
+        n_frames = 50
+        fps = 50.0
+
+        root_position = np.zeros((n_frames, 3))
+        root_position[:, 2] = 0.9
+        # Identity orientation for simplicity
+        root_quaternion = np.tile([0, 0, 0, 1], (n_frames, 1)).astype(np.float64)
+
+        body_links_pos_w = np.zeros((n_frames, 2, 3))
+        body_links_quat_w = np.zeros((n_frames, 2, 4))
+
+        # Anchor at root
+        body_links_pos_w[:, 0, :] = root_position
+        body_links_quat_w[:, 0, :] = root_quaternion
+
+        # Body1 at constant offset [0.5, 0.3, -0.2] from anchor in world frame
+        # With identity anchor rotation, relative = world offset
+        offset = np.array([0.5, 0.3, -0.2])
+        body_links_pos_w[:, 1, :] = root_position + offset
+        body_links_quat_w[:, 1, :] = root_quaternion
+
+        data = {
+            "joints_list": ["j0"],
+            "joint_positions": np.zeros((n_frames, 1)),
+            "root_position": root_position,
+            "root_quaternion": root_quaternion,
+            "fps": np.array([fps]),
+            "body_links_list": ["anchor", "body1"],
+            "body_links_pos_w": body_links_pos_w,
+            "body_links_quat_w": body_links_quat_w,
+        }
+        np.save(str(tmpdir / "test.npy"), data)
+
+        loader = AMPLoader(
+            device="cpu",
+            dataset_path_root=tmpdir,
+            datasets={"test": 1.0},
+            simulation_dt=0.02,
+            slow_down_factor=1,
+            amp_obs_components=["body_pos_b"],
+            body_links_names=["anchor", "body1"],
+            anchor_body_name="anchor",
+        )
+
+        expected = torch.tensor(offset, dtype=torch.float32)
+        for i in range(loader.all_obs.shape[0]):
+            assert torch.allclose(loader.all_obs[i], expected, atol=1e-4), (
+                f"Frame {i}: {loader.all_obs[i]} != {expected}"
+            )