滴啦嘟啦哒头像
【具身笔记】【Piper+PI0.5】二、基于自有数据集实现pi05微调与评估封面图

【具身笔记】【Piper+PI0.5】二、基于自有数据集实现pi05微调与评估

概述与QA

本系列为我在vla上的第二个实践:基于piper机械臂完成数据采集--格式转换--pi05微调 --评估--panda支线对比。本文主要介绍微调的过程。

下面列出几个QA来引入:

Q:什么叫微调?

A:简单来说就是让已经预训练过的模型在自己的数据集上再训练,使其能学会自己数据集上的东西


1. 第一次微调:只给动作专家进行lora微调

什么是lora,有很多文章或视频进行了讲解,如:

【LoRA微调】从原理到调参,7 个问题彻底理解LoRA,不懂线性代数也没问题_大模型微调_低秩适配_哔哩哔哩_bilibili

我对于lora的理解很简单,就是只关注每次更新权重的增量,降这个增量大矩阵拆成两个小矩阵,每次只更新这两个个小矩阵的参数,从而实现训练参数的大幅减少。

1.1. 训练指令

使用如下指令开启只针对动作专家的lora微调:

source /root/miniconda3/etc/profile.d/conda.sh
conda activate lerobot

export HF_HOME=/root/autodl-tmp/huggingface
export HF_HUB_OFFLINE=1

cd /root/autodl-tmp/workspace
mkdir -p outputs

b4_job="b4_piper_lora_19020steps_$(date +%Y%m%d_%H%M%S)"

set -o pipefail

# 不指定 dataset.episodes,使用全部 60 条演示
# time 的计时结果也会保存到日志
{
  time -p lerobot-train \
    --dataset.repo_id=dukiya/piper_black_bowl_b1_v2 \
    --dataset.root=/root/autodl-tmp/workspace/datasets/week5/piper_black_bowl_b1_v2_lerobot \
    --dataset.video_backend=pyav \
    --policy.type=pi05 \
    --policy.pretrained_path=/root/autodl-tmp/huggingface/hub/models--lerobot--pi05_libero_base/snapshots/a217bfd3b14673cf2ce597e69997ab21866438dd \
    --policy.normalization_mapping='{"ACTION":"MEAN_STD","STATE":"MEAN_STD","VISUAL":"IDENTITY"}' \
    --policy.use_relative_actions=false \
    --policy.n_action_steps=10 \
    --policy.empty_cameras=1 \
    --policy.freeze_vision_encoder=true \
    --policy.train_expert_only=true \
    --policy.gradient_checkpointing=true \
    --policy.dtype=bfloat16 \
    --policy.device=cuda \
    --policy.push_to_hub=false \
    --policy.optimizer_lr=2.5e-5 \
    --policy.scheduler_warmup_steps=100 \
    --policy.scheduler_decay_steps=19020 \
    --policy.scheduler_decay_lr=2.5e-6 \
    --peft.method_type=LORA \
    --peft.r=16 \
    --peft.lora_alpha=16 \
    --output_dir="outputs/${b4_job}" \
    --job_name="${b4_job}" \
    --batch_size=4 \
    --num_workers=0 \
    --steps=19020 \
    --log_freq=50 \
    --save_freq=2000 \
    --seed=1000 \
    --wandb.enable=false
} 2>&1 | tee "outputs/${b4_job}.log"

其中:

1. dataset.root=/root/autodl-tmp/workspace/datasets/week5/piper_black_bowl_b1_v2_lerobot表明训练集是我们自己的数据集

2. --policy.freeze_vision_encoder=true表明视觉模块冻结,不参与训练

3. --policy.train_expert_only=true表明只训练动作专家部分

4. --peft.method_type=LORA表明使用LORA方法做高效微调

1.2. 训练结果

如下:

训练了19000步左右,最终的loss为0.2。

1.3. 评估结果

通过如下指令让训练好的模型在自己搭建的50个同分布样本和50个OOD样本上评估,看成功率:

source /root/miniconda3/etc/profile.d/conda.sh
conda activate lerobot
cd /root/autodl-tmp/workspace

export HF_HOME=/root/autodl-tmp/huggingface
export HF_HUB_OFFLINE=1
export PYTHONUNBUFFERED=1
export OMP_NUM_THREADS=8
export OPENBLAS_NUM_THREADS=8
export MUJOCO_GL=egl
export PYOPENGL_PLATFORM=egl
export LIBERO_CONFIG_PATH=/root/autodl-tmp/libero-config
set -o pipefail

