{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# 部署已训练策略\n", "\n", "\n", "\n", "在仿真中部署已训练策略。\n" ] }, { "cell_type": "code", "execution_count": 1, "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.act.configuration_act import ACTConfig\n", "from lerobot.common.policies.act.modeling_act import ACTPolicy\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", "metadata": {}, "source": [ "## 加载策略\n" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "device = 'cuda'" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "WARNING:root:Device 'None' is not available. Switching to 'cuda'.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Loading weights from local directory\n" ] }, { "data": { "text/plain": [ "ACTPolicy(\n", " (normalize_inputs): Normalize(\n", " (buffer_observation_image): ParameterDict(\n", " (mean): Parameter containing: [torch.cuda.FloatTensor of size 3x1x1 (cuda:0)]\n", " (std): Parameter containing: [torch.cuda.FloatTensor of size 3x1x1 (cuda:0)]\n", " )\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): ACT(\n", " (vae_encoder): ACTEncoder(\n", " (layers): ModuleList(\n", " (0-3): 4 x ACTEncoderLayer(\n", " (self_attn): MultiheadAttention(\n", " (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n", " )\n", " (linear1): Linear(in_features=512, out_features=3200, bias=True)\n", " (dropout): Dropout(p=0.1, inplace=False)\n", " (linear2): Linear(in_features=3200, out_features=512, bias=True)\n", " (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n", " (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n", " (dropout1): Dropout(p=0.1, inplace=False)\n", " (dropout2): Dropout(p=0.1, inplace=False)\n", " )\n", " )\n", " (norm): Identity()\n", " )\n", " (vae_encoder_cls_embed): Embedding(1, 512)\n", " (vae_encoder_robot_state_input_proj): Linear(in_features=6, out_features=512, bias=True)\n", " (vae_encoder_action_input_proj): Linear(in_features=7, out_features=512, bias=True)\n", " (vae_encoder_latent_output_proj): Linear(in_features=512, out_features=64, bias=True)\n", " (backbone): IntermediateLayerGetter(\n", " (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n", " (bn1): FrozenBatchNorm2d(64, eps=1e-05)\n", " (relu): ReLU(inplace=True)\n", " (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n", " (layer1): Sequential(\n", " (0): BasicBlock(\n", " (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n", " (bn1): FrozenBatchNorm2d(64, eps=1e-05)\n", " (relu): ReLU(inplace=True)\n", " (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n", " (bn2): FrozenBatchNorm2d(64, eps=1e-05)\n", " )\n", " (1): BasicBlock(\n", " (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n", " (bn1): FrozenBatchNorm2d(64, eps=1e-05)\n", " (relu): ReLU(inplace=True)\n", " (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n", " (bn2): FrozenBatchNorm2d(64, eps=1e-05)\n", " )\n", " )\n", " (layer2): Sequential(\n", " (0): BasicBlock(\n", " (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n", " (bn1): FrozenBatchNorm2d(128, eps=1e-05)\n", " (relu): ReLU(inplace=True)\n", " (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n", " (bn2): FrozenBatchNorm2d(128, eps=1e-05)\n", " (downsample): Sequential(\n", " (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n", " (1): FrozenBatchNorm2d(128, eps=1e-05)\n", " )\n", " )\n", " (1): BasicBlock(\n", " (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n", " (bn1): FrozenBatchNorm2d(128, eps=1e-05)\n", " (relu): ReLU(inplace=True)\n", " (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n", " (bn2): FrozenBatchNorm2d(128, eps=1e-05)\n", " )\n", " )\n", " (layer3): Sequential(\n", " (0): BasicBlock(\n", " (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n", " (bn1): FrozenBatchNorm2d(256, eps=1e-05)\n", " (relu): ReLU(inplace=True)\n", " (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n", " (bn2): FrozenBatchNorm2d(256, eps=1e-05)\n", " (downsample): Sequential(\n", " (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n", " (1): FrozenBatchNorm2d(256, eps=1e-05)\n", " )\n", " )\n", " (1): BasicBlock(\n", " (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n", " (bn1): FrozenBatchNorm2d(256, eps=1e-05)\n", " (relu): ReLU(inplace=True)\n", " (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n", " (bn2): FrozenBatchNorm2d(256, eps=1e-05)\n", " )\n", " )\n", " (layer4): Sequential(\n", " (0): BasicBlock(\n", " (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n", " (bn1): FrozenBatchNorm2d(512, eps=1e-05)\n", " (relu): ReLU(inplace=True)\n", " (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n", " (bn2): FrozenBatchNorm2d(512, eps=1e-05)\n", " (downsample): Sequential(\n", " (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n", " (1): FrozenBatchNorm2d(512, eps=1e-05)\n", " )\n", " )\n", " (1): BasicBlock(\n", " (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n", " (bn1): FrozenBatchNorm2d(512, eps=1e-05)\n", " (relu): ReLU(inplace=True)\n", " (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n", " (bn2): FrozenBatchNorm2d(512, eps=1e-05)\n", " )\n", " )\n", " )\n", " (encoder): ACTEncoder(\n", " (layers): ModuleList(\n", " (0-3): 4 x ACTEncoderLayer(\n", " (self_attn): MultiheadAttention(\n", " (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n", " )\n", " (linear1): Linear(in_features=512, out_features=3200, bias=True)\n", " (dropout): Dropout(p=0.1, inplace=False)\n", " (linear2): Linear(in_features=3200, out_features=512, bias=True)\n", " (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n", " (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n", " (dropout1): Dropout(p=0.1, inplace=False)\n", " (dropout2): Dropout(p=0.1, inplace=False)\n", " )\n", " )\n", " (norm): Identity()\n", " )\n", " (decoder): ACTDecoder(\n", " (layers): ModuleList(\n", " (0): ACTDecoderLayer(\n", " (self_attn): MultiheadAttention(\n", " (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n", " )\n", " (multihead_attn): MultiheadAttention(\n", " (out_proj): NonDynamicallyQuantizableLinear(in_features=512, out_features=512, bias=True)\n", " )\n", " (linear1): Linear(in_features=512, out_features=3200, bias=True)\n", " (dropout): Dropout(p=0.1, inplace=False)\n", " (linear2): Linear(in_features=3200, out_features=512, bias=True)\n", " (norm1): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n", " (norm2): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n", " (norm3): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n", " (dropout1): Dropout(p=0.1, inplace=False)\n", " (dropout2): Dropout(p=0.1, inplace=False)\n", " (dropout3): Dropout(p=0.1, inplace=False)\n", " )\n", " )\n", " (norm): LayerNorm((512,), eps=1e-05, elementwise_affine=True)\n", " )\n", " (encoder_robot_state_input_proj): Linear(in_features=6, out_features=512, bias=True)\n", " (encoder_latent_input_proj): Linear(in_features=32, out_features=512, bias=True)\n", " (encoder_img_feat_input_proj): Conv2d(512, 512, kernel_size=(1, 1), stride=(1, 1))\n", " (encoder_1d_feature_pos_embed): Embedding(2, 512)\n", " (encoder_cam_feat_pos_embed): ACTSinusoidalPositionEmbedding2d()\n", " (decoder_pos_embed): Embedding(10, 512)\n", " (action_head): Linear(in_features=512, out_features=7, bias=True)\n", " )\n", ")" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "dataset_metadata = LeRobotDatasetMetadata(\"datawhale_eai_pnp\", root='./demo_data')\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", "input_features.pop(\"observation.wrist_image\")\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 = ACTConfig(input_features=input_features, output_features=output_features, chunk_size= 10, n_action_steps=1, temporal_ensemble_coeff = 0.9)\n", "delta_timestamps = resolve_delta_timestamps(cfg, dataset_metadata)\n", "# We can now instantiate our policy with this config and the dataset stats.\n", "policy = ACTPolicy.from_pretrained('./ckpt/act_y', config = cfg, dataset_stats=dataset_metadata.stats)\n", "policy.to(device)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 加载环境\n" ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "scrolled": true }, "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": [ "from mujoco_env.y_env import SimpleEnv\n", "xml_path = './asset/example_scene_y.xml'\n", "PnPEnv = SimpleEnv(xml_path, action_type='joint_angle')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Roll-Out 策略\n" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "DONE INITIALIZATION\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/tmp/ipykernel_797907/978506407.py:30: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at /pytorch/torch/csrc/utils/tensor_new.cpp:254.)\n", " 'observation.state': torch.tensor([state]).to(device),\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "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 = torchvision.transforms.ToTensor()\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(seed=0)\n", " step = 0\n", " save_image = False\n", " # Get the current state of the environment\n", " state = PnPEnv.get_ee_pose()\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': ['Put mug cup on the plate'],\n", " 'timestamp': torch.tensor([step/20]).to(device)\n", " }\n", " # Select an action\n", " action = policy.select_action(data)\n", " action = action[0].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": null, "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": null, "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": 4 }