{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# 键盘采集示教数据\n", "\n", "默认机器人是 ROBOTIS OMY 六轴机械臂及夹爪的 MuJoCo 模型，场景为 `asset/example_scene_y.xml`。运行采集单元后，让 MuJoCo 窗口获得键盘焦点再操作。\n", "\n", "使用本章固定的 LeRobot 版本，以 v2.1 格式、20 Hz 保存至 `./demo_data`。`observation.state` 是下发本帧命令前的末端位姿；`action` 是目标关节角加归一化夹爪命令，由 `get_commanded_joint_action()` 读取，不能用 `step()` 返回的当前关节状态代替。\n", "\n", "在给定环境中采集示教数据。\n", "任务是抓起杯子并放到盘子上。当杯子在盘子上、夹爪打开且末端执行器位于杯子上方时，环境会判定成功。\n", "\n", "

\n", "\n", "键位说明：WASD 控制 x-y 平面移动，R/F 控制 z 轴，Q/E 控制倾斜，方向键控制其余旋转。\n", "\n", "空格键切换夹爪开合状态，Z 键会重置环境并丢弃当前回合缓存数据。\n", "\n", "叠加图像说明：\n", "- 右上：Agent 视角\n", "- 右下：腕部（第一人称）视角\n", "- 左上：侧视图\n", "- 左下：俯视图\n", "\n", "\n", "注意采集数据时，尽量避免像视频中那样操作，需更精准地采集。" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# pip install -r requirements.txt\n", "# 推荐python310环境" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "C:\\Users\\kewei\\micromamba\\envs\\py311\\Lib\\site-packages\\torch\\cuda\\__init__.py:61: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.\n", " import pynvml # type: ignore[import]\n" ] } ], "source": [ "import sys\n", "import random\n", "import numpy as np\n", "import os\n", "from PIL import Image\n", "from mujoco_env.y_env import SimpleEnv\n", "from lerobot.common.datasets.lerobot_dataset import LeRobotDataset" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "# If you want to randomize the object positions, set this to None\n", "# If you fix the seed, the object positions will be the same every time\n", "SEED = 0 \n", "# SEED = None <- Uncomment this line to randomize the object positions\n", "\n", "REPO_NAME = 'datawhale_eai_pnp' #pick-and-place\n", "NUM_DEMO = 1 # Number of demonstrations to collect\n", "ROOT = \"./demo_data\" # The root directory to save the demonstrations" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[Env Info] 运行环境已就绪，当前工作目录: C:\\Users\\kewei\\lerobot-mujoco-tutorial\n" ] } ], "source": [ "import env_config" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "记得解压这个文件夹\n", "\n", "![fig1](./media/fig1.png)" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "-----------------------------------------------------------------------------\n", "name:[Tabletop] dt:[0.002] HZ:[500]\n", " n_qpos:[24] n_qvel:[22] n_qacc:[22] n_ctrl:[10]\n", " integrator:[IMPLICITFAST]\n", "\n", "n_body:[21]\n", " [0/21] [world] mass:[0.00]kg\n", " [1/21] [front_object_table] mass:[1.00]kg\n", " [2/21] [camera] mass:[0.00]kg\n", " [3/21] [camera2] mass:[0.00]kg\n", " [4/21] [camera3] mass:[0.00]kg\n", " [5/21] [link1] mass:[2.06]kg\n", " [6/21] [link2] mass:[3.68]kg\n", " [7/21] [link3] mass:[2.39]kg\n", " [8/21] [link4] mass:[1.40]kg\n", " [9/21] [link5] mass:[1.40]kg\n", " [10/21] [link6] mass:[0.65]kg\n", " [11/21] [camera_center] mass:[0.00]kg\n", " [12/21] [tcp_link] mass:[0.32]kg\n", " [13/21] [rh_p12_rn_r1] mass:[0.07]kg\n", " [14/21] [rh_p12_rn_r2] mass:[0.02]kg\n", " [15/21] [rh_p12_rn_l1] mass:[0.07]kg\n", " [16/21] [rh_p12_rn_l2] mass:[0.02]kg\n", " [17/21] [body_obj_mug_5] mass:[0.00]kg\n", " [18/21] [object_mug_5] mass:[0.08]kg\n", " [19/21] [body_obj_plate_11] mass:[0.00]kg\n", " [20/21] [object_plate_11] mass:[0.10]kg\n", "body_total_mass:[13.27]kg\n", "\n", "n_geom:[83]\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]\n", "\n", "n_mesh:[79]\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']\n", "\n", "n_joint:[12]\n", " [0/12] [joint1] axis:[0. 0. 1.]\n", " [1/12] [joint2] axis:[0. 1. 0.]\n", " [2/12] [joint3] axis:[0. 1. 0.]\n", " [3/12] [joint4] axis:[0. 1. 0.]\n", " [4/12] [joint5] axis:[0. 0. 1.]\n", " [5/12] [joint6] axis:[0. 1. 0.]\n", " [6/12] [rh_r1] axis:[1. 0. 0.]\n", " [7/12] [rh_r2] axis:[-1. 0. 0.]\n", " [8/12] [rh_l1] axis:[-1. 0. 0.]\n", " [9/12] [rh_l2] axis:[1. 0. 0.]\n", " [10/12] [None] axis:[0. 0. 1.]\n", " [11/12] [None] axis:[0. 0. 1.]\n", "\n", "n_dof:[22] (=number of rows of Jacobian)\n", " [0/22] [None] attached joint:[joint1] body:[link1]\n", " [1/22] [None] attached joint:[joint2] body:[link2]\n", " [2/22] [None] attached joint:[joint3] body:[link3]\n", " [3/22] [None] attached joint:[joint4] body:[link4]\n", " [4/22] [None] attached joint:[joint5] body:[link5]\n", " [5/22] [None] attached joint:[joint6] body:[link6]\n", " [6/22] [None] attached joint:[rh_r1] body:[rh_p12_rn_r1]\n", " [7/22] [None] attached joint:[rh_r2] body:[rh_p12_rn_r2]\n", " [8/22] [None] attached joint:[rh_l1] body:[rh_p12_rn_l1]\n", " [9/22] [None] attached joint:[rh_l2] body:[rh_p12_rn_l2]\n", " [10/22] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [11/22] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [12/22] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [13/22] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [14/22] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [15/22] [None] attached joint:[None] body:[body_obj_mug_5]\n", " [16/22] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [17/22] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [18/22] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [19/22] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [20/22] [None] attached joint:[None] body:[body_obj_plate_11]\n", " [21/22] [None] attached joint:[None] body:[body_obj_plate_11]\n", "\n", "Free joint information. n_free_joint:[2]\n", " [0/2] [None] body_name_attached:[body_obj_mug_5]\n", " [1/2] [None] body_name_attached:[body_obj_plate_11]\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:[6]\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']\n", "-----------------------------------------------------------------------------\n", "env:[Tabletop] reset\n", "env:[Tabletop] reset\n", "env:[Tabletop] initalize viewer\n", "DONE INITIALIZATION\n" ] } ], "source": [ "TASK_NAME = 'Put mug cup on the plate' \n", "xml_path = './asset/example_scene_y.xml'\n", "# xml_path = 'asset/example_scene_y.xml'\n", "# Define the environment\n", "PnPEnv = SimpleEnv(xml_path, seed = SEED, state_type = 'joint_angle')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 定义数据集特征并创建你的数据集\n", "数据集结构如下：\n", "```\n", "fps = 20,\n", "features={\n", " \"observation.image\": {\n", " \"dtype\": \"image\",\n", " \"shape\": (256, 256, 3),\n", " \"names\": [\"height\", \"width\", \"channels\"],\n", " },\n", " \"observation.wrist_image\": {\n", " \"dtype\": \"image\",\n", " \"shape\": (256, 256, 3),\n", " \"names\": [\"height\", \"width\", \"channel\"],\n", " },\n", " \"observation.state\": {\n", " \"dtype\": \"float32\",\n", " \"shape\": (6,),\n", " \"names\": [\"state\"], # x, y, z, roll, pitch, yaw\n", " },\n", " \"action\": {\n", " \"dtype\": \"float32\",\n", " \"shape\": (7,),\n", " \"names\": [\"action\"], # 6 joint angles and 1 gripper\n", " },\n", " \"obj_init\": {\n", " \"dtype\": \"float32\",\n", " \"shape\": (6,),\n", " \"names\": [\"obj_init\"], # just the initial position of the object. Not used in training.\n", " },\n", "},\n", "```\n", "\n", "数据将保存在 `./demo_data` 目录中。\n" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Directory ./demo_data already exists.\n" ] }, { "name": "stdin", "output_type": "stream", "text": [ "Do you want to delete it? (y/n) y\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "WARNING:root:'torchcodec' is not available in your platform, falling back to 'pyav' as a default decoder\n" ] } ], "source": [ "create_new = True\n", "if os.path.exists(ROOT):\n", " print(f\"Directory {ROOT} already exists.\")\n", " ans = input(\"Do you want to delete it? (y/n) \")\n", " if ans == 'y':\n", " import shutil\n", " shutil.rmtree(ROOT)\n", " else:\n", " create_new = False\n", "\n", "\n", "if create_new:\n", " dataset = LeRobotDataset.create(\n", " repo_id=REPO_NAME,\n", " root = ROOT, \n", " robot_type=\"omy\",\n", " fps=20, # 20 frames per second\n", " features={\n", " \"observation.image\": {\n", " \"dtype\": \"image\",\n", " \"shape\": (256, 256, 3),\n", " \"names\": [\"height\", \"width\", \"channels\"],\n", " },\n", " \"observation.wrist_image\": {\n", " \"dtype\": \"image\",\n", " \"shape\": (256, 256, 3),\n", " \"names\": [\"height\", \"width\", \"channel\"],\n", " },\n", " \"observation.state\": {\n", " \"dtype\": \"float32\",\n", " \"shape\": (6,),\n", " \"names\": [\"state\"], # x, y, z, roll, pitch, yaw\n", " },\n", " \"action\": {\n", " \"dtype\": \"float32\",\n", " \"shape\": (7,),\n", " \"names\": [\"action\"], # 6 joint angles and 1 gripper\n", " },\n", " \"obj_init\": {\n", " \"dtype\": \"float32\",\n", " \"shape\": (6,),\n", " \"names\": [\"obj_init\"], # just the initial position of the object. Not used in training.\n", " },\n", " },\n", " image_writer_threads=10,\n", " image_writer_processes=5,\n", " )\n", "else:\n", " print(\"Load from previous dataset\")\n", " dataset = LeRobotDataset(REPO_NAME, root=ROOT)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 键盘控制\n", "你可以用键盘遥操作机械臂并采集数据\n", "```\n", "--------- -----------------------\n", " w -> backward\n", "s a d left forward right\n", "--------- -----------------------\n", "In x, y plane\n", "\n", "---------\n", "R: Moving Up\n", "F: Moving Down\n", "---------\n", "In z axis\n", "\n", "---------\n", "Q: Tilt left\n", "E: Tilt right\n", "UP: Look Upward\n", "Down: Look Donward\n", "Right: Turn right\n", "Left: Turn left\n", "---------\n", "For rotation\n", "\n", "---------\n", "SPACEBAR: Toggle Gripper\n", "--------\n", "\n", "---------\n", "z: reset\n", "--------\n", "```\n", "重置环境会删除当前示教回合的缓存数据，并重新开始采集。\n", "\n", "\n", "关于 Rotation（旋转）部分的详细解释：\n", "在机器人控制中，这部分通常对应 3D 空间中的欧拉角（Euler Angles）：\n", "\n", "Tilt (Q/E): 对应 Roll（翻滚角）。想象机械臂末端像拧螺丝一样左右倾斜。\n", "\n", "Look Upward/Downward (UP/Down): 对应 Pitch（俯仰角）。控制机械臂末端的“头”抬起或低下。\n", "\n", "Turn Left/Right (Left/Right): 对应 Yaw（偏航角）。控制机械臂末端在水平方向上左右摆动。" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### 现在开始遥操作并采集数据\n", "\n", "**要收到成功信号，你需要松开夹爪并将末端执行器上移到杯子上方。**\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "杯子和盘子中心在 XY 上要足够近（< 0.1）\n", "\n", "夹爪要“真正打开”（rh_r1 < 0.1，这个很容易卡住）\n", "\n", "末端执行器（tcp_link）要抬高到 z > 0.9\n", "\n", "以上都满足后，才会触发 done=True" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "action = np.zeros(7)\n", "episode_id = 0\n", "record_flag = False # Start recording when the robot starts moving\n", "while PnPEnv.env.is_viewer_alive() and episode_id < NUM_DEMO:\n", " PnPEnv.step_env()\n", " if PnPEnv.env.loop_every(HZ=20):\n", " # check if the episode is done\n", " done = PnPEnv.check_success()\n", " if done and record_flag:\n", " # Save the episode data and reset the environment\n", " dataset.save_episode()\n", " episode_id += 1\n", " record_flag = False\n", " if episode_id >= NUM_DEMO:\n", " break\n", " PnPEnv.reset(seed=SEED)\n", " continue\n", " # Teleoperate the robot and get delta end-effector pose with gripper\n", " action, reset = PnPEnv.teleop_robot()\n", " if not record_flag and np.any(action != 0):\n", " record_flag = True\n", " print(\"Start recording\")\n", " if reset:\n", " # Reset the environment and clear the episode buffer\n", " # This can be done by pressing 'z' key\n", " PnPEnv.reset(seed=SEED)\n", " # PnPEnv.reset()\n", " dataset.clear_episode_buffer()\n", " record_flag = False\n", " continue\n", " # Step the environment\n", " # Get the end-effector pose and images\n", " ee_pose = PnPEnv.get_ee_pose()\n", " agent_image,wrist_image = PnPEnv.grab_image()\n", " # # resize to 256x256\n", " agent_image = Image.fromarray(agent_image)\n", " wrist_image = Image.fromarray(wrist_image)\n", " agent_image = agent_image.resize((256, 256))\n", " wrist_image = wrist_image.resize((256, 256))\n", " agent_image = np.array(agent_image)\n", " wrist_image = np.array(wrist_image)\n", " PnPEnv.step(action)\n", " joint_action = PnPEnv.get_commanded_joint_action()\n", " if record_flag:\n", " # Add the frame to the dataset\n", " dataset.add_frame( {\n", " \"observation.image\": agent_image,\n", " \"observation.wrist_image\": wrist_image,\n", " \"observation.state\": ee_pose, \n", " \"action\": joint_action,\n", " \"obj_init\": PnPEnv.obj_init_pose,\n", " # \"task\": TASK_NAME,\n", " }, task = TASK_NAME\n", " )\n", " PnPEnv.render(teleop=True)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "这里相当于只用采集一条数据" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [ "PnPEnv.env.close_viewer()" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [], "source": [ "# Clean up the images folder\n", "import shutil\n", "shutil.rmtree(dataset.root / 'images')" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python (py311)", "language": "python", "name": "py311" }, "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.11.13" } }, "nbformat": 4, "nbformat_minor": 4 }