CHECKPOINT=/root/autodl-tmp/workspace/outputs/b4_piper_lora_19020steps_20261016_213459/checkpoints/019020/pretrained_model
NAME="piper_action_lora_19020_same50_$(date +%Y%m%d_%H%M%S)"
mkdir -p outputs/week05

python -u my_test/week5/A4_3_piper_lerobot.py \
  --policy.path="$CHECKPOINT" \
  --policy.use_peft=true \
  --policy.device=cuda \
  --policy.dtype=bfloat16 \
  --policy.n_action_steps=10 \
  --policy.use_relative_actions=false \
  --env.type=piper_libero \
  --env.task=piper_black_bowl_test_same_dist \
  --env.task_ids='[0]' \
  --env.control_mode=relative \
  --env.observation_height=256 \
  --env.observation_width=256 \
  --env.episode_length=1000 \
  --env.init_states=true \
  --env.hard_reset=true \
  --env.max_parallel_tasks=1 \
  --eval.batch_size=1 \
  --eval.use_async_envs=false \
  --eval.n_episodes=50 \
  --seed=1000 \
  --output_dir="outputs/week05/$NAME" \
  --job_name="$NAME" \
  2>&1 | tee "outputs/week05/$NAME.log"

最终结果,同分布样本成功率2/50,OOD成功率0/50,效果很烂。

2. 基于第一次做优化

查阅了相关资料:

openpi 0.5 libero lora微调记录 - 知乎 (zhihu.com)

以及在网上发了求助,最后新增了4种训练方式:

B. LLM和动作专家部分使用lora

C. 仅动作专家全量,其余模块冻结

D. 仅LLM用lora,其余部分全量

E. VLM用lora,动作专家全量

2.1. 为什么要拆分成这样

为什么要拆分成这样,不同模块有的全量,有的用lora呢?

第一,关注pi05的模型结构图:

简单理解:

1. pi0.5需要训练的核心部分为VLM(视觉语言模型) + action expert(动作专家)

2. VLM可以分为视觉部分和语言部分(llm)

3. 我的五次训练中,分别对这些部分进行了全量和lora微调尝试,评估后得到了不同的结果

第二,pi05官方推荐的微调方法如下:

其中2对应我的B方法,3对应C,4对应E,5对应A,在1的基础上,我将VLM全参改成了把里面的LLM挑出来单独做lora,其余都用全参,因为我的算力难以支撑全量微调语言大模型。

2.2. 五次训练和评估的结果对比

相关图如下:

2.3. 对结果的思考

在评估结果中,我重点关注了方法D:仅LLM lora,其余部分全量微调。可以看到大量模块的全量微调,让其在同分布样本上有不错的表现,但OOD样本上却一次都没成功,我最开始对此很疑惑。

查阅了相关资料:

[2202.10054] Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution (arxiv.org)

[2405.09673] LoRA Learns Less and Forgets Less (arxiv.org)

以及在网上看了相关帖子,我对此现象做出如下理解:

1. 我的数据集太小,在此上做全量微调,大量更新参数,会导致模型过拟合,并且只学会了在我的轨迹上完成任务,原本预训练时的泛化能力被抹掉了。

2. 适当的lora可以保持记忆,能保持预训练时的泛化能力。

总之就是小数据集上需要一定程度的lora和全量相结合,可以看到方法E正是如此的结合,大参数的VLM使用lora,在视觉和语言能力上保持一定的泛化性,专门负责动作输出的动作专家模块则做全量微调,让其的输出更适配当前的机械臂和场景。


3. 用panda做支线对比

完成piper的完整流程后,我又用panda走了一遍相同的流程,其中在数据采集部分需要重点说一下。

3.1. 采集数据踩坑

第一次采集panda的数据时,由于robosuite已提前集成好了panda的接口,因此我便使用默认的接口进行了数据采集:

import numpy as np
from robosuite.devices import Keyboard
from robosuite.utils.input_utils import input2action

robot = env.robots[0]
controller = robot.controller

# 使用 Panda 原生相对控制
assert controller.use_delta
assert controller.uncoupling

device = Keyboard(
    pos_sensitivity=0.1,
    rot_sensitivity=0.2,
)
device.start_control()

states = []
actions = []
recording = False

