简介

旋转表示互转算子，支持 axis_angle / euler_angles / quaternion / rotation_6d / matrix 五种表示两两互转，全部经旋转矩阵中转。移植 GR00T-N1.5 gr00t/data/transform/state_action.py 的 RotationTransform（原版来自 diffusion_policy 的 rotation_transformer），不依赖 pytorch3d。

功能描述

• 五种表示的维度固定：axis_angle 3、euler_angles 3、quaternion 4、rotation_6d 6、matrix 9（3×3 行优先展平）
• 一律经 matrix 中转，forward/inverse 对称：库函数 convert_rotation(values, from_rep, to_rep) 把两个表示名对调就是逆变换。单射域内数值严格可逆（轴角范数 < π、规范 6d、matrix）；euler_angles 同一旋转有多组解、quaternion 有 ± 双覆盖，往返只保证旋转等价，不保证逐位相同
• 不引 pytorch3d（运行机未装）：欧拉角 / 轴角 / 四元数 ↔ matrix 走 scipy.spatial.transform.Rotation（pytorch3d 的同串大写 intrinsic 欧拉约定与 scipy 一致），rotation_6d ↔ matrix 自实现——6d 即旋转矩阵前两行展平，反向按 b1=normalize(a1)、b2=Gram-Schmidt(a2, b1)、b3=b1×b2
• 四元数沿用 GR00T / pytorch3d 的 wxyz（scalar-first）布局，进出 scipy 时重排为 xyzw；输入四元数先归一化
• 欧拉约定解析：euler_angles（默认 XYZ）、euler_angles_XYZ（大写轴串原样）、euler_angles_zyx（小写轴名转大写）、GR00T 缩写 euler_angles_rpy（r→X、p→Y、y→Z）；其它写法报错
• joint_dims 按给定顺序取出选中列组成旋转段，段长须是源表示维度的整数倍，因此一段里可以有多组旋转（如 36 维 = 6 组 rotation_6d），输出恒为展平的 组数 × 目标维度
• keep_other_dims=True 时未选中列按原相对顺序保留，转换段整体插在最左选中列的位置
• rotation_6d → matrix 的归一化用 eps 兜住零向量，不会除零
• from_rep 与 to_rep 相同（含欧拉约定相同）在构造时报错
• 行级容错：空行 / null / 段长不是源表示维度整数倍 / 下标越界，该行返回空列表，不打断整批

算子参数

输入

输出

参数

注意事项

• 输入段长必须是源表示维度的整数倍，否则该行返回空列表（例如 from_rep="axis_angle" 时 2 维的行会被判为坏行）。
• rotation_6d 的转换隐含正交化：任意 6 维向量都会被 Gram-Schmidt 投影成合法旋转，第二次转换起是幂等的，但第一次会改值。上游若不是真旋转数据（例如把关节角当 6d 用），结果没有物理含义。
• matrix 是 3×3 按行优先展平的 9 维，不要按列优先解释。
• keep_other_dims=True 的输出不是「按原下标替换」：转换段整体落在最左选中列的位置，未选中列只保持彼此的相对顺序。
• 输出列长度随 from_rep / to_rep 变化，下游 schema 要同步（如 3 维轴角 → 6 维 6d）。

调用示例

1from __future__ import annotations
2
3import os
4
5import daft
6import pyarrow.parquet as pq
7from daft import col
8
9from daft.aihc.common.udf import aihc_udf
10from daft.aihc.functions.embodied.lerobot.rotation_rep_convert import RotationRepConvert
11
12ROOT = "/path/to/lerobot_dataset"
13
14if __name__ == "__main__":
15    if os.getenv("DAFT_RUNNER", "native") == "ray":
16        import ray
17        ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
18        daft.set_runner_ray()
19    daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
20
21table = pq.read_table(
22        f"{ROOT}/data/chunk-000/episode_000000.parquet", columns=["observation.state"]
23    )
24    ds = daft.from_pydict({"state": table.column("observation.state").to_pylist()})
25    ds = ds.with_column(
26        "result",
27        aihc_udf(
28            RotationRepConvert,
29            construct_args={
30                "from_rep": "rotation_6d",
31                "to_rep": "matrix",
32                "joint_dims": list(range(6)),
33                "keep_other_dims": True,
34                "eps": 1e-8,
35            },
36            num_cpus=1,
37            concurrency=1,
38            batch_size=256,
39        )(col("state")),
40    )
41    ds.show()