{ "cells": [ { "cell_type": "markdown", "id": "0ac2d99d", "metadata": {}, "source": [ "# 部署已训练的 pi_0 策略\n", "\n", "\n", "\n", "在仿真中部署已训练策略。\n" ] }, { "cell_type": "code", "execution_count": null, "id": "eba71c8f", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Requirement already satisfied: pytest in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (8.3.5)\n", "Requirement already satisfied: exceptiongroup>=1.0.0rc8 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from pytest) (1.2.2)\n", "Requirement already satisfied: iniconfig in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from pytest) (2.1.0)\n", "Requirement already satisfied: packaging in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from pytest) (24.2)\n", "Requirement already satisfied: pluggy<2,>=1.5 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from pytest) (1.5.0)\n", "Requirement already satisfied: tomli>=1 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from pytest) (2.2.1)\n", "\n", "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m25.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.1.1\u001b[0m\n", "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n", "Requirement already satisfied: transformers==4.48.0 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (4.48.0)\n", "Requirement already satisfied: filelock in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from transformers==4.48.0) (3.18.0)\n", "Requirement already satisfied: huggingface-hub<1.0,>=0.24.0 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from transformers==4.48.0) (0.29.3)\n", "Requirement already satisfied: numpy>=1.17 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from transformers==4.48.0) (1.26.4)\n", "Requirement already satisfied: packaging>=20.0 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from transformers==4.48.0) (24.2)\n", "Requirement already satisfied: pyyaml>=5.1 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from transformers==4.48.0) (6.0.2)\n", "Requirement already satisfied: regex!=2019.12.17 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from transformers==4.48.0) (2024.11.6)\n", "Requirement already satisfied: requests in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from transformers==4.48.0) (2.32.3)\n", "Requirement already satisfied: tokenizers<0.22,>=0.21 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from transformers==4.48.0) (0.21.1)\n", "Requirement already satisfied: safetensors>=0.4.1 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from transformers==4.48.0) (0.5.3)\n", "Requirement already satisfied: tqdm>=4.27 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from transformers==4.48.0) (4.67.1)\n", "Requirement already satisfied: fsspec>=2023.5.0 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.24.0->transformers==4.48.0) (2024.12.0)\n", "Requirement already satisfied: typing-extensions>=3.7.4.3 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.24.0->transformers==4.48.0) (4.12.2)\n", "Requirement already satisfied: charset-normalizer<4,>=2 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from requests->transformers==4.48.0) (3.4.1)\n", "Requirement already satisfied: idna<4,>=2.5 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from requests->transformers==4.48.0) (3.10)\n", "Requirement already satisfied: urllib3<3,>=1.21.1 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from requests->transformers==4.48.0) (2.3.0)\n", "Requirement already satisfied: certifi>=2017.4.17 in /home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages (from requests->transformers==4.48.0) (2025.1.31)\n", "\n", "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m25.0.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.1.1\u001b[0m\n", "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n" ] } ], "source": [ "!pip install pytest\n", "!pip install transformers==4.50.3" ] }, { "cell_type": "markdown", "id": "f8c0deb6", "metadata": {}, "source": [ "# 训练并部署 pi_0\n" ] }, { "cell_type": "markdown", "id": "e381eebd", "metadata": {}, "source": [ "### 【可选】下载数据集\n" ] }, { "cell_type": "code", "execution_count": 2, "id": "2f3e4d81", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "fatal: destination path 'datawhale_eai_pnp_language' already exists and is not an empty directory.\n" ] } ], "source": [ "'''\n", "If you want to use the collected dataset, please download it from Hugging Face.\n", "'''\n", "# Optional: skip cloning if demo_data_language already contains your collected data.\n", "!git lfs install\n", "!git clone https://huggingface.co/datasets/Datawhale/datawhale_eai_pnp_language demo_data_language\n", "!git -C demo_data_language lfs pull" ] }, { "cell_type": "markdown", "id": "4b6de363", "metadata": {}, "source": [ "## 第 1 步：修改配置文件 `pi0_datawhale_eai.yaml`\n" ] }, { "cell_type": "markdown", "id": "233d18e1", "metadata": {}, "source": [ "使用本目录的 [pi0_datawhale_eai.yaml](pi0_datawhale_eai.yaml)，配置说明见 [README](README.md)。公开样例和第 5 节自行采集的数据均放在 ./demo_data_language，无需修改配置。训练前在本目录执行 python check_training_configs.py。" ] }, { "cell_type": "markdown", "id": "e2525163", "metadata": {}, "source": [ "## 第 2 步：训练模型\n", "代码在 A100 上测试通过。\n" ] }, { "cell_type": "code", "execution_count": null, "id": "ba97ebae", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "INFO 2025-06-28 20:05:05 ils/utils.py:46 Cuda backend detected, using cuda.\n", "WARNING 2025-06-28 20:05:05 /policies.py:67 Device 'None' is not available. Switching to 'cuda'.\n", "Traceback (most recent call last):\n", " File \"/home/jeongeun/Dropbox/code/vla_prj/lerobot-mujoco-tutorial/train_pi0.py\", line 289, in \n", " train()\n", " File \"/home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages/lerobot/configs/parser.py\", line 226, in wrapper_inner\n", " response = fn(cfg, *args, **kwargs)\n", " File \"/home/jeongeun/Dropbox/code/vla_prj/lerobot-mujoco-tutorial/train_pi0.py\", line 110, in train\n", " cfg.validate()\n", " File \"/home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages/lerobot/configs/train.py\", line 101, in validate\n", " raise FileExistsError(\n", "FileExistsError: Output directory ckpt/pi0_datawhale_eai already exists and resume is False. Please change your output directory so that ckpt/pi0_datawhale_eai is not overwritten.\n" ] } ], "source": [ "!python train_model.py --config_path pi0_datawhale_eai.yaml" ] }, { "cell_type": "markdown", "id": "61c74aea", "metadata": {}, "source": [ "## 第 3 步：部署\n" ] }, { "cell_type": "code", "execution_count": 1, "id": "d198aa7b", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/home/jeongeun/.pyenv/versions/3.10.2/envs/lerobot/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", " from .autonotebook import tqdm as notebook_tqdm\n" ] } ], "source": [ "from lerobot.common.datasets.lerobot_dataset import LeRobotDataset, LeRobotDatasetMetadata\n", "import numpy as np\n", "from lerobot.common.datasets.utils import write_json, serialize_dict\n", "from lerobot.common.policies.pi0.configuration_pi0 import PI0Config\n", "from lerobot.common.policies.pi0.modeling_pi0 import PI0Policy\n", "from lerobot.configs.types import FeatureType\n", "from lerobot.common.datasets.factory import resolve_delta_timestamps\n", "from lerobot.common.datasets.utils import dataset_to_policy_features\n", "import torch\n", "from PIL import Image\n", "import torchvision" ] }, { "cell_type": "markdown", "id": "5d189f29", "metadata": {}, "source": [ "### 加载策略\n" ] }, { "cell_type": "code", "execution_count": 2, "id": "53895453", "metadata": {}, "outputs": [], "source": [ "device = 'cuda'" ] }, { "cell_type": "code", "execution_count": 4, "id": "dee90c0e", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "WARNING:root:Device 'None' is not available. Switching to 'cuda'.\n" ] } ], "source": [ "dataset_metadata = LeRobotDatasetMetadata(\"datawhale_eai_pnp_language\", root='./demo_data_language')\n", "features = dataset_to_policy_features(dataset_metadata.features)\n", "output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}\n", "input_features = {key: ft for key, ft in features.items() if key not in output_features}\n", "# Policies are initialized with a configuration class, in this case `DiffusionConfig`. For this example,\n", "# we'll just use the defaults and so no arguments other than input/output features need to be passed.\n", "# Temporal ensemble to make smoother trajectory predictions\n", "cfg = PI0Config(input_features=input_features, output_features=output_features, chunk_size= 5, n_action_steps=5)\n", "delta_timestamps = resolve_delta_timestamps(cfg, dataset_metadata)" ] }, { "cell_type": "code", "execution_count": 5, "id": "00c5180d", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Loading weights from local directory\n" ] }, { "data": { "text/plain": [ "PI0Policy(\n", " (normalize_inputs): Normalize(\n", " (buffer_observation_state): ParameterDict(\n", " (mean): Parameter containing: [torch.cuda.FloatTensor of size 6 (cuda:0)]\n", " (std): Parameter containing: [torch.cuda.FloatTensor of size 6 (cuda:0)]\n", " )\n", " )\n", " (normalize_targets): Normalize(\n", " (buffer_action): ParameterDict(\n", " (mean): Parameter containing: [torch.cuda.FloatTensor of size 7 (cuda:0)]\n", " (std): Parameter containing: [torch.cuda.FloatTensor of size 7 (cuda:0)]\n", " )\n", " )\n", " (unnormalize_outputs): Unnormalize(\n", " (buffer_action): ParameterDict(\n", " (mean): Parameter containing: [torch.cuda.FloatTensor of size 7 (cuda:0)]\n", " (std): Parameter containing: [torch.cuda.FloatTensor of size 7 (cuda:0)]\n", " )\n", " )\n", " (model): PI0FlowMatching(\n", " (paligemma_with_expert): PaliGemmaWithExpertModel(\n", " (paligemma): PaliGemmaForConditionalGeneration(\n", " (vision_tower): SiglipVisionModel(\n", " (vision_model): SiglipVisionTransformer(\n", " (embeddings): SiglipVisionEmbeddings(\n", " (patch_embedding): Conv2d(3, 1152, kernel_size=(14, 14), stride=(14, 14), padding=valid)\n", " (position_embedding): Embedding(256, 1152)\n", " )\n", " (encoder): SiglipEncoder(\n", " (layers): ModuleList(\n", " (0-26): 27 x SiglipEncoderLayer(\n", " (self_attn): SiglipSdpaAttention(\n", " (k_proj): Linear(in_features=1152, out_features=1152, bias=True)\n", " (v_proj): Linear(in_features=1152, out_features=1152, bias=True)\n", " (q_proj): Linear(in_features=1152, out_features=1152, bias=True)\n", " (out_proj): Linear(in_features=1152, out_features=1152, bias=True)\n", " )\n", " (layer_norm1): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)\n", " (mlp): SiglipMLP(\n", " (activation_fn): PytorchGELUTanh()\n", " (fc1): Linear(in_features=1152, out_features=4304, bias=True)\n", " (fc2): Linear(in_features=4304, out_features=1152, bias=True)\n", " )\n", " (layer_norm2): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)\n", " )\n", " )\n", " )\n", " (post_layernorm): LayerNorm((1152,), eps=1e-06, elementwise_affine=True)\n", " )\n", " )\n", " (multi_modal_projector): PaliGemmaMultiModalProjector(\n", " (linear): Linear(in_features=1152, out_features=2048, bias=True)\n", " )\n", " (language_model): GemmaForCausalLM(\n", " (model): GemmaModel(\n", " (embed_tokens): Embedding(257152, 2048, padding_idx=0)\n", " (layers): ModuleList(\n", " (0-17): 18 x GemmaDecoderLayer(\n", " (self_attn): GemmaAttention(\n", " (q_proj): Linear(in_features=2048, out_features=2048, bias=False)\n", " (k_proj): Linear(in_features=2048, out_features=256, bias=False)\n", " (v_proj): Linear(in_features=2048, out_features=256, bias=False)\n", " (o_proj): Linear(in_features=2048, out_features=2048, bias=False)\n", " )\n", " (mlp): GemmaMLP(\n", " (gate_proj): Linear(in_features=2048, out_features=16384, bias=False)\n", " (up_proj): Linear(in_features=2048, out_features=16384, bias=False)\n", " (down_proj): Linear(in_features=16384, out_features=2048, bias=False)\n", " (act_fn): PytorchGELUTanh()\n", " )\n", " (input_layernorm): GemmaRMSNorm((2048,), eps=1e-06)\n", " (post_attention_layernorm): GemmaRMSNorm((2048,), eps=1e-06)\n", " )\n", " )\n", " (norm): GemmaRMSNorm((2048,), eps=1e-06)\n", " (rotary_emb): GemmaRotaryEmbedding()\n", " )\n", " (lm_head): Linear(in_features=2048, out_features=257152, bias=False)\n", " )\n", " )\n", " (gemma_expert): GemmaForCausalLM(\n", " (model): GemmaModel(\n", " (embed_tokens): None\n", " (layers): ModuleList(\n", " (0-17): 18 x GemmaDecoderLayer(\n", " (self_attn): GemmaAttention(\n", " (q_proj): Linear(in_features=1024, out_features=2048, bias=False)\n", " (k_proj): Linear(in_features=1024, out_features=256, bias=False)\n", " (v_proj): Linear(in_features=1024, out_features=256, bias=False)\n", " (o_proj): Linear(in_features=2048, out_features=1024, bias=False)\n", " )\n", " (mlp): GemmaMLP(\n", " (gate_proj): Linear(in_features=1024, out_features=4096, bias=False)\n", " (up_proj): Linear(in_features=1024, out_features=4096, bias=False)\n", " (down_proj): Linear(in_features=4096, out_features=1024, bias=False)\n", " (act_fn): PytorchGELUTanh()\n", " )\n", " (input_layernorm): GemmaRMSNorm((1024,), eps=1e-06)\n", " (post_attention_layernorm): GemmaRMSNorm((1024,), eps=1e-06)\n", " )\n", " )\n", " (norm): GemmaRMSNorm((1024,), eps=1e-06)\n", " (rotary_emb): GemmaRotaryEmbedding()\n", " )\n", " (lm_head): Linear(in_features=1024, out_features=257152, bias=False)\n", " )\n", " )\n", " (state_proj): Linear(in_features=32, out_features=1024, bias=True)\n", " (action_in_proj): Linear(in_features=32, out_features=1024, bias=True)\n", " (action_out_proj): Linear(in_features=1024, out_features=32, bias=True)\n", " (action_time_mlp_in): Linear(in_features=2048, out_features=1024, bias=True)\n", " (action_time_mlp_out): Linear(in_features=1024, out_features=1024, bias=True)\n", " )\n", ")" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# We can now instantiate our policy with this config and the dataset stats.\n", "policy = PI0Policy.from_pretrained('./ckpt/pi0_datawhale_eai/checkpoints/last/pretrained_model', dataset_stats=dataset_metadata.stats)\n", "# You can load the trained policy from hub if you don't have the resources to train it.\n", "# policy = PI0Policy.from_pretrained(\"datawhale-eai/pi0_datawhale_eai\", config=cfg, dataset_stats=dataset_metadata.stats)\n", "policy.to(device)\n", "\n" ] }, { "cell_type": "markdown", "id": "0e5938d2", "metadata": {}, "source": [ "### 部署策略\n" ] }, { "cell_type": "code", "execution_count": 15, "id": "f2152c6f", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "-----------------------------------------------------------------------------\n", "name:[Tabletop] dt:[0.002] HZ:[500]\n", " n_qpos:[31] n_qvel:[28] n_qacc:[28] n_ctrl:[10]\n", " integrator:[IMPLICITFAST]\n", "\n", "n_body:[23]\n", " [0/23] [world] mass:[0.00]kg\n", " [1/23] [front_object_table] mass:[1.00]kg\n", " [2/23] [camera] mass:[0.00]kg\n", " [3/23] [camera2] mass:[0.00]kg\n", " [4/23] [camera3] mass:[0.00]kg\n", " [5/23] [link1] mass:[2.06]kg\n", " [6/23] [link2] mass:[3.68]kg\n", " [7/23] [link3] mass:[2.39]kg\n", " [8/23] [link4] mass:[1.40]kg\n", " [9/23] [link5] mass:[1.40]kg\n", " [10/23] [link6] mass:[0.65]kg\n", " [11/23] [camera_center] mass:[0.00]kg\n", " [12/23] [tcp_link] mass:[0.32]kg\n", " [13/23] [rh_p12_rn_r1] mass:[0.07]kg\n", " [14/23] [rh_p12_rn_r2] mass:[0.02]kg\n", " [15/23] [rh_p12_rn_l1] mass:[0.07]kg\n", " [16/23] [rh_p12_rn_l2] mass:[0.02]kg\n", " [17/23] [body_obj_mug_5] mass:[0.00]kg\n", " [18/23] [object_mug_5] mass:[0.08]kg\n", " [19/23] [body_obj_plate_11] mass:[0.00]kg\n", " [20/23] [object_plate_11] mass:[0.10]kg\n", " [21/23] [body_obj_mug_6] mass:[0.00]kg\n", " [22/23] [object_mug_6] mass:[0.08]kg\n", "body_total_mass:[13.35]kg\n", "\n", "n_geom:[116]\n", "geom_names:['floor', None, 'front_object_table', None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None]\n", "\n", "n_mesh:[112]\n", "mesh_names:['base_unit', 'link1', 'link2', 'link3', 'link4', 'link5', 'link6', 'flange', 'base', 'r1', 'r2', 'l1', 'l2', 'mug_5_normalized_0_vis', 'mug_5_normalized_collision_22._coll', 'mug_5_normalized_collision_23._coll', 'mug_5_normalized_collision_21._coll', 'mug_5_normalized_collision_20._coll', 'mug_5_normalized_collision_24._coll', 'mug_5_normalized_collision_30._coll', 'mug_5_normalized_collision_18._coll', 'mug_5_normalized_collision_19._coll', 'mug_5_normalized_collision_31._coll', 'mug_5_normalized_collision_25._coll', 'mug_5_normalized_collision_27._coll', 'mug_5_normalized_collision_26._coll', 'mug_5_normalized_collision_9._coll', 'mug_5_normalized_collision_8._coll', 'mug_5_normalized_collision_6._coll', 'mug_5_normalized_collision_7._coll', 'mug_5_normalized_collision_5._coll', 'mug_5_normalized_collision_4._coll', 'mug_5_normalized_collision_0._coll', 'mug_5_normalized_collision_1._coll', 'mug_5_normalized_collision_3._coll', 'mug_5_normalized_collision_2._coll', 'mug_5_normalized_collision_17._coll', 'mug_5_normalized_collision_16._coll', 'mug_5_normalized_collision_28._coll', 'mug_5_normalized_collision_14._coll', 'mug_5_normalized_collision_15._coll', 'mug_5_normalized_collision_29._coll', 'mug_5_normalized_collision_11._coll', 'mug_5_normalized_collision_10._coll', 'mug_5_normalized_collision_12._coll', 'mug_5_normalized_collision_13._coll', 'plate_11_normalized_0_vis', 'plate_11_normalized_collision_22._coll', 'plate_11_normalized_collision_23._coll', 'plate_11_normalized_collision_21._coll', 'plate_11_normalized_collision_20._coll', 'plate_11_normalized_collision_24._coll', 'plate_11_normalized_collision_30._coll', 'plate_11_normalized_collision_18._coll', 'plate_11_normalized_collision_19._coll', 'plate_11_normalized_collision_31._coll', 'plate_11_normalized_collision_25._coll', 'plate_11_normalized_collision_27._coll', 'plate_11_normalized_collision_26._coll', 'plate_11_normalized_collision_9._coll', 'plate_11_normalized_collision_8._coll', 'plate_11_normalized_collision_6._coll', 'plate_11_normalized_collision_7._coll', 'plate_11_normalized_collision_5._coll', 'plate_11_normalized_collision_4._coll', 'plate_11_normalized_collision_0._coll', 'plate_11_normalized_collision_1._coll', 'plate_11_normalized_collision_3._coll', 'plate_11_normalized_collision_2._coll', 'plate_11_normalized_collision_17._coll', 'plate_11_normalized_collision_16._coll', 'plate_11_normalized_collision_28._coll', 'plate_11_normalized_collision_14._coll', 'plate_11_normalized_collision_15._coll', 'plate_11_normalized_collision_29._coll', 'plate_11_normalized_collision_11._coll', 'plate_11_normalized_collision_10._coll', 'plate_11_normalized_collision_12._coll', 'plate_11_normalized_collision_13._coll', 'mug_6_normalized_0_vis', 'mug_6_normalized_collision_22._coll', 'mug_6_normalized_collision_23._coll', 'mug_6_normalized_collision_21._coll', 'mug_6_normalized_collision_20._coll', 'mug_6_normalized_collision_24._coll', 'mug_6_normalized_collision_30._coll', 'mug_6_normalized_collision_18._coll', 'mug_6_normalized_collision_19._coll', 'mug_6_normalized_collision_31._coll', 'mug_6_normalized_collision_25._coll', 'mug_6_normalized_collision_27._coll', 'mug_6_normalized_collision_26._coll', 'mug_6_normalized_collision_9._coll', 'mug_6_normalized_collision_8._coll', 'mug_6_normalized_collision_6._coll', 'mug_6_normalized_collision_7._coll', 'mug_6_normalized_collision_5._coll', 'mug_6_normalized_collision_4._coll', 'mug_6_normalized_collision_0._coll', 'mug_6_normalized_collision_1._coll', 'mug_6_normalized_collision_3._coll', 'mug_6_normalized_collision_2._coll', 'mug_6_normalized_collision_17._coll', 'mug_6_normalized_collision_16._coll', 'mug_6_normalized_collision_28._coll', 'mug_6_normalized_collision_14._coll', 'mug_6_normalized_collision_15._coll', 'mug_6_normalized_collision_29._coll', 'mug_6_normalized_collision_11._coll', 'mug_6_normalized_collision_10._coll', 'mug_6_normalized_collision_12._coll', 'mug_6_normalized_collision_13._coll']\n", "\n", "n_joint:[13]\n", " [0/13] [joint1] axis:[0. 0. 1.]\n", " [1/13] [joint2] axis:[0. 1. 0.]\n", " [2/13] [joint3] axis:[0. 1. 0.]\n", " [3/13] [joint4] axis:[0. 1. 0.]\n", " [4/13] [joint5] axis:[0. 0. 1.]\n", " [5/13] [joint6] axis:[0. 1. 0.]\n", " [6/13] [rh_r1] axis:[1. 0. 0.]\n", " [7/13] [rh_r2] axis:[-1. 0. 0.]\n", " [8/13] [rh_l1] axis:[-1. 0. 0.]\n", " [9/13] [rh_l2] axis:[1. 0. 0.]\n", " [10/13] [None] axis:[0. 0. 1.]\n", " [11/13] [None] axis:[0. 0. 1.]\n", " [12/13] [None] axis:[0. 0. 1.]\n", "\n", "n_dof:[28] (=number of rows of Jacobian)\n", " [0/28] [None] attached joint:[joint1] body:[link1]\n", " [1/28] [None] attached joint:[joint2] body:[link2]\n", " [2/28] [None] attached joint:[joint3] body:[link3]\n", " [3/28] [None] attached joint:[joint4] body:[link4]\n", " [4/28] [None] attached joint:[joint5] body:[link5]\n", " [5/28] [None] attached joint:[joint6] body:[link6]\n", " [6/28] [None] attached joint:[rh_r1] body:[rh_p12_rn_r1]\n", " [7/28] [None] attached joint:[rh_r2] body:[rh_p12_rn_r2]\n", " [8/28] [None] attached joint:[rh_l1] body:[rh_p12_rn_l1]\n", " [9/28] [None] attached joint:[rh_l2] body:[rh_p12_rn_l2]\n", " [10/28] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [11/28] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [12/28] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [13/28] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [14/28] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [15/28] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [16/28] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [17/28] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [18/28] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [19/28] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [20/28] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [21/28] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [22/28] [None] attached joint:[None] body:[body_obj_mug_6]\n", " [23/28] [None] attached joint:[None] body:[body_obj_mug_6]\n", " [24/28] [None] attached joint:[None] body:[body_obj_mug_6]\n", " [25/28] [None] attached joint:[None] body:[body_obj_mug_6]\n", " [26/28] [None] attached joint:[None] body:[body_obj_mug_6]\n", " [27/28] [None] attached joint:[None] body:[body_obj_mug_6]\n", "\n", "Free joint information. n_free_joint:[3]\n", " [0/3] [None] body_name_attached:[body_obj_mug_5]\n", " [1/3] [None] body_name_attached:[body_obj_plate_11]\n", " [2/3] [None] body_name_attached:[body_obj_mug_6]\n", "\n", "Revolute joint information. n_rev_joint:[10]\n", " [0/10] [joint1] range:[-6.283]~[6.283]\n", " [1/10] [joint2] range:[-6.283]~[6.283]\n", " [2/10] [joint3] range:[-6.283]~[6.283]\n", " [3/10] [joint4] range:[-6.283]~[6.283]\n", " [4/10] [joint5] range:[-6.283]~[6.283]\n", " [5/10] [joint6] range:[-6.283]~[6.283]\n", " [6/10] [rh_r1] range:[0.000]~[1.100]\n", " [7/10] [rh_r2] range:[0.000]~[1.000]\n", " [8/10] [rh_l1] range:[0.000]~[1.100]\n", " [9/10] [rh_l2] range:[0.000]~[1.000]\n", "\n", "Prismatic joint information. n_pri_joint:[0]\n", "\n", "Control information. n_ctrl:[10]\n", " [0/10] [actuator_joint1] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [1/10] [actuator_joint2] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [2/10] [actuator_joint3] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [3/10] [actuator_joint4] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [4/10] [actuator_joint5] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [5/10] [actuator_joint6] range:[-6.283]~[6.283] gear:[1.00] type:[JOINT]\n", " [6/10] [actuator_rh_r1] range:[0.000]~[1.100] gear:[1.00] type:[JOINT]\n", " [7/10] [actuator_rh_r2] range:[0.000]~[1.000] gear:[1.00] type:[JOINT]\n", " [8/10] [actuator_rh_l1] range:[0.000]~[1.100] gear:[1.00] type:[JOINT]\n", " [9/10] [actuator_rh_l2] range:[0.000]~[1.000] gear:[1.00] type:[JOINT]\n", "\n", "Camera information. n_cam:[4]\n", " [0/4] [agentview] fov:[60.0]\n", " [1/4] [topview] fov:[90.0]\n", " [2/4] [sideview] fov:[90.0]\n", " [3/4] [egocentric] fov:[90.0]\n", "\n", "n_sensor:[0]\n", "sensor_names:[]\n", "n_site:[9]\n", "site_names:['bottom_site_mug_5', 'top_site_mug_5', 'horizontal_radius_site_mug_5', 'bottom_site_plate_11', 'top_site_plate_11', 'horizontal_radius_site_plate_11', 'bottom_site_mug_6', 'top_site_mug_6', 'horizontal_radius_site_mug_6']\n", "-----------------------------------------------------------------------------\n", "env:[Tabletop] reset\n", "env:[Tabletop] reset\n", "env:[Tabletop] initalize viewer\n", "DONE INITIALIZATION\n" ] } ], "source": [ "from mujoco_env.y_env2 import SimpleEnv2\n", "xml_path = './asset/example_scene_y2.xml'\n", "PnPEnv = SimpleEnv2(xml_path, action_type='joint_angle')" ] }, { "cell_type": "code", "execution_count": null, "id": "ac83797a", "metadata": {}, "outputs": [], "source": [ "from torchvision import transforms\n", "# Approach 1: Using torchvision.transforms\n", "def get_default_transform(image_size: int = 224):\n", " \"\"\"\n", " Returns a torchvision transform that:\n", " Converts to a FloatTensor and scales pixel values [0,255] -> [0.0,1.0]\n", " \"\"\"\n", " return transforms.Compose([\n", " transforms.ToTensor(), # PIL [0–255] -> FloatTensor [0.0–1.0], shape C×H×W\n", " ])" ] }, { "cell_type": "code", "execution_count": 17, "id": "82e9e6fa", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "DONE INITIALIZATION\n", "Success\n", "DONE INITIALIZATION\n", "Success\n", "DONE INITIALIZATION\n", "Success\n", "DONE INITIALIZATION\n", "Success\n", "DONE INITIALIZATION\n", "Success\n", "DONE INITIALIZATION\n" ] } ], "source": [ "step = 0\n", "PnPEnv.reset(seed=0)\n", "policy.reset()\n", "policy.eval()\n", "save_image = True\n", "IMG_TRANSFORM = get_default_transform()\n", "while PnPEnv.env.is_viewer_alive():\n", " PnPEnv.step_env()\n", " if PnPEnv.env.loop_every(HZ=20):\n", " # Check if the task is completed\n", " success = PnPEnv.check_success()\n", " if success:\n", " print('Success')\n", " # Reset the environment and action queue\n", " policy.reset()\n", " PnPEnv.reset()\n", " step = 0\n", " save_image = False\n", " # Get the current state of the environment\n", " state = PnPEnv.get_joint_state()[:6]\n", " # Get the current image from the environment\n", " image, wirst_image = PnPEnv.grab_image()\n", " image = Image.fromarray(image)\n", " image = image.resize((256, 256))\n", " image = IMG_TRANSFORM(image)\n", " wrist_image = Image.fromarray(wirst_image)\n", " wrist_image = wrist_image.resize((256, 256))\n", " wrist_image = IMG_TRANSFORM(wrist_image)\n", " data = {\n", " 'observation.state': torch.tensor([state]).to(device),\n", " 'observation.image': image.unsqueeze(0).to(device),\n", " 'observation.wrist_image': wrist_image.unsqueeze(0).to(device),\n", " 'task': [PnPEnv.instruction],\n", " }\n", " # Select an action\n", " action = policy.select_action(data)\n", " action = action[0,:7].cpu().detach().numpy()\n", " # Take a step in the environment\n", " _ = PnPEnv.step(action)\n", " PnPEnv.render()\n", " step += 1\n", " success = PnPEnv.check_success()\n", " if success:\n", " print('Success')\n", " break" ] }, { "cell_type": "code", "execution_count": 18, "id": "b38e0469", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "model.safetensors: 100%|██████████| 7.54G/7.54G [14:34<00:00, 8.61MB/s] \n" ] }, { "data": { "text/plain": [ "CommitInfo(commit_url='https://huggingface.co/datawhale-eai/pi0_datawhale_eai/commit/85152b3cbefeb37f26d3d8682e2db187621894c3', commit_message='Add trained policy for PnP task', commit_description='', oid='85152b3cbefeb37f26d3d8682e2db187621894c3', pr_url=None, repo_url=RepoUrl('https://huggingface.co/datawhale-eai/pi0_datawhale_eai', endpoint='https://huggingface.co', repo_type='model', repo_id='datawhale-eai/pi0_datawhale_eai'), pr_revision=None, pr_num=None)" ] }, "execution_count": 18, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# policy.push_to_hub(\n", "# repo_id='datawhale-eai/pi0_datawhale_eai',\n", "# commit_message='Add trained policy for PnP task',\n", "# )" ] }, { "cell_type": "code", "execution_count": null, "id": "755d81b9", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "id": "92589fed-5b85-42e3-a815-947afcd22022", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "markdown", "id": "4781e821-d4d0-46fd-be18-9d662a85c9be", "metadata": {}, "source": [ "$\\pi_0$ 之所以能在具身智能圈引发巨大轰动，是因为它极其优雅地解决了一个历史难题：**如何让拥有“互联网级智慧”的视觉语言大模型（VLM）学会“干粗活（高频连续的底层控制）”而不变笨。**\n", "\n", "我们先用一张 Mermaid 流程图来看看它的数据流向和核心组件，然后再深入拆解。\n", "\n", "### 📊 $\\pi_0$ 架构工作流\n", "\n", "代码段\n", "\n", "```mermaid\n", "graph TD\n", " %% Inputs\n", " subgraph Inputs [输入层 Inputs]\n", " IMG[视觉输入 Wrist/Base Camera]\n", " TXT[文本指令 Language Prompt]\n", " STATE[本体状态 State e.g. 6D关节角]\n", " NOISE[随机噪声动作 Noisy Action + 时间步 t]\n", " end\n", "\n", " %% Vision-Language Backbone\n", " subgraph Brain [PaliGemma 大脑：语义理解]\n", " SIGLIP[SigLIP 视觉编码器]\n", " PROJ[多模态投影器 Multi-modal Projector]\n", " LLM_MAIN[Gemma 语言模型 2048-dim]\n", " end\n", "\n", " %% Action Expert\n", " subgraph Cerebellum [Gemma Expert 小脑：运动控制]\n", " STATE_PROJ[状态投影 State Proj 32->1024]\n", " ACT_PROJ[动作/时间投影 Action/Time Proj]\n", " EXPERT_LLM[Gemma 动作专家模型 1024-dim]\n", " end\n", "\n", " %% Flow Matching Output\n", " subgraph Outputs [输出层 Outputs]\n", " OUT_PROJ[Action Out Proj 1024->32]\n", " FLOW[Flow Matching 向量场预测]\n", " ACT_CLEAN[去噪后的连续动作 Continuous Actions]\n", " end\n", "\n", " %% Connections\n", " IMG --> SIGLIP --> PROJ --> LLM_MAIN\n", " TXT --> LLM_MAIN\n", "\n", " STATE --> STATE_PROJ --> EXPERT_LLM\n", " NOISE --> ACT_PROJ --> EXPERT_LLM\n", " \n", " %% The magic interaction\n", " LLM_MAIN -.->|深层特征融合与交叉注意力| EXPERT_LLM\n", "\n", " EXPERT_LLM --> OUT_PROJ --> FLOW\n", " FLOW -->|ODE求解与去噪迭代| ACT_CLEAN\n", "```\n", "\n", "------\n", "\n", "### 🧠 代码拆解：$\\pi_0$ 到底包含了什么？\n", "\n", "根据你提供的代码输出，$\\pi_0$ 模型主要由以下三大模块构成：\n", "\n", "#### 1. 互联网级的大脑：`PaliGemmaForConditionalGeneration`\n", "\n", "这是模型的基础（Base Model），负责“看”和“想”。\n", "\n", "- **眼睛 (`SiglipVisionModel`)**：将输入的图像切成 Patch，通过 27 层的 Vision Transformer 提取极其强大的视觉特征。\n", "- **翻译官 (`PaliGemmaMultiModalProjector`)**：将 1152 维的视觉特征映射到 2048 维，让语言模型能“看懂”图像。\n", "- **思考核心 (`GemmaForCausalLM` - 18层)**：处理文本指令（如“把苹果放进碗里”）和图像特征，理解环境的上下文。\n", "\n", "#### 2. 运动控制的小脑：`gemma_expert`\n", "\n", "这正是 $\\pi_0$ 最核心的创新点！\n", "\n", "在传统的 Vision-Language-Action (VLA) 模型中（比如 Google 的 RT-2），研究人员强行让主 LLM 把动作当成文本单词（Token）输出。这导致模型容易“灾难性遗忘”（学会了控制机器人，却不会回答普通问题了）。\n", "\n", "- $\\pi_0$ 引入了一个**专门的专家网络（18层，1024维度的轻量级 Gemma）**。\n", "- 主脑 PaliGemma 负责高层语义的理解，然后把特征传递给这个 `gemma_expert`。专家网络结合机器人当前的物理状态（`state_proj`）和历史动作，专注于计算如何驱动电机。这种“脑体分离又融合”的设计极其巧妙。\n", "\n", "#### 3. 连续生成的魔法：`PI0FlowMatching`\n", "\n", "注意看代码外层的 `PI0FlowMatching`。模型并不是直接输出下一个动作的坐标，而是采用了**流匹配（Flow Matching）**。\n", "\n", "- **什么是 Flow Matching？** 它和 Diffusion（扩散模型）原理类似，是一种连续时间的生成模型。你给它一个随机噪声动作，它通过输入的时间步 $t$ (`action_time_mlp_in`)，一步步预测一个“向量场（Vector Field）”，将噪声推向真实、平滑的专家动作分布。\n", "- **为什么不用离散 Token？** 机器人的控制需要极高的精度和频率（比如 50Hz 到 100Hz）。把动作变成离散的词汇会损失精度。Flow Matching 能够原生地生成高维的、连续的、极其平滑的动作轨迹。\n", "\n", "------\n", "\n", "### 🚀 为什么 $\\pi_0$ 效果这么好，影响这么广泛？\n", "\n", "1. **破解了“语义”与“控制”的冲突（Alignment Tax）**\n", "\n", " 过去，用极其庞大的互联网图文数据训练出的 VLM，在微调适配小规模的机器人运动数据时，很容易顾此失彼。$\\pi_0$ 的 `PaliGemmaWithExpertModel` 结构，相当于给爱因斯坦（VLM）配了一个专业的赛车手替身（Expert）。爱因斯坦负责指挥方向，赛车手负责踩油门打方向盘，各司其职。\n", "\n", "2. **极强的泛化能力（Generalist Policy）**\n", "\n", " 因为底座是强大的 PaliGemma，$\\pi_0$ 具备开箱即用的 Zero-shot 空间推理能力。你换一个颜色不同的碗，或者换一个背景，它的“大脑”依然能认出来，并稳定地将信号传给“小脑”执行抓取。\n", "\n", "3. **支持多机器人平台（跨具身本体）**\n", "\n", " 通过灵活的 `state_proj` 和 `action_in/out_proj` 线性层（代码中是 32 维特征转换），它可以轻松适配 6 自由度机械臂、7 自由度机械臂、甚至是灵巧手和四足机器狗，只需要更换极其轻量的外围投影层即可。\n", "\n", "4. **Flow Matching 带来的丝滑控制**\n", "\n", " 传统的 MSE（均方误差）回归在遇到“同一种任务有多种做法”时，会输出平均值（导致动作四不像）。Flow Matching 这种生成式策略能够捕捉“多峰分布”，让机器人做出的动作像人类一样自然、连贯，且具有极强的抗干扰能力。" ] }, { "cell_type": "code", "execution_count": null, "id": "6fb7a792-ba68-4f41-8083-9991687ac124", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.3" } }, "nbformat": 4, "nbformat_minor": 5 }