while True:
    # 直接取得键盘产生的相对动作
    action, _ = input2action(
        device=device,
        robot=robot,
        active_arm="right",
    )

    if action is None:
        break

    # 首次移动或闭合夹爪后开始录制
    if not recording:
        if np.any(action[:6] != 0) or action[6] > 0:
            recording = True
        else:
            continue

    # 保存动作执行前的状态
    state_before = env.get_sim_state().copy()

    # 不累积目标,直接执行键盘相对动作
    observation, _, _, _ = env.step(action)

    # 开始录制后,每一步都保存,包括前6维为0的动作
    states.append(state_before)
    actions.append(action.copy())

其中关键设置controller.use_delta = True表明每一步的控制都是相对数值,所以我长按键盘让机械臂移动时,大量的数据值都为0。之后,我用此数据做了ACT的训练和评估:

# 1. 环境准备
source /root/miniconda3/etc/profile.d/conda.sh
conda activate lerobot

BASE=/root/panda-runs/pipeline_20261019
RUN="panda_act_v2_40000steps_$(date +%Y%m%d_%H%M%S)"
OUT="$BASE/$RUN"

export HF_HOME=/root/autodl-tmp/huggingface
export HF_HUB_OFFLINE=1
export PYTHONUNBUFFERED=1
export OMP_NUM_THREADS=8
export OPENBLAS_NUM_THREADS=8
set -o pipefail

# 2. ACT 训练:40,000 步
lerobot-train \
  --dataset.repo_id=dukiya/panda_black_bowl_v1 \
  --dataset.root=/root/panda-runs/v2_panda_black_bowl_lerobot \
  --dataset.video_backend=pyav \
  --policy.type=act \
  --policy.device=cuda \
  --policy.push_to_hub=false \
  --policy.chunk_size=20 \
  --policy.n_action_steps=5 \
  --policy.vision_backbone=resnet18 \
  --policy.pretrained_backbone_weights=ResNet18_Weights.IMAGENET1K_V1 \
  --policy.normalization_mapping='{"VISUAL":"MEAN_STD","STATE":"MEAN_STD","ACTION":"MEAN_STD"}' \
  --policy.use_vae=true \
  --policy.kl_weight=10 \
  --policy.optimizer_lr=1e-5 \
  --policy.optimizer_lr_backbone=1e-5 \
  --policy.optimizer_weight_decay=1e-4 \
  --batch_size=8 \
  --num_workers=0 \
  --steps=40000 \
  --log_freq=50 \
  --save_freq=5000 \
  --save_checkpoint=true \
  --env_eval_freq=0 \
  --seed=1000 \
  --wandb.enable=false \
  --output_dir="$OUT" \
  --job_name="$RUN" \
  2>&1 | tee "$BASE/${RUN}_train.log"

什么是ACT?为什么要做ACT?

ACT是一种纯模仿学习,它比pi05轻量得多。在这里做这个流程是为了用ACT来验证这批数据的质量够不够好,能不能教会机械臂做任务,如果数据集本身没有问题,那么模型模仿训练到一定程度,就一定能学会动作。

第一次采集的数据集经过ACT训练,并做了评测,结果为:0/30。

也就是说,模型学了这批数据,然后在这批数据上试着自己推理走一遍,发现压根没有学会,这批数据是有问题的!

进一步排查,复现了模型的推理结果,发现很多帧的动作都是0,机械臂一动不动。此时才注意到是数据集由于记录相对变量的设置原因,在大部分时间里数据都为0,模型学会了通过输出0来降低loss,导致最后训练好推理时也频繁输出0,评估效果不佳。

基于此,我将panda的数据采集方式改成了跟piper一致:

# B2_2_panda_collect.py
# 一次运行采一条:
# /home/dukiya/venvs/libero/bin/python week5/B2_2_panda_collect.py --init-index 0
# 每个训练开局采两条,编号 0~29,共 60 条。

import os

# Ubuntu Wayland 桌面通过 XWayland 显示 OpenCV 窗口
os.environ.setdefault("QT_QPA_PLATFORM", "xcb")

import argparse
import hashlib
import json
import time
import xml.etree.ElementTree as ET
from datetime import datetime
from pathlib import Path

import cv2
import h5py
import numpy as np

from robosuite import macros
from robosuite.devices import Keyboard
from robosuite.utils import transform_utils as T
from robosuite.utils.input_utils import input2action
from libero.libero import benchmark, get_libero_path
from libero.libero.envs.env_wrapper import ControlEnv

