restructure code
This commit is contained in:
@@ -1,461 +0,0 @@
|
|||||||
import logging
|
|
||||||
import time
|
|
||||||
from dataclasses import asdict
|
|
||||||
from pprint import pformat
|
|
||||||
from pprint import pprint
|
|
||||||
|
|
||||||
# from safetensors.torch import load_file, save_file
|
|
||||||
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
|
|
||||||
from lerobot.common.policies.factory import make_policy
|
|
||||||
from lerobot.common.robot_devices.control_configs import (
|
|
||||||
CalibrateControlConfig,
|
|
||||||
ControlPipelineConfig,
|
|
||||||
RecordControlConfig,
|
|
||||||
RemoteRobotConfig,
|
|
||||||
ReplayControlConfig,
|
|
||||||
TeleoperateControlConfig,
|
|
||||||
)
|
|
||||||
from lerobot.common.robot_devices.control_utils import (
|
|
||||||
# init_keyboard_listener,
|
|
||||||
record_episode,
|
|
||||||
stop_recording,
|
|
||||||
is_headless
|
|
||||||
)
|
|
||||||
from lerobot.common.robot_devices.robots.utils import Robot, make_robot_from_config
|
|
||||||
from lerobot.common.robot_devices.utils import busy_wait, safe_disconnect
|
|
||||||
from lerobot.common.utils.utils import has_method, init_logging, log_say
|
|
||||||
from lerobot.common.utils.utils import get_safe_torch_device
|
|
||||||
from contextlib import nullcontext
|
|
||||||
from copy import copy
|
|
||||||
import torch
|
|
||||||
import rospy
|
|
||||||
import cv2
|
|
||||||
from lerobot.configs import parser
|
|
||||||
from common.agilex_robot import AgilexRobot
|
|
||||||
from common.rosrobot_factory import RobotFactory
|
|
||||||
|
|
||||||
|
|
||||||
########################################################################################
|
|
||||||
# Control modes
|
|
||||||
########################################################################################
|
|
||||||
|
|
||||||
|
|
||||||
def predict_action(observation, policy, device, use_amp):
|
|
||||||
observation = copy(observation)
|
|
||||||
with (
|
|
||||||
torch.inference_mode(),
|
|
||||||
torch.autocast(device_type=device.type) if device.type == "cuda" and use_amp else nullcontext(),
|
|
||||||
):
|
|
||||||
# Convert to pytorch format: channel first and float32 in [0,1] with batch dimension
|
|
||||||
for name in observation:
|
|
||||||
if "image" in name:
|
|
||||||
observation[name] = observation[name].type(torch.float32) / 255
|
|
||||||
observation[name] = observation[name].permute(2, 0, 1).contiguous()
|
|
||||||
observation[name] = observation[name].unsqueeze(0)
|
|
||||||
observation[name] = observation[name].to(device)
|
|
||||||
|
|
||||||
# Compute the next action with the policy
|
|
||||||
# based on the current observation
|
|
||||||
action = policy.select_action(observation)
|
|
||||||
|
|
||||||
# Remove batch dimension
|
|
||||||
action = action.squeeze(0)
|
|
||||||
|
|
||||||
# Move to cpu, if not already the case
|
|
||||||
action = action.to("cpu")
|
|
||||||
|
|
||||||
return action
|
|
||||||
|
|
||||||
def control_loop(
|
|
||||||
robot,
|
|
||||||
control_time_s=None,
|
|
||||||
teleoperate=False,
|
|
||||||
display_cameras=False,
|
|
||||||
dataset: LeRobotDataset | None = None,
|
|
||||||
events=None,
|
|
||||||
policy = None,
|
|
||||||
fps: int | None = None,
|
|
||||||
single_task: str | None = None,
|
|
||||||
):
|
|
||||||
# TODO(rcadene): Add option to record logs
|
|
||||||
# if not robot.is_connected:
|
|
||||||
# robot.connect()
|
|
||||||
|
|
||||||
if events is None:
|
|
||||||
events = {"exit_early": False}
|
|
||||||
|
|
||||||
if control_time_s is None:
|
|
||||||
control_time_s = float("inf")
|
|
||||||
|
|
||||||
if dataset is not None and single_task is None:
|
|
||||||
raise ValueError("You need to provide a task as argument in `single_task`.")
|
|
||||||
|
|
||||||
if dataset is not None and fps is not None and dataset.fps != fps:
|
|
||||||
raise ValueError(f"The dataset fps should be equal to requested fps ({dataset['fps']} != {fps}).")
|
|
||||||
|
|
||||||
timestamp = 0
|
|
||||||
start_episode_t = time.perf_counter()
|
|
||||||
rate = rospy.Rate(fps)
|
|
||||||
print_flag = True
|
|
||||||
while timestamp < control_time_s and not rospy.is_shutdown():
|
|
||||||
# print(timestamp < control_time_s)
|
|
||||||
# print(rospy.is_shutdown())
|
|
||||||
start_loop_t = time.perf_counter()
|
|
||||||
|
|
||||||
if teleoperate:
|
|
||||||
observation, action = robot.teleop_step()
|
|
||||||
if observation is None or action is None:
|
|
||||||
if print_flag:
|
|
||||||
print("sync data fail, retrying...\n")
|
|
||||||
print_flag = False
|
|
||||||
rate.sleep()
|
|
||||||
continue
|
|
||||||
else:
|
|
||||||
# pass
|
|
||||||
observation = robot.capture_observation()
|
|
||||||
if policy is not None:
|
|
||||||
pred_action = predict_action(
|
|
||||||
observation, policy, get_safe_torch_device(policy.config.device), policy.config.use_amp
|
|
||||||
)
|
|
||||||
# Action can eventually be clipped using `max_relative_target`,
|
|
||||||
# so action actually sent is saved in the dataset.
|
|
||||||
action = robot.send_action(pred_action)
|
|
||||||
action = {"action": action}
|
|
||||||
|
|
||||||
if dataset is not None:
|
|
||||||
frame = {**observation, **action, "task": single_task}
|
|
||||||
dataset.add_frame(frame)
|
|
||||||
|
|
||||||
# if display_cameras and not is_headless():
|
|
||||||
# image_keys = [key for key in observation if "image" in key]
|
|
||||||
# for key in image_keys:
|
|
||||||
# if "depth" in key:
|
|
||||||
# pass
|
|
||||||
# else:
|
|
||||||
# cv2.imshow(key, cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR))
|
|
||||||
|
|
||||||
# print(1)
|
|
||||||
# cv2.waitKey(1)
|
|
||||||
|
|
||||||
if display_cameras and not is_headless():
|
|
||||||
image_keys = [key for key in observation if "image" in key]
|
|
||||||
|
|
||||||
# 获取屏幕分辨率（假设屏幕分辨率为 1920x1080，可以根据实际情况调整）
|
|
||||||
screen_width = 1920
|
|
||||||
screen_height = 1080
|
|
||||||
|
|
||||||
# 计算窗口的排列方式
|
|
||||||
num_images = len(image_keys)
|
|
||||||
max_columns = int(screen_width / 640) # 假设每个窗口宽度为 640
|
|
||||||
rows = (num_images + max_columns - 1) // max_columns # 计算需要的行数
|
|
||||||
columns = min(num_images, max_columns) # 实际使用的列数
|
|
||||||
|
|
||||||
# 遍历所有图像键并显示
|
|
||||||
for idx, key in enumerate(image_keys):
|
|
||||||
if "depth" in key:
|
|
||||||
continue # 跳过深度图像
|
|
||||||
|
|
||||||
# 将图像从 RGB 转换为 BGR 格式
|
|
||||||
image = cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR)
|
|
||||||
|
|
||||||
# 创建窗口
|
|
||||||
cv2.imshow(key, image)
|
|
||||||
|
|
||||||
# 计算窗口位置
|
|
||||||
window_width = 640
|
|
||||||
window_height = 480
|
|
||||||
row = idx // max_columns
|
|
||||||
col = idx % max_columns
|
|
||||||
x_position = col * window_width
|
|
||||||
y_position = row * window_height
|
|
||||||
|
|
||||||
# 移动窗口到指定位置
|
|
||||||
cv2.moveWindow(key, x_position, y_position)
|
|
||||||
|
|
||||||
# 等待 1 毫秒以处理事件
|
|
||||||
cv2.waitKey(1)
|
|
||||||
|
|
||||||
if fps is not None:
|
|
||||||
dt_s = time.perf_counter() - start_loop_t
|
|
||||||
busy_wait(1 / fps - dt_s)
|
|
||||||
|
|
||||||
dt_s = time.perf_counter() - start_loop_t
|
|
||||||
# log_control_info(robot, dt_s, fps=fps)
|
|
||||||
|
|
||||||
timestamp = time.perf_counter() - start_episode_t
|
|
||||||
if events["exit_early"]:
|
|
||||||
events["exit_early"] = False
|
|
||||||
break
|
|
||||||
|
|
||||||
|
|
||||||
def init_keyboard_listener():
|
|
||||||
# Allow to exit early while recording an episode or resetting the environment,
|
|
||||||
# by tapping the right arrow key '->'. This might require a sudo permission
|
|
||||||
# to allow your terminal to monitor keyboard events.
|
|
||||||
events = {}
|
|
||||||
events["exit_early"] = False
|
|
||||||
events["record_start"] = False
|
|
||||||
events["rerecord_episode"] = False
|
|
||||||
events["stop_recording"] = False
|
|
||||||
|
|
||||||
if is_headless():
|
|
||||||
logging.warning(
|
|
||||||
"Headless environment detected. On-screen cameras display and keyboard inputs will not be available."
|
|
||||||
)
|
|
||||||
listener = None
|
|
||||||
return listener, events
|
|
||||||
|
|
||||||
# Only import pynput if not in a headless environment
|
|
||||||
from pynput import keyboard
|
|
||||||
|
|
||||||
def on_press(key):
|
|
||||||
try:
|
|
||||||
if key == keyboard.Key.right:
|
|
||||||
print("Right arrow key pressed. Exiting loop...")
|
|
||||||
events["exit_early"] = True
|
|
||||||
events["record_start"] = False
|
|
||||||
elif key == keyboard.Key.left:
|
|
||||||
print("Left arrow key pressed. Exiting loop and rerecord the last episode...")
|
|
||||||
events["rerecord_episode"] = True
|
|
||||||
events["exit_early"] = True
|
|
||||||
elif key == keyboard.Key.esc:
|
|
||||||
print("Escape key pressed. Stopping data recording...")
|
|
||||||
events["stop_recording"] = True
|
|
||||||
events["exit_early"] = True
|
|
||||||
elif key == keyboard.Key.up:
|
|
||||||
print("Up arrow pressed. Start data recording...")
|
|
||||||
events["record_start"] = True
|
|
||||||
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
print(f"Error handling key press: {e}")
|
|
||||||
|
|
||||||
listener = keyboard.Listener(on_press=on_press)
|
|
||||||
listener.start()
|
|
||||||
|
|
||||||
return listener, events
|
|
||||||
|
|
||||||
|
|
||||||
def stop_recording(robot, listener, display_cameras):
|
|
||||||
|
|
||||||
if not is_headless():
|
|
||||||
if listener is not None:
|
|
||||||
listener.stop()
|
|
||||||
|
|
||||||
if display_cameras:
|
|
||||||
cv2.destroyAllWindows()
|
|
||||||
|
|
||||||
|
|
||||||
def record_episode(
|
|
||||||
robot,
|
|
||||||
dataset,
|
|
||||||
events,
|
|
||||||
episode_time_s,
|
|
||||||
display_cameras,
|
|
||||||
policy,
|
|
||||||
fps,
|
|
||||||
single_task,
|
|
||||||
):
|
|
||||||
control_loop(
|
|
||||||
robot=robot,
|
|
||||||
control_time_s=episode_time_s,
|
|
||||||
display_cameras=display_cameras,
|
|
||||||
dataset=dataset,
|
|
||||||
events=events,
|
|
||||||
policy=policy,
|
|
||||||
fps=fps,
|
|
||||||
teleoperate=policy is None,
|
|
||||||
single_task=single_task,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def record(
|
|
||||||
robot,
|
|
||||||
cfg
|
|
||||||
) -> LeRobotDataset:
|
|
||||||
# TODO(rcadene): Add option to record logs
|
|
||||||
if cfg.resume:
|
|
||||||
dataset = LeRobotDataset(
|
|
||||||
cfg.repo_id,
|
|
||||||
root=cfg.root,
|
|
||||||
)
|
|
||||||
if len(robot.cameras) > 0:
|
|
||||||
dataset.start_image_writer(
|
|
||||||
num_processes=cfg.num_image_writer_processes,
|
|
||||||
num_threads=cfg.num_image_writer_threads_per_camera * len(robot.cameras),
|
|
||||||
)
|
|
||||||
# sanity_check_dataset_robot_compatibility(dataset, robot, cfg.fps, cfg.video)
|
|
||||||
else:
|
|
||||||
# Create empty dataset or load existing saved episodes
|
|
||||||
# sanity_check_dataset_name(cfg.repo_id, cfg.policy)
|
|
||||||
dataset = LeRobotDataset.create(
|
|
||||||
cfg.repo_id,
|
|
||||||
cfg.fps,
|
|
||||||
root=cfg.root,
|
|
||||||
robot=None,
|
|
||||||
features=robot.features,
|
|
||||||
use_videos=cfg.video,
|
|
||||||
image_writer_processes=cfg.num_image_writer_processes,
|
|
||||||
image_writer_threads=cfg.num_image_writer_threads_per_camera * len(robot.cameras),
|
|
||||||
)
|
|
||||||
|
|
||||||
# Load pretrained policy
|
|
||||||
policy = None if cfg.policy is None else make_policy(cfg.policy, ds_meta=dataset.meta)
|
|
||||||
# policy = None
|
|
||||||
|
|
||||||
# if not robot.is_connected:
|
|
||||||
# robot.connect()
|
|
||||||
|
|
||||||
listener, events = init_keyboard_listener()
|
|
||||||
|
|
||||||
# Execute a few seconds without recording to:
|
|
||||||
# 1. teleoperate the robot to move it in starting position if no policy provided,
|
|
||||||
# 2. give times to the robot devices to connect and start synchronizing,
|
|
||||||
# 3. place the cameras windows on screen
|
|
||||||
enable_teleoperation = policy is None
|
|
||||||
log_say("Warmup record", cfg.play_sounds)
|
|
||||||
print()
|
|
||||||
print(f"开始记录轨迹，共需要记录{cfg.num_episodes}条\n每条轨迹的最长时间为{cfg.episode_time_s}frame\n按右方向键代表当前轨迹结束录制\n按上方面键代表当前轨迹开始录制\n按左方向键代表当前轨迹重新录制\n按ESC方向键代表退出轨迹录制\n")
|
|
||||||
# warmup_record(robot, events, enable_teleoperation, cfg.warmup_time_s, cfg.display_cameras, cfg.fps)
|
|
||||||
|
|
||||||
# if has_method(robot, "teleop_safety_stop"):
|
|
||||||
# robot.teleop_safety_stop()
|
|
||||||
|
|
||||||
recorded_episodes = 0
|
|
||||||
while True:
|
|
||||||
if recorded_episodes >= cfg.num_episodes:
|
|
||||||
break
|
|
||||||
|
|
||||||
# if events["record_start"]:
|
|
||||||
log_say(f"Recording episode {dataset.num_episodes}", cfg.play_sounds)
|
|
||||||
pprint(f"Recording episode {dataset.num_episodes}, total episodes is {cfg.num_episodes}")
|
|
||||||
record_episode(
|
|
||||||
robot=robot,
|
|
||||||
dataset=dataset,
|
|
||||||
events=events,
|
|
||||||
episode_time_s=cfg.episode_time_s,
|
|
||||||
display_cameras=cfg.display_cameras,
|
|
||||||
policy=policy,
|
|
||||||
fps=cfg.fps,
|
|
||||||
single_task=cfg.single_task,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Execute a few seconds without recording to give time to manually reset the environment
|
|
||||||
# Current code logic doesn't allow to teleoperate during this time.
|
|
||||||
# TODO(rcadene): add an option to enable teleoperation during reset
|
|
||||||
# Skip reset for the last episode to be recorded
|
|
||||||
if not events["stop_recording"] and (
|
|
||||||
(recorded_episodes < cfg.num_episodes - 1) or events["rerecord_episode"]
|
|
||||||
):
|
|
||||||
log_say("Reset the environment", cfg.play_sounds)
|
|
||||||
pprint("Reset the environment, stop recording")
|
|
||||||
# reset_environment(robot, events, cfg.reset_time_s, cfg.fps)
|
|
||||||
|
|
||||||
if events["rerecord_episode"]:
|
|
||||||
log_say("Re-record episode", cfg.play_sounds)
|
|
||||||
pprint("Re-record episode")
|
|
||||||
events["rerecord_episode"] = False
|
|
||||||
events["exit_early"] = False
|
|
||||||
dataset.clear_episode_buffer()
|
|
||||||
continue
|
|
||||||
|
|
||||||
dataset.save_episode()
|
|
||||||
recorded_episodes += 1
|
|
||||||
|
|
||||||
if events["stop_recording"]:
|
|
||||||
break
|
|
||||||
|
|
||||||
log_say("Stop recording", cfg.play_sounds, blocking=True)
|
|
||||||
stop_recording(robot, listener, cfg.display_cameras)
|
|
||||||
|
|
||||||
if cfg.push_to_hub:
|
|
||||||
dataset.push_to_hub(tags=cfg.tags, private=cfg.private)
|
|
||||||
|
|
||||||
log_say("Exiting", cfg.play_sounds)
|
|
||||||
return dataset
|
|
||||||
|
|
||||||
|
|
||||||
def replay(
|
|
||||||
robot: AgilexRobot,
|
|
||||||
cfg,
|
|
||||||
):
|
|
||||||
# TODO(rcadene, aliberts): refactor with control_loop, once `dataset` is an instance of LeRobotDataset
|
|
||||||
# TODO(rcadene): Add option to record logs
|
|
||||||
|
|
||||||
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root, episodes=[cfg.episode])
|
|
||||||
actions = dataset.hf_dataset.select_columns("action")
|
|
||||||
|
|
||||||
# if not robot.is_connected:
|
|
||||||
# robot.connect()
|
|
||||||
|
|
||||||
log_say("Replaying episode", cfg.play_sounds, blocking=True)
|
|
||||||
for idx in range(dataset.num_frames):
|
|
||||||
start_episode_t = time.perf_counter()
|
|
||||||
|
|
||||||
action = actions[idx]["action"]
|
|
||||||
robot.send_action(action)
|
|
||||||
|
|
||||||
dt_s = time.perf_counter() - start_episode_t
|
|
||||||
busy_wait(1 / cfg.fps - dt_s)
|
|
||||||
|
|
||||||
dt_s = time.perf_counter() - start_episode_t
|
|
||||||
# log_control_info(robot, dt_s, fps=cfg.fps)
|
|
||||||
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
def get_arguments():
|
|
||||||
parser = argparse.ArgumentParser()
|
|
||||||
args = parser.parse_args()
|
|
||||||
args.fps = 30
|
|
||||||
args.resume = False
|
|
||||||
args.repo_id = "move_the_bottle_from_the_right_to_the_scale_right"
|
|
||||||
args.root = "./data5"
|
|
||||||
args.episode = 0 # replay episode
|
|
||||||
args.num_image_writer_processes = 0
|
|
||||||
args.num_image_writer_threads_per_camera = 4
|
|
||||||
args.video = True
|
|
||||||
args.num_episodes = 100
|
|
||||||
args.episode_time_s = 30000
|
|
||||||
args.play_sounds = False
|
|
||||||
args.display_cameras = True
|
|
||||||
args.single_task = "move the bottle from the right to the scale right"
|
|
||||||
args.use_depth_image = False
|
|
||||||
args.use_base = False
|
|
||||||
args.push_to_hub = False
|
|
||||||
args.policy = None
|
|
||||||
# args.teleoprate = True
|
|
||||||
args.control_type = "record"
|
|
||||||
# args.control_type = "replay"
|
|
||||||
return args
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
# @parser.wrap()
|
|
||||||
def control_robot(cfg):
|
|
||||||
# 使用工厂模式创建机器人实例
|
|
||||||
robot = RobotFactory.create(config_file="/home/ubuntu/LYT/lerobot_aloha/lerobot_aloha/configs/agilex.yaml", args=cfg)
|
|
||||||
|
|
||||||
if cfg.control_type == "record":
|
|
||||||
record(robot, cfg)
|
|
||||||
elif cfg.control_type == "replay":
|
|
||||||
replay(robot, cfg)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
cfg = get_arguments()
|
|
||||||
control_robot(cfg)
|
|
||||||
# control_robot()
|
|
||||||
# 使用工厂模式创建机器人实例
|
|
||||||
# robot = RobotFactory.create(config_file="/home/ubuntu/LYT/aloha_lerobot/collect_data/agilex.yaml", args=cfg)
|
|
||||||
# print(robot.features.items())
|
|
||||||
# print([key for key, ft in robot.features.items() if ft["dtype"] == "video"])
|
|
||||||
# record(robot, cfg)
|
|
||||||
# capture = robot.capture_observation()
|
|
||||||
# import torch
|
|
||||||
# torch.save(capture, "test.pt")
|
|
||||||
# action = torch.tensor([[ 0.0277, 0.0167, 0.0142, -0.1628, 0.1473, -0.0296, 0.0238, -0.1094,
|
|
||||||
# 0.0109, 0.0139, -0.1591, -0.1490, -0.1650, -0.0980]],
|
|
||||||
# device='cpu')
|
|
||||||
# robot.send_action(action.squeeze(0))
|
|
||||||
# print()
|
|
||||||
12
lerobot_aloha/common/utils/__init__.py
Normal file
12
lerobot_aloha/common/utils/__init__.py
Normal file
@@ -0,0 +1,12 @@
|
|||||||
|
# Import utility functions for easy access
|
||||||
|
from .control_utils import (
|
||||||
|
predict_action,
|
||||||
|
control_loop,
|
||||||
|
init_keyboard_listener,
|
||||||
|
stop_recording,
|
||||||
|
record_episode,
|
||||||
|
is_headless,
|
||||||
|
busy_wait
|
||||||
|
)
|
||||||
|
from .data_utils import record
|
||||||
|
from .replay_utils import replay
|
||||||
BIN
lerobot_aloha/common/utils/__pycache__/__init__.cpython-310.pyc
Normal file
BIN
lerobot_aloha/common/utils/__pycache__/__init__.cpython-310.pyc
Normal file
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
303
lerobot_aloha/common/utils/control_utils.py
Normal file
303
lerobot_aloha/common/utils/control_utils.py
Normal file
@@ -0,0 +1,303 @@
|
|||||||
|
import logging
|
||||||
|
import time
|
||||||
|
import torch
|
||||||
|
import rospy
|
||||||
|
import cv2
|
||||||
|
from contextlib import nullcontext
|
||||||
|
from copy import copy
|
||||||
|
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
|
||||||
|
from lerobot.common.utils.utils import get_safe_torch_device
|
||||||
|
|
||||||
|
|
||||||
|
def is_headless():
|
||||||
|
"""
|
||||||
|
Check if the environment is headless (no display available).
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
bool: True if the environment is headless, False otherwise.
|
||||||
|
"""
|
||||||
|
try:
|
||||||
|
import tkinter as tk
|
||||||
|
root = tk.Tk()
|
||||||
|
root.withdraw()
|
||||||
|
root.update()
|
||||||
|
root.destroy()
|
||||||
|
return False
|
||||||
|
except:
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def predict_action(observation, policy, device, use_amp):
|
||||||
|
"""
|
||||||
|
Predict action based on observation using the policy.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
observation: Current observation
|
||||||
|
policy: Policy model
|
||||||
|
device: Torch device
|
||||||
|
use_amp: Whether to use automatic mixed precision
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
torch.Tensor: Predicted action
|
||||||
|
"""
|
||||||
|
observation = copy(observation)
|
||||||
|
with (
|
||||||
|
torch.inference_mode(),
|
||||||
|
torch.autocast(device_type=device.type) if device.type == "cuda" and use_amp else nullcontext(),
|
||||||
|
):
|
||||||
|
# Convert to pytorch format: channel first and float32 in [0,1] with batch dimension
|
||||||
|
for name in observation:
|
||||||
|
if "image" in name:
|
||||||
|
observation[name] = observation[name].type(torch.float32) / 255
|
||||||
|
observation[name] = observation[name].permute(2, 0, 1).contiguous()
|
||||||
|
observation[name] = observation[name].unsqueeze(0)
|
||||||
|
observation[name] = observation[name].to(device)
|
||||||
|
|
||||||
|
# Compute the next action with the policy
|
||||||
|
# based on the current observation
|
||||||
|
action = policy.select_action(observation)
|
||||||
|
|
||||||
|
# Remove batch dimension
|
||||||
|
action = action.squeeze(0)
|
||||||
|
|
||||||
|
# Move to cpu, if not already the case
|
||||||
|
action = action.to("cpu")
|
||||||
|
|
||||||
|
return action
|
||||||
|
|
||||||
|
|
||||||
|
def control_loop(
|
||||||
|
robot,
|
||||||
|
control_time_s=None,
|
||||||
|
teleoperate=False,
|
||||||
|
display_cameras=False,
|
||||||
|
dataset: LeRobotDataset | None = None,
|
||||||
|
events=None,
|
||||||
|
policy = None,
|
||||||
|
fps: int | None = None,
|
||||||
|
single_task: str | None = None,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Main control loop for robot operation.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
robot: Robot instance
|
||||||
|
control_time_s: Control time in seconds
|
||||||
|
teleoperate: Whether to use teleoperation
|
||||||
|
display_cameras: Whether to display camera feeds
|
||||||
|
dataset: Dataset for recording
|
||||||
|
events: Event dictionary
|
||||||
|
policy: Policy model
|
||||||
|
fps: Frames per second
|
||||||
|
single_task: Task name
|
||||||
|
"""
|
||||||
|
if events is None:
|
||||||
|
events = {"exit_early": False}
|
||||||
|
|
||||||
|
if control_time_s is None:
|
||||||
|
control_time_s = float("inf")
|
||||||
|
|
||||||
|
if dataset is not None and single_task is None:
|
||||||
|
raise ValueError("You need to provide a task as argument in `single_task`.")
|
||||||
|
|
||||||
|
if dataset is not None and fps is not None and dataset.fps != fps:
|
||||||
|
raise ValueError(f"The dataset fps should be equal to requested fps ({dataset['fps']} != {fps}).")
|
||||||
|
|
||||||
|
timestamp = 0
|
||||||
|
start_episode_t = time.perf_counter()
|
||||||
|
rate = rospy.Rate(fps)
|
||||||
|
print_flag = True
|
||||||
|
while timestamp < control_time_s and not rospy.is_shutdown():
|
||||||
|
start_loop_t = time.perf_counter()
|
||||||
|
|
||||||
|
if teleoperate:
|
||||||
|
observation, action = robot.teleop_step()
|
||||||
|
if observation is None or action is None:
|
||||||
|
if print_flag:
|
||||||
|
print("sync data fail, retrying...\n")
|
||||||
|
print_flag = False
|
||||||
|
rate.sleep()
|
||||||
|
continue
|
||||||
|
else:
|
||||||
|
observation = robot.capture_observation()
|
||||||
|
if policy is not None:
|
||||||
|
pred_action = predict_action(
|
||||||
|
observation, policy, get_safe_torch_device(policy.config.device), policy.config.use_amp
|
||||||
|
)
|
||||||
|
# Action can eventually be clipped using `max_relative_target`,
|
||||||
|
# so action actually sent is saved in the dataset.
|
||||||
|
action = robot.send_action(pred_action)
|
||||||
|
action = {"action": action}
|
||||||
|
|
||||||
|
if dataset is not None:
|
||||||
|
frame = {**observation, **action, "task": single_task}
|
||||||
|
dataset.add_frame(frame)
|
||||||
|
|
||||||
|
if display_cameras and not is_headless():
|
||||||
|
image_keys = [key for key in observation if "image" in key]
|
||||||
|
|
||||||
|
# 获取屏幕分辨率（假设屏幕分辨率为 1920x1080，可以根据实际情况调整）
|
||||||
|
screen_width = 1920
|
||||||
|
screen_height = 1080
|
||||||
|
|
||||||
|
# 计算窗口的排列方式
|
||||||
|
num_images = len(image_keys)
|
||||||
|
max_columns = int(screen_width / 640) # 假设每个窗口宽度为 640
|
||||||
|
rows = (num_images + max_columns - 1) // max_columns # 计算需要的行数
|
||||||
|
columns = min(num_images, max_columns) # 实际使用的列数
|
||||||
|
|
||||||
|
# 遍历所有图像键并显示
|
||||||
|
for idx, key in enumerate(image_keys):
|
||||||
|
if "depth" in key:
|
||||||
|
continue # 跳过深度图像
|
||||||
|
|
||||||
|
# 将图像从 RGB 转换为 BGR 格式
|
||||||
|
image = cv2.cvtColor(observation[key].numpy(), cv2.COLOR_RGB2BGR)
|
||||||
|
|
||||||
|
# 创建窗口
|
||||||
|
cv2.imshow(key, image)
|
||||||
|
|
||||||
|
# 计算窗口位置
|
||||||
|
window_width = 640
|
||||||
|
window_height = 480
|
||||||
|
row = idx // max_columns
|
||||||
|
col = idx % max_columns
|
||||||
|
x_position = col * window_width
|
||||||
|
y_position = row * window_height
|
||||||
|
|
||||||
|
# 移动窗口到指定位置
|
||||||
|
cv2.moveWindow(key, x_position, y_position)
|
||||||
|
|
||||||
|
# 等待 1 毫秒以处理事件
|
||||||
|
cv2.waitKey(1)
|
||||||
|
|
||||||
|
if fps is not None:
|
||||||
|
dt_s = time.perf_counter() - start_loop_t
|
||||||
|
busy_wait(1 / fps - dt_s)
|
||||||
|
|
||||||
|
dt_s = time.perf_counter() - start_loop_t
|
||||||
|
|
||||||
|
timestamp = time.perf_counter() - start_episode_t
|
||||||
|
if events["exit_early"]:
|
||||||
|
events["exit_early"] = False
|
||||||
|
break
|
||||||
|
|
||||||
|
|
||||||
|
def init_keyboard_listener():
|
||||||
|
"""
|
||||||
|
Initialize keyboard listener for control events.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
tuple: (listener, events) - Keyboard listener and events dictionary
|
||||||
|
"""
|
||||||
|
# Allow to exit early while recording an episode or resetting the environment,
|
||||||
|
# by tapping the right arrow key '->'. This might require a sudo permission
|
||||||
|
# to allow your terminal to monitor keyboard events.
|
||||||
|
events = {}
|
||||||
|
events["exit_early"] = False
|
||||||
|
events["record_start"] = False
|
||||||
|
events["rerecord_episode"] = False
|
||||||
|
events["stop_recording"] = False
|
||||||
|
|
||||||
|
if is_headless():
|
||||||
|
logging.warning(
|
||||||
|
"Headless environment detected. On-screen cameras display and keyboard inputs will not be available."
|
||||||
|
)
|
||||||
|
listener = None
|
||||||
|
return listener, events
|
||||||
|
|
||||||
|
# Only import pynput if not in a headless environment
|
||||||
|
from pynput import keyboard
|
||||||
|
|
||||||
|
def on_press(key):
|
||||||
|
try:
|
||||||
|
if key == keyboard.Key.right:
|
||||||
|
print("Right arrow key pressed. Exiting loop...")
|
||||||
|
events["exit_early"] = True
|
||||||
|
events["record_start"] = False
|
||||||
|
elif key == keyboard.Key.left:
|
||||||
|
print("Left arrow key pressed. Exiting loop and rerecord the last episode...")
|
||||||
|
events["rerecord_episode"] = True
|
||||||
|
events["exit_early"] = True
|
||||||
|
elif key == keyboard.Key.esc:
|
||||||
|
print("Escape key pressed. Stopping data recording...")
|
||||||
|
events["stop_recording"] = True
|
||||||
|
events["exit_early"] = True
|
||||||
|
elif key == keyboard.Key.up:
|
||||||
|
print("Up arrow pressed. Start data recording...")
|
||||||
|
events["record_start"] = True
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Error handling key press: {e}")
|
||||||
|
|
||||||
|
listener = keyboard.Listener(on_press=on_press)
|
||||||
|
listener.start()
|
||||||
|
|
||||||
|
return listener, events
|
||||||
|
|
||||||
|
|
||||||
|
def stop_recording(robot, listener, display_cameras):
|
||||||
|
"""
|
||||||
|
Stop recording and clean up resources.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
robot: Robot instance
|
||||||
|
listener: Keyboard listener
|
||||||
|
display_cameras: Whether cameras are being displayed
|
||||||
|
"""
|
||||||
|
if not is_headless():
|
||||||
|
if listener is not None:
|
||||||
|
listener.stop()
|
||||||
|
|
||||||
|
if display_cameras:
|
||||||
|
cv2.destroyAllWindows()
|
||||||
|
|
||||||
|
|
||||||
|
def record_episode(
|
||||||
|
robot,
|
||||||
|
dataset,
|
||||||
|
events,
|
||||||
|
episode_time_s,
|
||||||
|
display_cameras,
|
||||||
|
policy,
|
||||||
|
fps,
|
||||||
|
single_task,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Record a single episode.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
robot: Robot instance
|
||||||
|
dataset: Dataset for recording
|
||||||
|
events: Event dictionary
|
||||||
|
episode_time_s: Episode time in seconds
|
||||||
|
display_cameras: Whether to display camera feeds
|
||||||
|
policy: Policy model
|
||||||
|
fps: Frames per second
|
||||||
|
single_task: Task name
|
||||||
|
"""
|
||||||
|
control_loop(
|
||||||
|
robot=robot,
|
||||||
|
control_time_s=episode_time_s,
|
||||||
|
display_cameras=display_cameras,
|
||||||
|
dataset=dataset,
|
||||||
|
events=events,
|
||||||
|
policy=policy,
|
||||||
|
fps=fps,
|
||||||
|
teleoperate=policy is None,
|
||||||
|
single_task=single_task,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def busy_wait(seconds):
|
||||||
|
"""
|
||||||
|
Busy wait for a specified number of seconds.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
seconds: Number of seconds to wait
|
||||||
|
"""
|
||||||
|
if seconds <= 0:
|
||||||
|
return
|
||||||
|
start_time = time.perf_counter()
|
||||||
|
while time.perf_counter() - start_time < seconds:
|
||||||
|
pass
|
||||||
105
lerobot_aloha/common/utils/data_utils.py
Normal file
105
lerobot_aloha/common/utils/data_utils.py
Normal file
@@ -0,0 +1,105 @@
|
|||||||
|
import logging
|
||||||
|
import time
|
||||||
|
from pprint import pprint
|
||||||
|
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
|
||||||
|
from lerobot.common.policies.factory import make_policy
|
||||||
|
from lerobot.common.utils.utils import log_say, has_method
|
||||||
|
from common.utils.control_utils import init_keyboard_listener, stop_recording, record_episode
|
||||||
|
|
||||||
|
|
||||||
|
def record(
|
||||||
|
robot,
|
||||||
|
cfg
|
||||||
|
) -> LeRobotDataset:
|
||||||
|
"""
|
||||||
|
Record robot data according to configuration.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
robot: Robot instance
|
||||||
|
cfg: Configuration object
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
LeRobotDataset: Dataset with recorded episodes
|
||||||
|
"""
|
||||||
|
# Initialize or load dataset
|
||||||
|
if cfg.resume:
|
||||||
|
dataset = LeRobotDataset(
|
||||||
|
cfg.repo_id,
|
||||||
|
root=cfg.root,
|
||||||
|
)
|
||||||
|
if len(robot.cameras) > 0:
|
||||||
|
dataset.start_image_writer(
|
||||||
|
num_processes=cfg.num_image_writer_processes,
|
||||||
|
num_threads=cfg.num_image_writer_threads_per_camera * len(robot.cameras),
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
# Create empty dataset or load existing saved episodes
|
||||||
|
dataset = LeRobotDataset.create(
|
||||||
|
cfg.repo_id,
|
||||||
|
cfg.fps,
|
||||||
|
root=cfg.root,
|
||||||
|
robot=None,
|
||||||
|
features=robot.features,
|
||||||
|
use_videos=cfg.video,
|
||||||
|
image_writer_processes=cfg.num_image_writer_processes,
|
||||||
|
image_writer_threads=cfg.num_image_writer_threads_per_camera * len(robot.cameras),
|
||||||
|
)
|
||||||
|
|
||||||
|
# Load pretrained policy
|
||||||
|
policy = None if cfg.policy is None else make_policy(cfg.policy, ds_meta=dataset.meta)
|
||||||
|
|
||||||
|
# Initialize keyboard listener
|
||||||
|
listener, events = init_keyboard_listener()
|
||||||
|
|
||||||
|
# Print recording instructions
|
||||||
|
print()
|
||||||
|
print(f"开始记录轨迹，共需要记录{cfg.num_episodes}条\n每条轨迹的最长时间为{cfg.episode_time_s}frame\n按右方向键代表当前轨迹结束录制\n按上方面键代表当前轨迹开始录制\n按左方向键代表当前轨迹重新录制\n按ESC方向键代表退出轨迹录制\n")
|
||||||
|
|
||||||
|
# Record episodes
|
||||||
|
recorded_episodes = 0
|
||||||
|
while True:
|
||||||
|
if recorded_episodes >= cfg.num_episodes:
|
||||||
|
break
|
||||||
|
|
||||||
|
log_say(f"Recording episode {dataset.num_episodes}", cfg.play_sounds)
|
||||||
|
pprint(f"Recording episode {dataset.num_episodes}, total episodes is {cfg.num_episodes}")
|
||||||
|
record_episode(
|
||||||
|
robot=robot,
|
||||||
|
dataset=dataset,
|
||||||
|
events=events,
|
||||||
|
episode_time_s=cfg.episode_time_s,
|
||||||
|
display_cameras=cfg.display_cameras,
|
||||||
|
policy=policy,
|
||||||
|
fps=cfg.fps,
|
||||||
|
single_task=cfg.single_task,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Skip reset for the last episode to be recorded
|
||||||
|
if not events["stop_recording"] and (
|
||||||
|
(recorded_episodes < cfg.num_episodes - 1) or events["rerecord_episode"]
|
||||||
|
):
|
||||||
|
log_say("Reset the environment", cfg.play_sounds)
|
||||||
|
pprint("Reset the environment, stop recording")
|
||||||
|
|
||||||
|
if events["rerecord_episode"]:
|
||||||
|
log_say("Re-record episode", cfg.play_sounds)
|
||||||
|
pprint("Re-record episode")
|
||||||
|
events["rerecord_episode"] = False
|
||||||
|
events["exit_early"] = False
|
||||||
|
dataset.clear_episode_buffer()
|
||||||
|
continue
|
||||||
|
|
||||||
|
dataset.save_episode()
|
||||||
|
recorded_episodes += 1
|
||||||
|
|
||||||
|
if events["stop_recording"]:
|
||||||
|
break
|
||||||
|
|
||||||
|
log_say("Stop recording", cfg.play_sounds, blocking=True)
|
||||||
|
stop_recording(robot, listener, cfg.display_cameras)
|
||||||
|
|
||||||
|
if cfg.push_to_hub:
|
||||||
|
dataset.push_to_hub(tags=cfg.tags, private=cfg.private)
|
||||||
|
|
||||||
|
log_say("Exiting", cfg.play_sounds)
|
||||||
|
return dataset
|
||||||
32
lerobot_aloha/common/utils/replay_utils.py
Normal file
32
lerobot_aloha/common/utils/replay_utils.py
Normal file
@@ -0,0 +1,32 @@
|
|||||||
|
import time
|
||||||
|
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
|
||||||
|
from common.utils.control_utils import busy_wait
|
||||||
|
|
||||||
|
|
||||||
|
def replay(
|
||||||
|
robot,
|
||||||
|
cfg,
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Replay recorded robot data according to configuration.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
robot: Robot instance
|
||||||
|
cfg: Configuration object
|
||||||
|
"""
|
||||||
|
# Load dataset
|
||||||
|
dataset = LeRobotDataset(cfg.repo_id, root=cfg.root, episodes=[cfg.episode])
|
||||||
|
actions = dataset.hf_dataset.select_columns("action")
|
||||||
|
|
||||||
|
print(f"Replaying episode {cfg.episode} from dataset {cfg.repo_id}")
|
||||||
|
print(f"Total frames: {dataset.num_frames}")
|
||||||
|
|
||||||
|
# Replay each frame
|
||||||
|
for idx in range(dataset.num_frames):
|
||||||
|
start_episode_t = time.perf_counter()
|
||||||
|
|
||||||
|
action = actions[idx]["action"]
|
||||||
|
robot.send_action(action)
|
||||||
|
|
||||||
|
dt_s = time.perf_counter() - start_episode_t
|
||||||
|
busy_wait(1 / cfg.fps - dt_s)
|
||||||
@@ -1,769 +0,0 @@
|
|||||||
#!/home/lin/software/miniconda3/envs/aloha/bin/python
|
|
||||||
# -- coding: UTF-8
|
|
||||||
"""
|
|
||||||
#!/usr/bin/python3
|
|
||||||
"""
|
|
||||||
|
|
||||||
import torch
|
|
||||||
import numpy as np
|
|
||||||
import os
|
|
||||||
import pickle
|
|
||||||
import argparse
|
|
||||||
from einops import rearrange
|
|
||||||
import collections
|
|
||||||
from collections import deque
|
|
||||||
|
|
||||||
import rospy
|
|
||||||
from std_msgs.msg import Header
|
|
||||||
from geometry_msgs.msg import Twist
|
|
||||||
from sensor_msgs.msg import JointState, Image
|
|
||||||
from nav_msgs.msg import Odometry
|
|
||||||
from cv_bridge import CvBridge
|
|
||||||
import time
|
|
||||||
import threading
|
|
||||||
import math
|
|
||||||
import threading
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
import sys
|
|
||||||
sys.path.append("./")
|
|
||||||
|
|
||||||
SEED = 42
|
|
||||||
torch.manual_seed(SEED)
|
|
||||||
np.random.seed(SEED)
|
|
||||||
|
|
||||||
task_config = {'camera_names': ['cam_high', 'cam_left_wrist', 'cam_right_wrist']}
|
|
||||||
|
|
||||||
inference_thread = None
|
|
||||||
inference_lock = threading.Lock()
|
|
||||||
inference_actions = None
|
|
||||||
inference_timestep = None
|
|
||||||
|
|
||||||
|
|
||||||
def actions_interpolation(args, pre_action, actions, stats):
|
|
||||||
steps = np.concatenate((np.array(args.arm_steps_length), np.array(args.arm_steps_length)), axis=0)
|
|
||||||
pre_process = lambda s_qpos: (s_qpos - stats['qpos_mean']) / stats['qpos_std']
|
|
||||||
post_process = lambda a: a * stats['action_std'] + stats['action_mean']
|
|
||||||
result = [pre_action]
|
|
||||||
post_action = post_process(actions[0])
|
|
||||||
# print("pre_action:", pre_action[7:])
|
|
||||||
# print("actions_interpolation1:", post_action[:, 7:])
|
|
||||||
max_diff_index = 0
|
|
||||||
max_diff = -1
|
|
||||||
for i in range(post_action.shape[0]):
|
|
||||||
diff = 0
|
|
||||||
for j in range(pre_action.shape[0]):
|
|
||||||
if j == 6 or j == 13:
|
|
||||||
continue
|
|
||||||
diff += math.fabs(pre_action[j] - post_action[i][j])
|
|
||||||
if diff > max_diff:
|
|
||||||
max_diff = diff
|
|
||||||
max_diff_index = i
|
|
||||||
|
|
||||||
for i in range(max_diff_index, post_action.shape[0]):
|
|
||||||
step = max([math.floor(math.fabs(result[-1][j] - post_action[i][j])/steps[j]) for j in range(pre_action.shape[0])])
|
|
||||||
inter = np.linspace(result[-1], post_action[i], step+2)
|
|
||||||
result.extend(inter[1:])
|
|
||||||
while len(result) < args.chunk_size+1:
|
|
||||||
result.append(result[-1])
|
|
||||||
result = np.array(result)[1:args.chunk_size+1]
|
|
||||||
# print("actions_interpolation2:", result.shape, result[:, 7:])
|
|
||||||
result = pre_process(result)
|
|
||||||
result = result[np.newaxis, :]
|
|
||||||
return result
|
|
||||||
|
|
||||||
|
|
||||||
def get_model_config(args):
|
|
||||||
# 设置随机种子，你可以确保在相同的初始条件下，每次运行代码时生成的随机数序列是相同的。
|
|
||||||
set_seed(1)
|
|
||||||
|
|
||||||
# 如果是ACT策略
|
|
||||||
# fixed parameters
|
|
||||||
if args.policy_class == 'ACT':
|
|
||||||
policy_config = {'lr': args.lr,
|
|
||||||
'lr_backbone': args.lr_backbone,
|
|
||||||
'backbone': args.backbone,
|
|
||||||
'masks': args.masks,
|
|
||||||
'weight_decay': args.weight_decay,
|
|
||||||
'dilation': args.dilation,
|
|
||||||
'position_embedding': args.position_embedding,
|
|
||||||
'loss_function': args.loss_function,
|
|
||||||
'chunk_size': args.chunk_size, # 查询
|
|
||||||
'camera_names': task_config['camera_names'],
|
|
||||||
'use_depth_image': args.use_depth_image,
|
|
||||||
'use_robot_base': args.use_robot_base,
|
|
||||||
'kl_weight': args.kl_weight, # kl散度权重
|
|
||||||
'hidden_dim': args.hidden_dim, # 隐藏层维度
|
|
||||||
'dim_feedforward': args.dim_feedforward,
|
|
||||||
'enc_layers': args.enc_layers,
|
|
||||||
'dec_layers': args.dec_layers,
|
|
||||||
'nheads': args.nheads,
|
|
||||||
'dropout': args.dropout,
|
|
||||||
'pre_norm': args.pre_norm
|
|
||||||
}
|
|
||||||
elif args.policy_class == 'CNNMLP':
|
|
||||||
policy_config = {'lr': args.lr,
|
|
||||||
'lr_backbone': args.lr_backbone,
|
|
||||||
'backbone': args.backbone,
|
|
||||||
'masks': args.masks,
|
|
||||||
'weight_decay': args.weight_decay,
|
|
||||||
'dilation': args.dilation,
|
|
||||||
'position_embedding': args.position_embedding,
|
|
||||||
'loss_function': args.loss_function,
|
|
||||||
'chunk_size': 1, # 查询
|
|
||||||
'camera_names': task_config['camera_names'],
|
|
||||||
'use_depth_image': args.use_depth_image,
|
|
||||||
'use_robot_base': args.use_robot_base
|
|
||||||
}
|
|
||||||
|
|
||||||
elif args.policy_class == 'Diffusion':
|
|
||||||
policy_config = {'lr': args.lr,
|
|
||||||
'lr_backbone': args.lr_backbone,
|
|
||||||
'backbone': args.backbone,
|
|
||||||
'masks': args.masks,
|
|
||||||
'weight_decay': args.weight_decay,
|
|
||||||
'dilation': args.dilation,
|
|
||||||
'position_embedding': args.position_embedding,
|
|
||||||
'loss_function': args.loss_function,
|
|
||||||
'chunk_size': args.chunk_size, # 查询
|
|
||||||
'camera_names': task_config['camera_names'],
|
|
||||||
'use_depth_image': args.use_depth_image,
|
|
||||||
'use_robot_base': args.use_robot_base,
|
|
||||||
'observation_horizon': args.observation_horizon,
|
|
||||||
'action_horizon': args.action_horizon,
|
|
||||||
'num_inference_timesteps': args.num_inference_timesteps,
|
|
||||||
'ema_power': args.ema_power
|
|
||||||
}
|
|
||||||
else:
|
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
config = {
|
|
||||||
'ckpt_dir': args.ckpt_dir,
|
|
||||||
'ckpt_name': args.ckpt_name,
|
|
||||||
'ckpt_stats_name': args.ckpt_stats_name,
|
|
||||||
'episode_len': args.max_publish_step,
|
|
||||||
'state_dim': args.state_dim,
|
|
||||||
'policy_class': args.policy_class,
|
|
||||||
'policy_config': policy_config,
|
|
||||||
'temporal_agg': args.temporal_agg,
|
|
||||||
'camera_names': task_config['camera_names'],
|
|
||||||
}
|
|
||||||
return config
|
|
||||||
|
|
||||||
|
|
||||||
def make_policy(policy_class, policy_config):
|
|
||||||
if policy_class == 'ACT':
|
|
||||||
policy = ACTPolicy(policy_config)
|
|
||||||
elif policy_class == 'CNNMLP':
|
|
||||||
policy = CNNMLPPolicy(policy_config)
|
|
||||||
elif policy_class == 'Diffusion':
|
|
||||||
policy = DiffusionPolicy(policy_config)
|
|
||||||
else:
|
|
||||||
raise NotImplementedError
|
|
||||||
return policy
|
|
||||||
|
|
||||||
|
|
||||||
def get_image(observation, camera_names):
|
|
||||||
curr_images = []
|
|
||||||
for cam_name in camera_names:
|
|
||||||
curr_image = rearrange(observation['images'][cam_name], 'h w c -> c h w')
|
|
||||||
|
|
||||||
curr_images.append(curr_image)
|
|
||||||
curr_image = np.stack(curr_images, axis=0)
|
|
||||||
curr_image = torch.from_numpy(curr_image / 255.0).float().cuda().unsqueeze(0)
|
|
||||||
return curr_image
|
|
||||||
|
|
||||||
|
|
||||||
def get_depth_image(observation, camera_names):
|
|
||||||
curr_images = []
|
|
||||||
for cam_name in camera_names:
|
|
||||||
curr_images.append(observation['images_depth'][cam_name])
|
|
||||||
curr_image = np.stack(curr_images, axis=0)
|
|
||||||
curr_image = torch.from_numpy(curr_image / 255.0).float().cuda().unsqueeze(0)
|
|
||||||
return curr_image
|
|
||||||
|
|
||||||
|
|
||||||
def inference_process(args, config, ros_operator, policy, stats, t, pre_action):
|
|
||||||
global inference_lock
|
|
||||||
global inference_actions
|
|
||||||
global inference_timestep
|
|
||||||
print_flag = True
|
|
||||||
pre_pos_process = lambda s_qpos: (s_qpos - stats['qpos_mean']) / stats['qpos_std']
|
|
||||||
pre_action_process = lambda next_action: (next_action - stats["action_mean"]) / stats["action_std"]
|
|
||||||
rate = rospy.Rate(args.publish_rate)
|
|
||||||
while True and not rospy.is_shutdown():
|
|
||||||
result = ros_operator.get_frame()
|
|
||||||
if not result:
|
|
||||||
if print_flag:
|
|
||||||
print("syn fail")
|
|
||||||
print_flag = False
|
|
||||||
rate.sleep()
|
|
||||||
continue
|
|
||||||
print_flag = True
|
|
||||||
(img_front, img_left, img_right, img_front_depth, img_left_depth, img_right_depth,
|
|
||||||
puppet_arm_left, puppet_arm_right, robot_base) = result
|
|
||||||
obs = collections.OrderedDict()
|
|
||||||
image_dict = dict()
|
|
||||||
|
|
||||||
image_dict[config['camera_names'][0]] = img_front
|
|
||||||
image_dict[config['camera_names'][1]] = img_left
|
|
||||||
image_dict[config['camera_names'][2]] = img_right
|
|
||||||
|
|
||||||
|
|
||||||
obs['images'] = image_dict
|
|
||||||
|
|
||||||
if args.use_depth_image:
|
|
||||||
image_depth_dict = dict()
|
|
||||||
image_depth_dict[config['camera_names'][0]] = img_front_depth
|
|
||||||
image_depth_dict[config['camera_names'][1]] = img_left_depth
|
|
||||||
image_depth_dict[config['camera_names'][2]] = img_right_depth
|
|
||||||
obs['images_depth'] = image_depth_dict
|
|
||||||
|
|
||||||
obs['qpos'] = np.concatenate(
|
|
||||||
(np.array(puppet_arm_left.position), np.array(puppet_arm_right.position)), axis=0)
|
|
||||||
obs['qvel'] = np.concatenate(
|
|
||||||
(np.array(puppet_arm_left.velocity), np.array(puppet_arm_right.velocity)), axis=0)
|
|
||||||
obs['effort'] = np.concatenate(
|
|
||||||
(np.array(puppet_arm_left.effort), np.array(puppet_arm_right.effort)), axis=0)
|
|
||||||
if args.use_robot_base:
|
|
||||||
obs['base_vel'] = [robot_base.twist.twist.linear.x, robot_base.twist.twist.angular.z]
|
|
||||||
obs['qpos'] = np.concatenate((obs['qpos'], obs['base_vel']), axis=0)
|
|
||||||
else:
|
|
||||||
obs['base_vel'] = [0.0, 0.0]
|
|
||||||
# qpos_numpy = np.array(obs['qpos'])
|
|
||||||
|
|
||||||
# 归一化处理qpos 并转到cuda
|
|
||||||
qpos = pre_pos_process(obs['qpos'])
|
|
||||||
qpos = torch.from_numpy(qpos).float().cuda().unsqueeze(0)
|
|
||||||
# 当前图像curr_image获取图像
|
|
||||||
curr_image = get_image(obs, config['camera_names'])
|
|
||||||
curr_depth_image = None
|
|
||||||
if args.use_depth_image:
|
|
||||||
curr_depth_image = get_depth_image(obs, config['camera_names'])
|
|
||||||
start_time = time.time()
|
|
||||||
all_actions = policy(curr_image, curr_depth_image, qpos)
|
|
||||||
end_time = time.time()
|
|
||||||
print("model cost time: ", end_time -start_time)
|
|
||||||
inference_lock.acquire()
|
|
||||||
inference_actions = all_actions.cpu().detach().numpy()
|
|
||||||
if pre_action is None:
|
|
||||||
pre_action = obs['qpos']
|
|
||||||
# print("obs['qpos']:", obs['qpos'][7:])
|
|
||||||
if args.use_actions_interpolation:
|
|
||||||
inference_actions = actions_interpolation(args, pre_action, inference_actions, stats)
|
|
||||||
inference_timestep = t
|
|
||||||
inference_lock.release()
|
|
||||||
break
|
|
||||||
|
|
||||||
|
|
||||||
def model_inference(args, config, ros_operator, save_episode=True):
|
|
||||||
global inference_lock
|
|
||||||
global inference_actions
|
|
||||||
global inference_timestep
|
|
||||||
global inference_thread
|
|
||||||
set_seed(1000)
|
|
||||||
|
|
||||||
# 1 创建模型数据 继承nn.Module
|
|
||||||
policy = make_policy(config['policy_class'], config['policy_config'])
|
|
||||||
# print("model structure\n", policy.model)
|
|
||||||
|
|
||||||
# 2 加载模型权重
|
|
||||||
ckpt_path = os.path.join(config['ckpt_dir'], config['ckpt_name'])
|
|
||||||
state_dict = torch.load(ckpt_path)
|
|
||||||
new_state_dict = {}
|
|
||||||
for key, value in state_dict.items():
|
|
||||||
if key in ["model.is_pad_head.weight", "model.is_pad_head.bias"]:
|
|
||||||
continue
|
|
||||||
if key in ["model.input_proj_next_action.weight", "model.input_proj_next_action.bias"]:
|
|
||||||
continue
|
|
||||||
new_state_dict[key] = value
|
|
||||||
loading_status = policy.deserialize(new_state_dict)
|
|
||||||
if not loading_status:
|
|
||||||
print("ckpt path not exist")
|
|
||||||
return False
|
|
||||||
|
|
||||||
# 3 模型设置为cuda模式和验证模式
|
|
||||||
policy.cuda()
|
|
||||||
policy.eval()
|
|
||||||
|
|
||||||
# 4 加载统计值
|
|
||||||
stats_path = os.path.join(config['ckpt_dir'], config['ckpt_stats_name'])
|
|
||||||
# 统计的数据 # 加载action_mean, action_std, qpos_mean, qpos_std 14维
|
|
||||||
with open(stats_path, 'rb') as f:
|
|
||||||
stats = pickle.load(f)
|
|
||||||
|
|
||||||
# 数据预处理和后处理函数定义
|
|
||||||
pre_process = lambda s_qpos: (s_qpos - stats['qpos_mean']) / stats['qpos_std']
|
|
||||||
post_process = lambda a: a * stats['action_std'] + stats['action_mean']
|
|
||||||
|
|
||||||
max_publish_step = config['episode_len']
|
|
||||||
chunk_size = config['policy_config']['chunk_size']
|
|
||||||
|
|
||||||
# 发布基础的姿态
|
|
||||||
left0 = [-0.00133514404296875, 0.00209808349609375, 0.01583099365234375, -0.032616615295410156, -0.00286102294921875, 0.00095367431640625, 3.557830810546875]
|
|
||||||
right0 = [-0.00133514404296875, 0.00438690185546875, 0.034523963928222656, -0.053597450256347656, -0.00476837158203125, -0.00209808349609375, 3.557830810546875]
|
|
||||||
left1 = [-0.00133514404296875, 0.00209808349609375, 0.01583099365234375, -0.032616615295410156, -0.00286102294921875, 0.00095367431640625, -0.3393220901489258]
|
|
||||||
right1 = [-0.00133514404296875, 0.00247955322265625, 0.01583099365234375, -0.032616615295410156, -0.00286102294921875, 0.00095367431640625, -0.3397035598754883]
|
|
||||||
|
|
||||||
ros_operator.puppet_arm_publish_continuous(left0, right0)
|
|
||||||
input("Enter any key to continue :")
|
|
||||||
ros_operator.puppet_arm_publish_continuous(left1, right1)
|
|
||||||
action = None
|
|
||||||
# 推理
|
|
||||||
with torch.inference_mode():
|
|
||||||
while True and not rospy.is_shutdown():
|
|
||||||
# 每个回合的步数
|
|
||||||
t = 0
|
|
||||||
max_t = 0
|
|
||||||
rate = rospy.Rate(args.publish_rate)
|
|
||||||
if config['temporal_agg']:
|
|
||||||
all_time_actions = np.zeros([max_publish_step, max_publish_step + chunk_size, config['state_dim']])
|
|
||||||
while t < max_publish_step and not rospy.is_shutdown():
|
|
||||||
# start_time = time.time()
|
|
||||||
# query policy
|
|
||||||
if config['policy_class'] == "ACT":
|
|
||||||
if t >= max_t:
|
|
||||||
pre_action = action
|
|
||||||
inference_thread = threading.Thread(target=inference_process,
|
|
||||||
args=(args, config, ros_operator,
|
|
||||||
policy, stats, t, pre_action))
|
|
||||||
inference_thread.start()
|
|
||||||
inference_thread.join()
|
|
||||||
inference_lock.acquire()
|
|
||||||
if inference_actions is not None:
|
|
||||||
inference_thread = None
|
|
||||||
all_actions = inference_actions
|
|
||||||
inference_actions = None
|
|
||||||
max_t = t + args.pos_lookahead_step
|
|
||||||
if config['temporal_agg']:
|
|
||||||
all_time_actions[[t], t:t + chunk_size] = all_actions
|
|
||||||
inference_lock.release()
|
|
||||||
if config['temporal_agg']:
|
|
||||||
actions_for_curr_step = all_time_actions[:, t]
|
|
||||||
actions_populated = np.all(actions_for_curr_step != 0, axis=1)
|
|
||||||
actions_for_curr_step = actions_for_curr_step[actions_populated]
|
|
||||||
k = 0.01
|
|
||||||
exp_weights = np.exp(-k * np.arange(len(actions_for_curr_step)))
|
|
||||||
exp_weights = exp_weights / exp_weights.sum()
|
|
||||||
exp_weights = exp_weights[:, np.newaxis]
|
|
||||||
raw_action = (actions_for_curr_step * exp_weights).sum(axis=0, keepdims=True)
|
|
||||||
else:
|
|
||||||
if args.pos_lookahead_step != 0:
|
|
||||||
raw_action = all_actions[:, t % args.pos_lookahead_step]
|
|
||||||
else:
|
|
||||||
raw_action = all_actions[:, t % chunk_size]
|
|
||||||
else:
|
|
||||||
raise NotImplementedError
|
|
||||||
action = post_process(raw_action[0])
|
|
||||||
left_action = action[:7] # 取7维度
|
|
||||||
right_action = action[7:14]
|
|
||||||
ros_operator.puppet_arm_publish(left_action, right_action) # puppet_arm_publish_continuous_thread
|
|
||||||
if args.use_robot_base:
|
|
||||||
vel_action = action[14:16]
|
|
||||||
ros_operator.robot_base_publish(vel_action)
|
|
||||||
t += 1
|
|
||||||
# end_time = time.time()
|
|
||||||
# print("publish: ", t)
|
|
||||||
# print("time:", end_time - start_time)
|
|
||||||
# print("left_action:", left_action)
|
|
||||||
# print("right_action:", right_action)
|
|
||||||
rate.sleep()
|
|
||||||
|
|
||||||
|
|
||||||
class RosOperator:
|
|
||||||
def __init__(self, args):
|
|
||||||
self.robot_base_deque = None
|
|
||||||
self.puppet_arm_right_deque = None
|
|
||||||
self.puppet_arm_left_deque = None
|
|
||||||
self.img_front_deque = None
|
|
||||||
self.img_right_deque = None
|
|
||||||
self.img_left_deque = None
|
|
||||||
self.img_front_depth_deque = None
|
|
||||||
self.img_right_depth_deque = None
|
|
||||||
self.img_left_depth_deque = None
|
|
||||||
self.bridge = None
|
|
||||||
self.puppet_arm_left_publisher = None
|
|
||||||
self.puppet_arm_right_publisher = None
|
|
||||||
self.robot_base_publisher = None
|
|
||||||
self.puppet_arm_publish_thread = None
|
|
||||||
self.puppet_arm_publish_lock = None
|
|
||||||
self.args = args
|
|
||||||
self.ctrl_state = False
|
|
||||||
self.ctrl_state_lock = threading.Lock()
|
|
||||||
self.init()
|
|
||||||
self.init_ros()
|
|
||||||
|
|
||||||
def init(self):
|
|
||||||
self.bridge = CvBridge()
|
|
||||||
self.img_left_deque = deque()
|
|
||||||
self.img_right_deque = deque()
|
|
||||||
self.img_front_deque = deque()
|
|
||||||
self.img_left_depth_deque = deque()
|
|
||||||
self.img_right_depth_deque = deque()
|
|
||||||
self.img_front_depth_deque = deque()
|
|
||||||
self.puppet_arm_left_deque = deque()
|
|
||||||
self.puppet_arm_right_deque = deque()
|
|
||||||
self.robot_base_deque = deque()
|
|
||||||
self.puppet_arm_publish_lock = threading.Lock()
|
|
||||||
self.puppet_arm_publish_lock.acquire()
|
|
||||||
|
|
||||||
def puppet_arm_publish(self, left, right):
|
|
||||||
joint_state_msg = JointState()
|
|
||||||
joint_state_msg.header = Header()
|
|
||||||
joint_state_msg.header.stamp = rospy.Time.now() # 设置时间戳
|
|
||||||
joint_state_msg.name = ['joint0', 'joint1', 'joint2', 'joint3', 'joint4', 'joint5', 'joint6'] # 设置关节名称
|
|
||||||
joint_state_msg.position = left
|
|
||||||
self.puppet_arm_left_publisher.publish(joint_state_msg)
|
|
||||||
joint_state_msg.position = right
|
|
||||||
self.puppet_arm_right_publisher.publish(joint_state_msg)
|
|
||||||
|
|
||||||
def robot_base_publish(self, vel):
|
|
||||||
vel_msg = Twist()
|
|
||||||
vel_msg.linear.x = vel[0]
|
|
||||||
vel_msg.linear.y = 0
|
|
||||||
vel_msg.linear.z = 0
|
|
||||||
vel_msg.angular.x = 0
|
|
||||||
vel_msg.angular.y = 0
|
|
||||||
vel_msg.angular.z = vel[1]
|
|
||||||
self.robot_base_publisher.publish(vel_msg)
|
|
||||||
|
|
||||||
def puppet_arm_publish_continuous(self, left, right):
|
|
||||||
rate = rospy.Rate(self.args.publish_rate)
|
|
||||||
left_arm = None
|
|
||||||
right_arm = None
|
|
||||||
while True and not rospy.is_shutdown():
|
|
||||||
if len(self.puppet_arm_left_deque) != 0:
|
|
||||||
left_arm = list(self.puppet_arm_left_deque[-1].position)
|
|
||||||
if len(self.puppet_arm_right_deque) != 0:
|
|
||||||
right_arm = list(self.puppet_arm_right_deque[-1].position)
|
|
||||||
if left_arm is None or right_arm is None:
|
|
||||||
rate.sleep()
|
|
||||||
continue
|
|
||||||
else:
|
|
||||||
break
|
|
||||||
left_symbol = [1 if left[i] - left_arm[i] > 0 else -1 for i in range(len(left))]
|
|
||||||
right_symbol = [1 if right[i] - right_arm[i] > 0 else -1 for i in range(len(right))]
|
|
||||||
flag = True
|
|
||||||
step = 0
|
|
||||||
while flag and not rospy.is_shutdown():
|
|
||||||
if self.puppet_arm_publish_lock.acquire(False):
|
|
||||||
return
|
|
||||||
left_diff = [abs(left[i] - left_arm[i]) for i in range(len(left))]
|
|
||||||
right_diff = [abs(right[i] - right_arm[i]) for i in range(len(right))]
|
|
||||||
flag = False
|
|
||||||
for i in range(len(left)):
|
|
||||||
if left_diff[i] < self.args.arm_steps_length[i]:
|
|
||||||
left_arm[i] = left[i]
|
|
||||||
else:
|
|
||||||
left_arm[i] += left_symbol[i] * self.args.arm_steps_length[i]
|
|
||||||
flag = True
|
|
||||||
for i in range(len(right)):
|
|
||||||
if right_diff[i] < self.args.arm_steps_length[i]:
|
|
||||||
right_arm[i] = right[i]
|
|
||||||
else:
|
|
||||||
right_arm[i] += right_symbol[i] * self.args.arm_steps_length[i]
|
|
||||||
flag = True
|
|
||||||
joint_state_msg = JointState()
|
|
||||||
joint_state_msg.header = Header()
|
|
||||||
joint_state_msg.header.stamp = rospy.Time.now() # 设置时间戳
|
|
||||||
joint_state_msg.name = ['joint0', 'joint1', 'joint2', 'joint3', 'joint4', 'joint5', 'joint6'] # 设置关节名称
|
|
||||||
joint_state_msg.position = left_arm
|
|
||||||
self.puppet_arm_left_publisher.publish(joint_state_msg)
|
|
||||||
joint_state_msg.position = right_arm
|
|
||||||
self.puppet_arm_right_publisher.publish(joint_state_msg)
|
|
||||||
step += 1
|
|
||||||
print("puppet_arm_publish_continuous:", step)
|
|
||||||
rate.sleep()
|
|
||||||
|
|
||||||
def puppet_arm_publish_linear(self, left, right):
|
|
||||||
num_step = 100
|
|
||||||
rate = rospy.Rate(200)
|
|
||||||
|
|
||||||
left_arm = None
|
|
||||||
right_arm = None
|
|
||||||
|
|
||||||
while True and not rospy.is_shutdown():
|
|
||||||
if len(self.puppet_arm_left_deque) != 0:
|
|
||||||
left_arm = list(self.puppet_arm_left_deque[-1].position)
|
|
||||||
if len(self.puppet_arm_right_deque) != 0:
|
|
||||||
right_arm = list(self.puppet_arm_right_deque[-1].position)
|
|
||||||
if left_arm is None or right_arm is None:
|
|
||||||
rate.sleep()
|
|
||||||
continue
|
|
||||||
else:
|
|
||||||
break
|
|
||||||
|
|
||||||
traj_left_list = np.linspace(left_arm, left, num_step)
|
|
||||||
traj_right_list = np.linspace(right_arm, right, num_step)
|
|
||||||
|
|
||||||
for i in range(len(traj_left_list)):
|
|
||||||
traj_left = traj_left_list[i]
|
|
||||||
traj_right = traj_right_list[i]
|
|
||||||
traj_left[-1] = left[-1]
|
|
||||||
traj_right[-1] = right[-1]
|
|
||||||
joint_state_msg = JointState()
|
|
||||||
joint_state_msg.header = Header()
|
|
||||||
joint_state_msg.header.stamp = rospy.Time.now() # 设置时间戳
|
|
||||||
joint_state_msg.name = ['joint0', 'joint1', 'joint2', 'joint3', 'joint4', 'joint5', 'joint6'] # 设置关节名称
|
|
||||||
joint_state_msg.position = traj_left
|
|
||||||
self.puppet_arm_left_publisher.publish(joint_state_msg)
|
|
||||||
joint_state_msg.position = traj_right
|
|
||||||
self.puppet_arm_right_publisher.publish(joint_state_msg)
|
|
||||||
rate.sleep()
|
|
||||||
|
|
||||||
def puppet_arm_publish_continuous_thread(self, left, right):
|
|
||||||
if self.puppet_arm_publish_thread is not None:
|
|
||||||
self.puppet_arm_publish_lock.release()
|
|
||||||
self.puppet_arm_publish_thread.join()
|
|
||||||
self.puppet_arm_publish_lock.acquire(False)
|
|
||||||
self.puppet_arm_publish_thread = None
|
|
||||||
self.puppet_arm_publish_thread = threading.Thread(target=self.puppet_arm_publish_continuous, args=(left, right))
|
|
||||||
self.puppet_arm_publish_thread.start()
|
|
||||||
|
|
||||||
def get_frame(self):
|
|
||||||
if len(self.img_left_deque) == 0 or len(self.img_right_deque) == 0 or len(self.img_front_deque) == 0 or \
|
|
||||||
(self.args.use_depth_image and (len(self.img_left_depth_deque) == 0 or len(self.img_right_depth_deque) == 0 or len(self.img_front_depth_deque) == 0)):
|
|
||||||
return False
|
|
||||||
if self.args.use_depth_image:
|
|
||||||
frame_time = min([self.img_left_deque[-1].header.stamp.to_sec(), self.img_right_deque[-1].header.stamp.to_sec(), self.img_front_deque[-1].header.stamp.to_sec(),
|
|
||||||
self.img_left_depth_deque[-1].header.stamp.to_sec(), self.img_right_depth_deque[-1].header.stamp.to_sec(), self.img_front_depth_deque[-1].header.stamp.to_sec()])
|
|
||||||
else:
|
|
||||||
frame_time = min([self.img_left_deque[-1].header.stamp.to_sec(), self.img_right_deque[-1].header.stamp.to_sec(), self.img_front_deque[-1].header.stamp.to_sec()])
|
|
||||||
|
|
||||||
if len(self.img_left_deque) == 0 or self.img_left_deque[-1].header.stamp.to_sec() < frame_time:
|
|
||||||
return False
|
|
||||||
if len(self.img_right_deque) == 0 or self.img_right_deque[-1].header.stamp.to_sec() < frame_time:
|
|
||||||
return False
|
|
||||||
if len(self.img_front_deque) == 0 or self.img_front_deque[-1].header.stamp.to_sec() < frame_time:
|
|
||||||
return False
|
|
||||||
if len(self.puppet_arm_left_deque) == 0 or self.puppet_arm_left_deque[-1].header.stamp.to_sec() < frame_time:
|
|
||||||
return False
|
|
||||||
if len(self.puppet_arm_right_deque) == 0 or self.puppet_arm_right_deque[-1].header.stamp.to_sec() < frame_time:
|
|
||||||
return False
|
|
||||||
if self.args.use_depth_image and (len(self.img_left_depth_deque) == 0 or self.img_left_depth_deque[-1].header.stamp.to_sec() < frame_time):
|
|
||||||
return False
|
|
||||||
if self.args.use_depth_image and (len(self.img_right_depth_deque) == 0 or self.img_right_depth_deque[-1].header.stamp.to_sec() < frame_time):
|
|
||||||
return False
|
|
||||||
if self.args.use_depth_image and (len(self.img_front_depth_deque) == 0 or self.img_front_depth_deque[-1].header.stamp.to_sec() < frame_time):
|
|
||||||
return False
|
|
||||||
if self.args.use_robot_base and (len(self.robot_base_deque) == 0 or self.robot_base_deque[-1].header.stamp.to_sec() < frame_time):
|
|
||||||
return False
|
|
||||||
|
|
||||||
while self.img_left_deque[0].header.stamp.to_sec() < frame_time:
|
|
||||||
self.img_left_deque.popleft()
|
|
||||||
img_left = self.bridge.imgmsg_to_cv2(self.img_left_deque.popleft(), 'passthrough')
|
|
||||||
|
|
||||||
while self.img_right_deque[0].header.stamp.to_sec() < frame_time:
|
|
||||||
self.img_right_deque.popleft()
|
|
||||||
img_right = self.bridge.imgmsg_to_cv2(self.img_right_deque.popleft(), 'passthrough')
|
|
||||||
|
|
||||||
while self.img_front_deque[0].header.stamp.to_sec() < frame_time:
|
|
||||||
self.img_front_deque.popleft()
|
|
||||||
img_front = self.bridge.imgmsg_to_cv2(self.img_front_deque.popleft(), 'passthrough')
|
|
||||||
|
|
||||||
while self.puppet_arm_left_deque[0].header.stamp.to_sec() < frame_time:
|
|
||||||
self.puppet_arm_left_deque.popleft()
|
|
||||||
puppet_arm_left = self.puppet_arm_left_deque.popleft()
|
|
||||||
|
|
||||||
while self.puppet_arm_right_deque[0].header.stamp.to_sec() < frame_time:
|
|
||||||
self.puppet_arm_right_deque.popleft()
|
|
||||||
puppet_arm_right = self.puppet_arm_right_deque.popleft()
|
|
||||||
|
|
||||||
img_left_depth = None
|
|
||||||
if self.args.use_depth_image:
|
|
||||||
while self.img_left_depth_deque[0].header.stamp.to_sec() < frame_time:
|
|
||||||
self.img_left_depth_deque.popleft()
|
|
||||||
img_left_depth = self.bridge.imgmsg_to_cv2(self.img_left_depth_deque.popleft(), 'passthrough')
|
|
||||||
|
|
||||||
img_right_depth = None
|
|
||||||
if self.args.use_depth_image:
|
|
||||||
while self.img_right_depth_deque[0].header.stamp.to_sec() < frame_time:
|
|
||||||
self.img_right_depth_deque.popleft()
|
|
||||||
img_right_depth = self.bridge.imgmsg_to_cv2(self.img_right_depth_deque.popleft(), 'passthrough')
|
|
||||||
|
|
||||||
img_front_depth = None
|
|
||||||
if self.args.use_depth_image:
|
|
||||||
while self.img_front_depth_deque[0].header.stamp.to_sec() < frame_time:
|
|
||||||
self.img_front_depth_deque.popleft()
|
|
||||||
img_front_depth = self.bridge.imgmsg_to_cv2(self.img_front_depth_deque.popleft(), 'passthrough')
|
|
||||||
|
|
||||||
robot_base = None
|
|
||||||
if self.args.use_robot_base:
|
|
||||||
while self.robot_base_deque[0].header.stamp.to_sec() < frame_time:
|
|
||||||
self.robot_base_deque.popleft()
|
|
||||||
robot_base = self.robot_base_deque.popleft()
|
|
||||||
|
|
||||||
return (img_front, img_left, img_right, img_front_depth, img_left_depth, img_right_depth,
|
|
||||||
puppet_arm_left, puppet_arm_right, robot_base)
|
|
||||||
|
|
||||||
def img_left_callback(self, msg):
|
|
||||||
if len(self.img_left_deque) >= 2000:
|
|
||||||
self.img_left_deque.popleft()
|
|
||||||
self.img_left_deque.append(msg)
|
|
||||||
|
|
||||||
def img_right_callback(self, msg):
|
|
||||||
if len(self.img_right_deque) >= 2000:
|
|
||||||
self.img_right_deque.popleft()
|
|
||||||
self.img_right_deque.append(msg)
|
|
||||||
|
|
||||||
def img_front_callback(self, msg):
|
|
||||||
if len(self.img_front_deque) >= 2000:
|
|
||||||
self.img_front_deque.popleft()
|
|
||||||
self.img_front_deque.append(msg)
|
|
||||||
|
|
||||||
def img_left_depth_callback(self, msg):
|
|
||||||
if len(self.img_left_depth_deque) >= 2000:
|
|
||||||
self.img_left_depth_deque.popleft()
|
|
||||||
self.img_left_depth_deque.append(msg)
|
|
||||||
|
|
||||||
def img_right_depth_callback(self, msg):
|
|
||||||
if len(self.img_right_depth_deque) >= 2000:
|
|
||||||
self.img_right_depth_deque.popleft()
|
|
||||||
self.img_right_depth_deque.append(msg)
|
|
||||||
|
|
||||||
def img_front_depth_callback(self, msg):
|
|
||||||
if len(self.img_front_depth_deque) >= 2000:
|
|
||||||
self.img_front_depth_deque.popleft()
|
|
||||||
self.img_front_depth_deque.append(msg)
|
|
||||||
|
|
||||||
def puppet_arm_left_callback(self, msg):
|
|
||||||
if len(self.puppet_arm_left_deque) >= 2000:
|
|
||||||
self.puppet_arm_left_deque.popleft()
|
|
||||||
self.puppet_arm_left_deque.append(msg)
|
|
||||||
|
|
||||||
def puppet_arm_right_callback(self, msg):
|
|
||||||
if len(self.puppet_arm_right_deque) >= 2000:
|
|
||||||
self.puppet_arm_right_deque.popleft()
|
|
||||||
self.puppet_arm_right_deque.append(msg)
|
|
||||||
|
|
||||||
def robot_base_callback(self, msg):
|
|
||||||
if len(self.robot_base_deque) >= 2000:
|
|
||||||
self.robot_base_deque.popleft()
|
|
||||||
self.robot_base_deque.append(msg)
|
|
||||||
|
|
||||||
def ctrl_callback(self, msg):
|
|
||||||
self.ctrl_state_lock.acquire()
|
|
||||||
self.ctrl_state = msg.data
|
|
||||||
self.ctrl_state_lock.release()
|
|
||||||
|
|
||||||
def get_ctrl_state(self):
|
|
||||||
self.ctrl_state_lock.acquire()
|
|
||||||
state = self.ctrl_state
|
|
||||||
self.ctrl_state_lock.release()
|
|
||||||
return state
|
|
||||||
|
|
||||||
def init_ros(self):
|
|
||||||
rospy.init_node('joint_state_publisher', anonymous=True)
|
|
||||||
rospy.Subscriber(self.args.img_left_topic, Image, self.img_left_callback, queue_size=1000, tcp_nodelay=True)
|
|
||||||
rospy.Subscriber(self.args.img_right_topic, Image, self.img_right_callback, queue_size=1000, tcp_nodelay=True)
|
|
||||||
rospy.Subscriber(self.args.img_front_topic, Image, self.img_front_callback, queue_size=1000, tcp_nodelay=True)
|
|
||||||
if self.args.use_depth_image:
|
|
||||||
rospy.Subscriber(self.args.img_left_depth_topic, Image, self.img_left_depth_callback, queue_size=1000, tcp_nodelay=True)
|
|
||||||
rospy.Subscriber(self.args.img_right_depth_topic, Image, self.img_right_depth_callback, queue_size=1000, tcp_nodelay=True)
|
|
||||||
rospy.Subscriber(self.args.img_front_depth_topic, Image, self.img_front_depth_callback, queue_size=1000, tcp_nodelay=True)
|
|
||||||
rospy.Subscriber(self.args.puppet_arm_left_topic, JointState, self.puppet_arm_left_callback, queue_size=1000, tcp_nodelay=True)
|
|
||||||
rospy.Subscriber(self.args.puppet_arm_right_topic, JointState, self.puppet_arm_right_callback, queue_size=1000, tcp_nodelay=True)
|
|
||||||
rospy.Subscriber(self.args.robot_base_topic, Odometry, self.robot_base_callback, queue_size=1000, tcp_nodelay=True)
|
|
||||||
self.puppet_arm_left_publisher = rospy.Publisher(self.args.puppet_arm_left_cmd_topic, JointState, queue_size=10)
|
|
||||||
self.puppet_arm_right_publisher = rospy.Publisher(self.args.puppet_arm_right_cmd_topic, JointState, queue_size=10)
|
|
||||||
self.robot_base_publisher = rospy.Publisher(self.args.robot_base_cmd_topic, Twist, queue_size=10)
|
|
||||||
|
|
||||||
|
|
||||||
def get_arguments():
|
|
||||||
parser = argparse.ArgumentParser()
|
|
||||||
parser.add_argument('--ckpt_dir', action='store', type=str, help='ckpt_dir', required=True)
|
|
||||||
parser.add_argument('--task_name', action='store', type=str, help='task_name', default='aloha_mobile_dummy', required=False)
|
|
||||||
parser.add_argument('--max_publish_step', action='store', type=int, help='max_publish_step', default=10000, required=False)
|
|
||||||
parser.add_argument('--ckpt_name', action='store', type=str, help='ckpt_name', default='policy_best.ckpt', required=False)
|
|
||||||
parser.add_argument('--ckpt_stats_name', action='store', type=str, help='ckpt_stats_name', default='dataset_stats.pkl', required=False)
|
|
||||||
parser.add_argument('--policy_class', action='store', type=str, help='policy_class, capitalize', default='ACT', required=False)
|
|
||||||
parser.add_argument('--batch_size', action='store', type=int, help='batch_size', default=8, required=False)
|
|
||||||
parser.add_argument('--seed', action='store', type=int, help='seed', default=0, required=False)
|
|
||||||
parser.add_argument('--num_epochs', action='store', type=int, help='num_epochs', default=2000, required=False)
|
|
||||||
parser.add_argument('--lr', action='store', type=float, help='lr', default=1e-5, required=False)
|
|
||||||
parser.add_argument('--weight_decay', type=float, help='weight_decay', default=1e-4, required=False)
|
|
||||||
parser.add_argument('--dilation', action='store_true',
|
|
||||||
help="If true, we replace stride with dilation in the last convolutional block (DC5)", required=False)
|
|
||||||
parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'),
|
|
||||||
help="Type of positional embedding to use on top of the image features", required=False)
|
|
||||||
parser.add_argument('--masks', action='store_true',
|
|
||||||
help="Train segmentation head if the flag is provided")
|
|
||||||
parser.add_argument('--kl_weight', action='store', type=int, help='KL Weight', default=10, required=False)
|
|
||||||
parser.add_argument('--hidden_dim', action='store', type=int, help='hidden_dim', default=512, required=False)
|
|
||||||
parser.add_argument('--dim_feedforward', action='store', type=int, help='dim_feedforward', default=3200, required=False)
|
|
||||||
parser.add_argument('--temporal_agg', action='store', type=bool, help='temporal_agg', default=True, required=False)
|
|
||||||
|
|
||||||
parser.add_argument('--state_dim', action='store', type=int, help='state_dim', default=14, required=False)
|
|
||||||
parser.add_argument('--lr_backbone', action='store', type=float, help='lr_backbone', default=1e-5, required=False)
|
|
||||||
parser.add_argument('--backbone', action='store', type=str, help='backbone', default='resnet18', required=False)
|
|
||||||
parser.add_argument('--loss_function', action='store', type=str, help='loss_function l1 l2 l1+l2', default='l1', required=False)
|
|
||||||
parser.add_argument('--enc_layers', action='store', type=int, help='enc_layers', default=4, required=False)
|
|
||||||
parser.add_argument('--dec_layers', action='store', type=int, help='dec_layers', default=7, required=False)
|
|
||||||
parser.add_argument('--nheads', action='store', type=int, help='nheads', default=8, required=False)
|
|
||||||
parser.add_argument('--dropout', default=0.1, type=float, help="Dropout applied in the transformer", required=False)
|
|
||||||
parser.add_argument('--pre_norm', action='store_true', required=False)
|
|
||||||
|
|
||||||
parser.add_argument('--img_front_topic', action='store', type=str, help='img_front_topic',
|
|
||||||
default='/camera_f/color/image_raw', required=False)
|
|
||||||
parser.add_argument('--img_left_topic', action='store', type=str, help='img_left_topic',
|
|
||||||
default='/camera_l/color/image_raw', required=False)
|
|
||||||
parser.add_argument('--img_right_topic', action='store', type=str, help='img_right_topic',
|
|
||||||
default='/camera_r/color/image_raw', required=False)
|
|
||||||
|
|
||||||
parser.add_argument('--img_front_depth_topic', action='store', type=str, help='img_front_depth_topic',
|
|
||||||
default='/camera_f/depth/image_raw', required=False)
|
|
||||||
parser.add_argument('--img_left_depth_topic', action='store', type=str, help='img_left_depth_topic',
|
|
||||||
default='/camera_l/depth/image_raw', required=False)
|
|
||||||
parser.add_argument('--img_right_depth_topic', action='store', type=str, help='img_right_depth_topic',
|
|
||||||
default='/camera_r/depth/image_raw', required=False)
|
|
||||||
|
|
||||||
parser.add_argument('--puppet_arm_left_cmd_topic', action='store', type=str, help='puppet_arm_left_cmd_topic',
|
|
||||||
default='/master/joint_left', required=False)
|
|
||||||
parser.add_argument('--puppet_arm_right_cmd_topic', action='store', type=str, help='puppet_arm_right_cmd_topic',
|
|
||||||
default='/master/joint_right', required=False)
|
|
||||||
parser.add_argument('--puppet_arm_left_topic', action='store', type=str, help='puppet_arm_left_topic',
|
|
||||||
default='/puppet/joint_left', required=False)
|
|
||||||
parser.add_argument('--puppet_arm_right_topic', action='store', type=str, help='puppet_arm_right_topic',
|
|
||||||
default='/puppet/joint_right', required=False)
|
|
||||||
|
|
||||||
parser.add_argument('--robot_base_topic', action='store', type=str, help='robot_base_topic',
|
|
||||||
default='/odom_raw', required=False)
|
|
||||||
parser.add_argument('--robot_base_cmd_topic', action='store', type=str, help='robot_base_topic',
|
|
||||||
default='/cmd_vel', required=False)
|
|
||||||
parser.add_argument('--use_robot_base', action='store', type=bool, help='use_robot_base',
|
|
||||||
default=False, required=False)
|
|
||||||
parser.add_argument('--publish_rate', action='store', type=int, help='publish_rate',
|
|
||||||
default=40, required=False)
|
|
||||||
parser.add_argument('--pos_lookahead_step', action='store', type=int, help='pos_lookahead_step',
|
|
||||||
default=0, required=False)
|
|
||||||
parser.add_argument('--chunk_size', action='store', type=int, help='chunk_size',
|
|
||||||
default=32, required=False)
|
|
||||||
parser.add_argument('--arm_steps_length', action='store', type=float, help='arm_steps_length',
|
|
||||||
default=[0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.2], required=False)
|
|
||||||
|
|
||||||
parser.add_argument('--use_actions_interpolation', action='store', type=bool, help='use_actions_interpolation',
|
|
||||||
default=False, required=False)
|
|
||||||
parser.add_argument('--use_depth_image', action='store', type=bool, help='use_depth_image',
|
|
||||||
default=False, required=False)
|
|
||||||
|
|
||||||
# for Diffusion
|
|
||||||
parser.add_argument('--observation_horizon', action='store', type=int, help='observation_horizon', default=1, required=False)
|
|
||||||
parser.add_argument('--action_horizon', action='store', type=int, help='action_horizon', default=8, required=False)
|
|
||||||
parser.add_argument('--num_inference_timesteps', action='store', type=int, help='num_inference_timesteps', default=10, required=False)
|
|
||||||
parser.add_argument('--ema_power', action='store', type=int, help='ema_power', default=0.75, required=False)
|
|
||||||
args = parser.parse_args()
|
|
||||||
return args
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
args = get_arguments()
|
|
||||||
ros_operator = RosOperator(args)
|
|
||||||
config = get_model_config(args)
|
|
||||||
model_inference(args, config, ros_operator, save_episode=True)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
|
||||||
main()
|
|
||||||
# python act/inference.py --ckpt_dir ~/train0314/
|
|
||||||
56
lerobot_aloha/main.py
Normal file
56
lerobot_aloha/main.py
Normal file
@@ -0,0 +1,56 @@
|
|||||||
|
import argparse
|
||||||
|
from common.rosrobot_factory import RobotFactory
|
||||||
|
from common.utils.data_utils import record
|
||||||
|
from common.utils.replay_utils import replay
|
||||||
|
|
||||||
|
|
||||||
|
def get_arguments():
|
||||||
|
"""
|
||||||
|
Parse command line arguments.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
argparse.Namespace: Parsed arguments
|
||||||
|
"""
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
args = parser.parse_args()
|
||||||
|
args.fps = 30
|
||||||
|
args.resume = False
|
||||||
|
args.repo_id = "move_the_bottle_from_the_right_to_the_scale_right"
|
||||||
|
args.root = "./data5"
|
||||||
|
args.episode = 0 # replay episode
|
||||||
|
args.num_image_writer_processes = 0
|
||||||
|
args.num_image_writer_threads_per_camera = 4
|
||||||
|
args.video = True
|
||||||
|
args.num_episodes = 100
|
||||||
|
args.episode_time_s = 30000
|
||||||
|
args.play_sounds = False
|
||||||
|
args.display_cameras = True
|
||||||
|
args.single_task = "move the bottle from the right to the scale right"
|
||||||
|
args.use_depth_image = False
|
||||||
|
args.use_base = False
|
||||||
|
args.push_to_hub = False
|
||||||
|
args.policy = None
|
||||||
|
args.control_type = "record"
|
||||||
|
return args
|
||||||
|
|
||||||
|
|
||||||
|
def control_robot(cfg):
|
||||||
|
"""
|
||||||
|
Control robot based on configuration.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
cfg: Configuration object
|
||||||
|
"""
|
||||||
|
# Create robot instance using factory pattern
|
||||||
|
robot = RobotFactory.create(config_file="/home/ubuntu/LYT/lerobot_aloha/lerobot_aloha/configs/agilex.yaml", args=cfg)
|
||||||
|
|
||||||
|
# Execute appropriate control mode
|
||||||
|
if cfg.control_type == "record":
|
||||||
|
record(robot, cfg)
|
||||||
|
elif cfg.control_type == "replay":
|
||||||
|
replay(robot, cfg)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
cfg = get_arguments()
|
||||||
|
control_robot(cfg)
|
||||||
@@ -1,112 +1,36 @@
|
|||||||
#coding=utf-8
|
# coding=utf-8
|
||||||
import os
|
|
||||||
import numpy as np
|
|
||||||
import cv2
|
|
||||||
import h5py
|
|
||||||
import argparse
|
import argparse
|
||||||
import rospy
|
import rospy
|
||||||
|
from common.rosrobot_factory import RobotFactory
|
||||||
from cv_bridge import CvBridge
|
from common.utils.replay_utils import replay
|
||||||
from std_msgs.msg import Header
|
|
||||||
from sensor_msgs.msg import Image, JointState
|
|
||||||
from geometry_msgs.msg import Twist
|
|
||||||
from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
|
|
||||||
|
|
||||||
|
|
||||||
|
def get_arguments():
|
||||||
|
"""
|
||||||
|
Parse command line arguments.
|
||||||
|
|
||||||
def main(args):
|
Returns:
|
||||||
rospy.init_node("replay_node")
|
argparse.Namespace: Parsed arguments
|
||||||
bridge = CvBridge()
|
"""
|
||||||
# img_left_publisher = rospy.Publisher(args.img_left_topic, Image, queue_size=10)
|
|
||||||
# img_right_publisher = rospy.Publisher(args.img_right_topic, Image, queue_size=10)
|
|
||||||
# img_front_publisher = rospy.Publisher(args.img_front_topic, Image, queue_size=10)
|
|
||||||
|
|
||||||
# puppet_arm_left_publisher = rospy.Publisher(args.puppet_arm_left_topic, JointState, queue_size=10)
|
|
||||||
# puppet_arm_right_publisher = rospy.Publisher(args.puppet_arm_right_topic, JointState, queue_size=10)
|
|
||||||
|
|
||||||
master_arm_left_publisher = rospy.Publisher(args.master_arm_left_topic, JointState, queue_size=10)
|
|
||||||
master_arm_right_publisher = rospy.Publisher(args.master_arm_right_topic, JointState, queue_size=10)
|
|
||||||
|
|
||||||
# robot_base_publisher = rospy.Publisher(args.robot_base_topic, Twist, queue_size=10)
|
|
||||||
|
|
||||||
|
|
||||||
# dataset_dir = args.dataset_dir
|
|
||||||
# episode_idx = args.episode_idx
|
|
||||||
# task_name = args.task_name
|
|
||||||
# dataset_name = f'episode_{episode_idx}'
|
|
||||||
|
|
||||||
dataset = LeRobotDataset(args.repo_id, root=args.root, episodes=[args.episode])
|
|
||||||
actions = dataset.hf_dataset.select_columns("action")
|
|
||||||
velocitys = dataset.hf_dataset.select_columns("observation.velocity")
|
|
||||||
efforts = dataset.hf_dataset.select_columns("observation.effort")
|
|
||||||
|
|
||||||
origin_left = [-0.0057,-0.031, -0.0122, -0.032, 0.0099, 0.0179, 0.2279]
|
|
||||||
origin_right = [ 0.0616, 0.0021, 0.0475, -0.1013, 0.1097, 0.0872, 0.2279]
|
|
||||||
|
|
||||||
joint_state_msg = JointState()
|
|
||||||
joint_state_msg.header = Header()
|
|
||||||
joint_state_msg.name = ['joint0', 'joint1', 'joint2', 'joint3', 'joint4', 'joint5', ''] # 设置关节名称
|
|
||||||
twist_msg = Twist()
|
|
||||||
|
|
||||||
rate = rospy.Rate(args.fps)
|
|
||||||
|
|
||||||
# qposs, qvels, efforts, actions, base_actions, image_dicts = load_hdf5(os.path.join(dataset_dir, task_name), dataset_name)
|
|
||||||
|
|
||||||
|
|
||||||
last_action = [-0.00019073486328125, 0.00934600830078125, 0.01354217529296875, -0.01049041748046875, -0.00057220458984375, -0.00057220458984375, -0.00526118278503418, -0.00095367431640625, 0.00705718994140625, 0.01239776611328125, -0.00705718994140625, -0.00019073486328125, -0.00057220458984375, -0.009171326644718647]
|
|
||||||
last_velocity = [-0.010990142822265625, -0.010990142822265625, -0.03296661376953125, 0.010990142822265625, -0.010990142822265625, -0.010990142822265625, -0.010990142822265625, -0.010990142822265625, -0.010990142822265625, -0.03296661376953125, -0.010990142822265625, -0.010990142822265625, -0.03296661376953125, -0.03296661376953125]
|
|
||||||
last_effort = [-0.021978378295898438, 0.2417583465576172, 4.320878982543945, 3.6527481079101562, -0.013187408447265625, -0.013187408447265625, 0.0, -0.010990142822265625, -0.010990142822265625, -0.03296661376953125, -0.010990142822265625, -0.010990142822265625, -0.03296661376953125, -0.03296661376953125]
|
|
||||||
rate = rospy.Rate(50)
|
|
||||||
for idx in range(len(actions)):
|
|
||||||
action = actions[idx]['action'].detach().cpu().numpy()
|
|
||||||
velocity = velocitys[idx]['observation.velocity'].detach().cpu().numpy()
|
|
||||||
effort = efforts[idx]['observation.effort'].detach().cpu().numpy()
|
|
||||||
if(rospy.is_shutdown()):
|
|
||||||
break
|
|
||||||
|
|
||||||
new_actions = np.linspace(last_action, action, 5) # 插值
|
|
||||||
new_velocitys = np.linspace(last_velocity, velocity, 5) # 插值
|
|
||||||
new_efforts = np.linspace(last_effort, effort, 5) # 插值
|
|
||||||
last_action = action
|
|
||||||
last_velocity = velocity
|
|
||||||
last_effort = effort
|
|
||||||
for act in new_actions:
|
|
||||||
print(np.round(act[:7], 4))
|
|
||||||
cur_timestamp = rospy.Time.now() # 设置时间戳
|
|
||||||
joint_state_msg.header.stamp = cur_timestamp
|
|
||||||
|
|
||||||
joint_state_msg.position = act[:7]
|
|
||||||
joint_state_msg.velocity = last_velocity[:7]
|
|
||||||
joint_state_msg.effort = last_effort[:7]
|
|
||||||
master_arm_left_publisher.publish(joint_state_msg)
|
|
||||||
|
|
||||||
joint_state_msg.position = act[7:]
|
|
||||||
joint_state_msg.velocity = last_velocity[:7]
|
|
||||||
joint_state_msg.effort = last_effort[7:]
|
|
||||||
master_arm_right_publisher.publish(joint_state_msg)
|
|
||||||
|
|
||||||
if(rospy.is_shutdown()):
|
|
||||||
break
|
|
||||||
rate.sleep()
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == '__main__':
|
|
||||||
parser = argparse.ArgumentParser()
|
parser = argparse.ArgumentParser()
|
||||||
# parser.add_argument('--master_arm_left_topic', action='store', type=str, help='master_arm_left_topic',
|
|
||||||
# default='/master/joint_left', required=False)
|
|
||||||
# parser.add_argument('--master_arm_right_topic', action='store', type=str, help='master_arm_right_topic',
|
|
||||||
# default='/master/joint_right', required=False)
|
|
||||||
|
|
||||||
|
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
args.repo_id = "tangger/test"
|
args.repo_id = "tangger/test"
|
||||||
args.root = "/home/ubuntu/LYT/aloha_lerobot/data1"
|
args.root = "/home/ubuntu/LYT/aloha_lerobot/data1"
|
||||||
args.episode = 1 # replay episode
|
args.episode = 1 # replay episode
|
||||||
args.master_arm_left_topic = "/master/joint_left"
|
|
||||||
args.master_arm_right_topic = "/master/joint_right"
|
|
||||||
args.fps = 30
|
args.fps = 30
|
||||||
|
args.use_depth_image = False
|
||||||
|
args.use_base = False
|
||||||
|
return args
|
||||||
|
|
||||||
main(args)
|
|
||||||
# python collect_data.py --max_timesteps 500 --is_compress --episode_idx 0
|
if __name__ == '__main__':
|
||||||
|
args = get_arguments()
|
||||||
|
|
||||||
|
# Initialize ROS node
|
||||||
|
rospy.init_node("replay_node")
|
||||||
|
|
||||||
|
# Create robot instance using factory pattern
|
||||||
|
robot = RobotFactory.create(config_file="/home/ubuntu/LYT/lerobot_aloha/lerobot_aloha/configs/agilex.yaml", args=args)
|
||||||
|
|
||||||
|
# Replay the specified episode
|
||||||
|
replay(robot, args)
|
||||||
|
|||||||
Reference in New Issue
Block a user