CONTROL_FREQ = 20
SEED = 1000
MAX_STEPS = 1000
SETTLE_STEPS = 10
SUCCESS_HOLD_STEPS = 10

# 沿用 Piper 已验证的持续目标控制方式(B1_piper_collect.py 同款参数)
POSITION_GAIN = 0.01
ROTATION_GAIN = 0.5
MAX_POSITION_LEAD = 0.03
MAX_ROTATION_LEAD = 0.2

# robosuite Keyboard 默认 pos_sensitivity=0.1,一次按键只推进 3.75mm;
# Piper 数据的逐帧目标推进中位是 20.58mm,因此这里提高到 0.5(18.75mm/次)。
POS_SENSITIVITY = 0.5
ROT_SENSITIVITY = 0.2
CAMERAS = ["agentview_image", "robot0_eye_in_hand_image"]
# 绝对动作重采版本单独放一个目录,旧 v1(相对动作)60 条保持不动
OUTPUT_DIR = (
    Path(__file__).resolve().parent / "datasets/b2_panda_black_bowl_v2_abs"
)

def display_image(image):
    # 仅转换显示方向和颜色,不修改原始观测
    if macros.IMAGE_CONVENTION == "opengl":
        image = image[::-1]
    # RGB改BGR,让cv2能更好的展示
    return cv2.cvtColor(np.ascontiguousarray(image), cv2.COLOR_RGB2BGR)

def make_command(position, orientation, gripper):
    # 世界坐标目标位置3 + 目标轴角3 + 夹爪指令1。
    command = np.concatenate([
        position,
        T.quat2axisangle(T.mat2quat(orientation)),
        [gripper],
    ])
    assert command.shape == (7,)
    assert np.isfinite(command).all()
    return command

def limit_target_gap(gap, max_lead):
    # 超出允许距离时只截断多余部分,不整步拒绝,避免退回指令被锁死
    distance = np.linalg.norm(gap)
    if distance > max_lead:
        return gap * (max_lead / distance)
    return gap.copy()

def advance_target(target_pos, target_ori, action, controller):
    """键盘增量 -> 持久目标。与 B1_piper_collect.py / B2_1 修复版一致。"""
    candidate_pos = target_pos + action[:3] * POSITION_GAIN
    position_gap = candidate_pos - controller.ee_pos
    if np.linalg.norm(position_gap) <= MAX_POSITION_LEAD:
        target_pos = candidate_pos
    else:
        # 只截断多余部分,保留持续目标,不整步丢弃
        target_pos = controller.ee_pos + limit_target_gap(
            position_gap, MAX_POSITION_LEAD,
        )

    candidate_ori = (
        T.quat2mat(T.axisangle2quat(action[3:6] * ROTATION_GAIN)) @ target_ori
    )
    rotation_quat = T.mat2quat(candidate_ori @ controller.ee_ori_mat.T)
    # q和-q代表同一旋转;取w非负,按不超过pi的短旋转计算误差
    if rotation_quat[3] < 0:
        rotation_quat = -rotation_quat
    rotation_gap = T.quat2axisangle(rotation_quat)
    if np.linalg.norm(rotation_gap) <= MAX_ROTATION_LEAD:
        target_ori = candidate_ori
    else:
        target_ori = T.quat2mat(
            T.axisangle2quat(limit_target_gap(rotation_gap, MAX_ROTATION_LEAD))
        ) @ controller.ee_ori_mat

    assert np.isfinite(target_pos).all()
    assert np.isfinite(target_ori).all()
    return target_pos, target_ori

class CollectKeyboard(Keyboard):
    def __init__(self, **kwargs):
        self.retry = False
        super().__init__(**kwargs)

    def on_release(self, key):
        # 抓取失败需要重试时,按 N 放弃本条,从原开局重新采。
        if getattr(key, "char", None) == "n":
            self.retry = True
            return
        super().on_release(key)

def progress(source_hash):
    counts = [0] * 30
    lengths = []
    for path in sorted(OUTPUT_DIR.glob("demo_train_*.hdf5")):
        with h5py.File(path, "r") as file:
            info = json.loads(file["data"].attrs["env_info"])
            demo = file["data/demo_0"]
            assert info["init_sha256"] == source_hash, path
            assert info["robots"] == ["Panda"], path
            assert bool(demo.attrs["success"]), path
            index = int(info["init_index"])
            assert 0 <= index < 30, path
            counts[index] += 1
            lengths.append(len(demo["actions"]))
    print("各训练开局成功条数(0~29):", counts)
    print("累计成功条数:", sum(counts), "/ 60")
    print("成功示范平均步数:", float(np.mean(lengths)) if lengths else None)
    return counts

def main():
    global OUTPUT_DIR

    parser = argparse.ArgumentParser()
    parser.add_argument("--init-index", type=int, required=True, choices=range(30))
    parser.add_argument(
        "--pos-sensitivity",
        type=float,
        default=POS_SENSITIVITY,
        help="键盘平移灵敏度;一次按键推进 = 0.05*sens*75*POSITION_GAIN 米",
    )
    parser.add_argument(
        "--output-dir",
        default=str(OUTPUT_DIR),
        help="新采集目录;默认 datasets/b2_panda_black_bowl_v2_abs",
    )
    args = parser.parse_args()

    OUTPUT_DIR = Path(args.output_dir).resolve()
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    run_id = datetime.now().astimezone().strftime("%Y%m%d_%H%M%S_%f")
    env = device = None
    states, actions, stationary_flags = [], [], []
    record_start = None
    success = False
    success_count = switches = 0
    previous_gripper = -1.0
    last_command = None
    pressed_frames = 0
    advance_total = 0.0
    lead_now = 0.0
    stop_reason = "用户按 Q 放弃"
    initial_state = final_state = None
    env_info = model_xml = None

    try:
        suite = benchmark.get_benchmark_dict()["libero_spatial"]()
        task = suite.get_task(0)
        print("任务:", task.language)
        assert task.language == (
            "pick up the black bowl between the plate "
            "and the ramekin and place it on the plate"
        )

        official_states = np.ascontiguousarray(
            suite.get_task_init_states(0), dtype=np.float64,
        )
        assert official_states.shape == (50, 92), official_states.shape
        source_hash = hashlib.sha256(official_states.tobytes()).hexdigest()

        # 固定保存开局,后续评估复用。
        # 官方 0~29 用于训练,30~49 留作同分布测试。
        # OOD 尚未生成,不能把随机留出的官方开局直接称为 OOD。
        pack_path = OUTPUT_DIR / "init_states.npz"
        if not pack_path.exists():
            with pack_path.open("xb") as file:
                np.savez(
                    file,
                    train=official_states[:30],
                    test_same_dist=official_states[30:],
                    train_official_indices=np.arange(30),
                    test_same_dist_official_indices=np.arange(30, 50),
                )
        with np.load(pack_path, allow_pickle=False) as pack:
            np.testing.assert_array_equal(pack["train"], official_states[:30])
            np.testing.assert_array_equal(pack["test_same_dist"], official_states[30:])
            selected = pack["train"][args.init_index].copy()

        counts = progress(source_hash)
        if counts[args.init_index] >= 2:
            print("这个开局已有两条成功演示,请换编号。")
            return

        bddl_path = (
            Path(get_libero_path("bddl_files"))
            / task.problem_folder
            / task.bddl_file
        )
        env = ControlEnv(
            bddl_file_name=str(bddl_path),
            robots=["Panda"],
            controller="OSC_POSE",
            initialization_noise=None,
            control_freq=CONTROL_FREQ,
            has_renderer=True,
            has_offscreen_renderer=True,
            render_camera="agentview",
            use_camera_obs=True,
            camera_names=["agentview", "robot0_eye_in_hand"],
            camera_widths=256,
            camera_heights=256,
            camera_depths=False,
            ignore_done=True,
        )
        env.seed(SEED)
        env.reset()

        # 此前当前配置 seed=1000 的原始 reset 开局实测互穿。
        # 这里恢复官方开局,不用原始 reset 的随机摆放。
        observation = env.set_init_state(selected)
        np.testing.assert_allclose(
            env.get_sim_state(), selected, rtol=0, atol=1e-8,
        )

        robot = env.robots[0]
        controller = robot.controller
        controller.update(force=True)
        assert robot.action_dim == 7

        # 与 Piper 采集一致:绝对目标 + 持久 target(松开键后守住同一目标)。
        # 相对控制会把目标每帧贴回当前位姿,导致 66% 帧动作为零、累积意图被抹掉。
        controller.uncoupling = False
        controller.use_delta = False
        assert not controller.use_delta
        assert not controller.uncoupling

        target_pos = controller.ee_pos.copy()
        target_ori = controller.ee_ori_mat.copy()

        # 只统一视野角度,不照搬 Piper 的相机位置与朝向
        wrist_id = env.sim.model.camera_name2id("robot0_eye_in_hand")
        env.sim.model.cam_fovy[wrist_id] = 90.0

        print("预热前盘子位置:", observation["plate_1_pos"])
        hold_command = make_command(target_pos, target_ori, -1.0)
        for _ in range(SETTLE_STEPS):
            observation, _, _, _ = env.step(hold_command)
            env.env.render()
        print("预热后盘子位置:", observation["plate_1_pos"])
        assert not env.check_success(), "开局已经成功,请检查。"

        initial_state = env.get_sim_state().copy()

        # 将运行时修改的 fovy 明确写入 XML,供后续重放。
        root = ET.fromstring(env.sim.model.get_xml())
        camera_node = root.find(".//camera[@name='robot0_eye_in_hand']")
        assert camera_node is not None
        camera_node.set("fovy", "90.0")
        model_xml = ET.tostring(root, encoding="unicode")

        env_info = {
            "environment": type(env.env).__name__,
            "suite": "libero_spatial",
            "task_id": 0,
            "task": task.language,
            "bddl_file": str(bddl_path),
            "robots": ["Panda"],
            "seed": SEED,
            "split": "train",
            "init_index": args.init_index,
            "official_init_index": args.init_index,
            "init_source": "init_states.npz/train",
            "init_sha256": source_hash,
            "control_freq": CONTROL_FREQ,
            "settle_steps": SETTLE_STEPS,
            "physics_dt": float(env.sim.model.opt.timestep),
            "action_representation": (
                "absolute EEF target (persistent) + gripper"
            ),
            "teleop": {
                "pos_sensitivity": float(args.pos_sensitivity),
                "rot_sensitivity": ROT_SENSITIVITY,
                "keyboard_pos_step_m": 0.05,
                "note": "一次按键目标推进 = 0.05*sens*75*POSITION_GAIN",
            },
            "action_scaling": {
                "position_gain": POSITION_GAIN,
                "rotation_gain": ROTATION_GAIN,
                "max_position_lead": MAX_POSITION_LEAD,
                "max_rotation_lead": MAX_ROTATION_LEAD,
                "position_unit": "meter",
                "rotation_unit": "axis-angle rad",
            },
            "controller": {
                "type": controller.name,
                "control_delta": bool(controller.use_delta),
                "uncouple_pos_ori": bool(controller.uncoupling),
                **{
                    key: np.asarray(getattr(controller, key)).tolist()
                    for key in [
                        "kp", "kd", "input_min", "input_max",
                        "output_min", "output_max",
                    ]
                },
            },
            "wrist_camera": {
                "pos": env.sim.model.cam_pos[wrist_id].tolist(),
                "quat_wxyz": env.sim.model.cam_quat[wrist_id].tolist(),
                "fovy": 90.0,
            },
            "image_convention": macros.IMAGE_CONVENTION,
            "camera_names": ["agentview", "robot0_eye_in_hand"],
            "camera_size": [256, 256],
            "initial_objects": {
                key: np.asarray(value).tolist()
                for key, value in observation.items()
                if not key.startswith("robot0_")
                and key.endswith(("_pos", "_quat"))
            },
            "stationary_definition": {
                "max_arm_joint_change_rad": 0.001,
                "max_gripper_joint_change_m": 0.0001,
            },
        }

        device = CollectKeyboard(
            pos_sensitivity=args.pos_sensitivity,
            rot_sensitivity=ROT_SENSITIVITY,
        )
        device.start_control()

        print("\n训练开局:", args.init_index)
        print("首次运动或闭合夹爪开始录制。")
        print("W/A/S/D 平移;R/F 升降;空格开合夹爪。")
        print("Z/X、T/G、C/V 旋转;Q 放弃;N 标记重试并放弃。")
        print("放盘松爪并连续成功 10 步后自动保存,最多 1000 步。")
        print("若开局异常或明显失控,按 Q 或 Ctrl+C。")
        last_log = time.perf_counter()

        while True:
            start = time.perf_counter()
            action, _ = input2action(
                device=device, robot=robot, active_arm="right",
            )
            if device.retry:
                stop_reason = "手动标记重试,放弃本条"
                break
            if action is None:
                break
            assert action.shape == (7,)
            assert np.isfinite(action).all()

            if record_start is None and (
                np.any(action[:6] != 0) or action[6] > 0
            ):
                record_start = time.perf_counter()
                print("RECORDING_STARTED")

            # 首次操作前只显示画面,不继续推进物理时间。
            if record_start is not None:
                # states[t] 是执行 actions[t] 之前的完整仿真状态。
                state_before = env.get_sim_state().copy()
                arm_before = observation["robot0_joint_pos"].copy()
                gripper_before = observation["robot0_gripper_qpos"].copy()

                # 键盘增量累积到持久目标,超过允许前置量才截断
                controller.update(force=True)
                target_pos, target_ori = advance_target(
                    target_pos, target_ori, action, controller,
                )
                command = make_command(
                    target_pos, target_ori, float(action[6]),
                )

                if last_command is None or not np.array_equal(
                    command, last_command
                ):
                    pressed_frames += 1
                    if last_command is not None:
                        advance_total += float(np.linalg.norm(
                            command[:3] - last_command[:3]
                        ))
                last_command = command

                observation, _, _, _ = env.step(command)
                states.append(state_before)
                actions.append(command.copy())
                final_state = env.get_sim_state().copy()
                assert np.isfinite(final_state).all()

                stationary_flags.append(bool(
                    np.max(np.abs(
                        observation["robot0_joint_pos"] - arm_before
                    )) <= 0.001
                    and np.max(np.abs(
                        observation["robot0_gripper_qpos"] - gripper_before
                    )) <= 0.0001
                ))
                if command[6] != previous_gripper:
                    switches += 1
                previous_gripper = float(command[6])

            env.env.render()
            for key in CAMERAS:
                cv2.imshow(key, display_image(observation[key]))
            if cv2.waitKey(1) & 0xFF == ord("q"):
                break

            if record_start is not None:
                success_count = (
                    success_count + 1
                    if env.check_success() and command[6] < 0 else 0
                )
                if success_count >= SUCCESS_HOLD_STEPS:
                    success = True
                    stop_reason = "连续成功并持续发送松爪指令"
                    break
                if len(actions) >= MAX_STEPS:
                    stop_reason = "达到 1000 步仍未满足成功条件"
                    break

            if time.perf_counter() - last_log >= 2:
                ratio = pressed_frames / max(len(actions), 1)
                mean_adv = 1000 * advance_total / max(len(actions) - 1, 1)
                lead_now = float(np.linalg.norm(
                    command[:3] - controller.ee_pos
                ))
                flag = "" if ratio >= 0.5 and mean_adv >= 10 else "  <- 再按住一点"
                print(
                    "步数:", len(actions),
                    "夹爪切换:", switches,
                    "连续成功:", success_count,
                    "| 按键帧占比: %.1f%%" % (100 * ratio),
                    "平均推进: %.2fmm/帧" % mean_adv,
                    "当前领先: %.1fmm" % (1000 * lead_now),
                    flag,
                )
                last_log = time.perf_counter()

            remaining = 1 / CONTROL_FREQ - (time.perf_counter() - start)
            if remaining > 0:
                time.sleep(remaining)

    except KeyboardInterrupt:
        success = False
        stop_reason = "用户通过 Ctrl+C 结束"
    except Exception as error:
        success = False
        stop_reason = f"{type(error).__name__}: {error}"
        raise
    finally:
        try:
            if record_start is not None:
                elapsed = time.perf_counter() - record_start
                attempt = {
                    "timestamp": run_id,
                    "split": "train",
                    "init_index": args.init_index,
                    "success": success,
                    "steps": len(actions),
                    "recording_seconds": elapsed,
                    "collection_fps": len(actions) / elapsed,
                    "stationary_ratio": (
                        float(np.mean(stationary_flags))
                        if stationary_flags else None
                    ),
                    "gripper_switch_count": switches,
                    "retry_count_manual": int(device.retry),
                    "stop_reason": stop_reason,
                    "hdf5_path": None,
                }

                if success:
                    state_array = np.asarray(states, dtype=np.float64)
                    action_array = np.asarray(actions, dtype=np.float64)
                    assert state_array.shape == (len(actions), 92)
                    assert action_array.shape == (len(actions), 7)
                    assert len(stationary_flags) == len(actions)
                    assert np.isfinite(state_array).all()
                    assert np.isfinite(action_array).all()
                    np.testing.assert_array_equal(state_array[0], initial_state)
                    np.testing.assert_allclose(
                        np.diff(np.r_[state_array[:, 0], final_state[0]]),
                        1 / CONTROL_FREQ, rtol=0, atol=1e-6,
                    )

                    path = OUTPUT_DIR / (
                        f"demo_train_{args.init_index:02d}_{run_id}.hdf5"
                    )
                    pending = path.with_suffix(".hdf5.partial")
                    assert not path.exists()
                    with h5py.File(pending, "x") as file:
                        data = file.create_group("data")
                        data.attrs["env"] = env_info["environment"]
                        data.attrs["env_info"] = json.dumps(
                            env_info, ensure_ascii=False,
                        )
                        demo = data.create_group("demo_0")
                        demo.create_dataset("states", data=state_array)
                        demo.create_dataset("actions", data=action_array)
                        demo.create_dataset("settled_initial_state", data=initial_state)
                        demo.create_dataset("final_state", data=final_state)
                        demo.create_dataset("stationary_flags", data=stationary_flags)
                        demo.attrs["model_file"] = model_xml
                        demo.attrs["num_samples"] = len(actions)
                        for key, value in attempt.items():
                            if value is not None:
                                demo.attrs[key] = value

                    with h5py.File(pending, "r") as file:
                        demo = file["data/demo_0"]
                        np.testing.assert_array_equal(demo["states"][:], state_array)
                        np.testing.assert_array_equal(demo["actions"][:], action_array)
                        print("states shape:", demo["states"].shape)
                        print("actions shape:", demo["actions"].shape)
                    pending.rename(path)
                    attempt["hdf5_path"] = str(path)
                    print("B2_2_HDF5_READBACK_OK")
                    ratio = pressed_frames / max(len(actions), 1)
                    mean_adv = 1000 * advance_total / max(len(actions) - 1, 1)
                    print(
                        "B2_2_DENSITY 按键帧占比=%.1f%% 平均推进=%.2fmm/帧"
                        % (100 * ratio, mean_adv),
                        "(Piper: 100.0%% / 20.58mm)",
                    )
                    print("已保存:", path)
                else:
                    print("本次不保存训练示范,只记尝试日志。")

                # 单个采集进程使用;不要同时开多个本脚本。
                log_path = OUTPUT_DIR / "attempts.jsonl"
                with log_path.open("a", encoding="utf-8") as file:
                    file.write(json.dumps(attempt, ensure_ascii=False) + "\n")
                print(json.dumps(attempt, ensure_ascii=False, indent=2))
                progress(source_hash)

                with log_path.open(encoding="utf-8") as file:
                    attempts = [json.loads(line) for line in file if line.strip()]
                successes = sum(item["success"] for item in attempts)
                print("累计成功 / 已录制尝试:", successes, "/", len(attempts))
                print("采集成功率:", successes / len(attempts))
            else:
                print("未开始录制,不计入采集尝试。")
        finally:
            try:
                if device is not None:
                    device.listener.stop()
            finally:
                try:
                    if env is not None:
                        env.close()
                finally:
                    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()

重新采了60条数据,然后走跟之前同样的流程。

3.2. ACT结果对比

第二次采集的数据同样进行了ACT训练、评测,与第一次的对比图如下:

可以发现,第二次采集的数据评估成功率达到了80,这说明之前的问题已解决,这批数据可以拿来给pi05训练。

3.3. pi05结果对比

同样,使用跟之前piper训练相同的E方法训练了两万步,然后进行评估,与piper对比如下:

可以发现,在自己构建的数据集上,piper和panda各自训练了20000步,最终piper的loss更低,但评估却显示panda的两项成果率都高于piper,尤其是OOD,达到了78%的成功率,这说明其有不错的泛化能力。

对此我的理解为,pi05在预训练时就在panda的大规模数据集上训练过了,VLM模块的lora适当保留了其预训练阶段的泛化性,因此有相当好的评估结果。


总结

本文介绍了pi05在自建的piper和panda数据集上微调、评估和采坑的记录。

文章仅做学习分享,如有谬误,欢迎指正。

转载自 CSDN-专业IT技术社区

原文链接:https://blog.csdn.net/m0_75114363/article/details/168662015

文章来源转载

评论

评论列表

微信小程序
QQ小程序

关于作者

点赞数:0
关注数:0
粉丝:0
文章:0
关注标签:0
加入于:--