add mot (#83)
* add mot * update libero example * translate zh to en * fix load model from hf * lint * lint --------- Co-authored-by: yangping <yangping@x2robot.com>
This commit is contained in:
suolyer
@@ -0,0 +1,323 @@
|
||||
"""Utils for evaluating policies in LIBERO simulation environments."""
|
||||
|
||||
import math
|
||||
import os
|
||||
from enum import Enum
|
||||
import imageio
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
import torch
|
||||
from transformers import BatchFeature
|
||||
import random
|
||||
import time
|
||||
|
||||
from libero.libero import get_libero_path
|
||||
from libero.libero.envs import OffScreenRenderEnv
|
||||
|
||||
|
||||
# Define task suite constants
|
||||
class TaskSuite(str, Enum):
|
||||
LIBERO_SPATIAL = "libero_spatial"
|
||||
LIBERO_OBJECT = "libero_object"
|
||||
LIBERO_GOAL = "libero_goal"
|
||||
LIBERO_10 = "libero_10"
|
||||
LIBERO_90 = "libero_90"
|
||||
|
||||
|
||||
# Define max steps for each task suite
|
||||
TASK_MAX_STEPS = {
|
||||
TaskSuite.LIBERO_SPATIAL: 220, # longest training demo has 193 steps
|
||||
TaskSuite.LIBERO_OBJECT: 280, # longest training demo has 254 steps
|
||||
TaskSuite.LIBERO_GOAL: 300, # longest training demo has 270 steps
|
||||
TaskSuite.LIBERO_10: 520, # longest training demo has 505 steps
|
||||
TaskSuite.LIBERO_90: 400, # longest training demo has 373 steps
|
||||
}
|
||||
|
||||
|
||||
# Initialize important constants
|
||||
ACTION_DIM = 7
|
||||
DATE = time.strftime("%Y_%m_%d")
|
||||
DATE_TIME = time.strftime("%Y_%m_%d-%H_%M_%S")
|
||||
|
||||
# Configure NumPy print settings
|
||||
np.set_printoptions(formatter={"float": lambda x: "{0:0.3f}".format(x)})
|
||||
|
||||
|
||||
def set_seed_everywhere(seed: int) -> None:
|
||||
"""
|
||||
Set random seed for all random number generators for reproducibility.
|
||||
|
||||
Args:
|
||||
seed: The random seed to use
|
||||
"""
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
np.random.seed(seed)
|
||||
random.seed(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
os.environ["PYTHONHASHSEED"] = str(seed)
|
||||
|
||||
|
||||
def normalize_gripper_action(action: np.ndarray, binarize: bool = True) -> np.ndarray:
|
||||
"""
|
||||
Normalize gripper action from [0,1] to [-1,+1] range.
|
||||
|
||||
This is necessary for some environments because the dataset wrapper
|
||||
standardizes gripper actions to [0,1]. Note that unlike the other action
|
||||
dimensions, the gripper action is not normalized to [-1,+1] by default.
|
||||
|
||||
Normalization formula: y = 2 * (x - orig_low) / (orig_high - orig_low) - 1
|
||||
|
||||
Args:
|
||||
action: Action array with gripper action in the last dimension
|
||||
binarize: Whether to binarize gripper action to -1 or +1
|
||||
|
||||
Returns:
|
||||
np.ndarray: Action array with normalized gripper action
|
||||
"""
|
||||
# Create a copy to avoid modifying the original
|
||||
normalized_action = action.copy()
|
||||
|
||||
# Normalize the last action dimension to [-1,+1]
|
||||
orig_low, orig_high = 0.0, 1.0
|
||||
normalized_action[..., -1] = (
|
||||
2 * (normalized_action[..., -1] - orig_low) / (orig_high - orig_low) - 1
|
||||
)
|
||||
|
||||
if binarize:
|
||||
# Binarize to -1 or +1
|
||||
normalized_action[..., -1] = np.sign(normalized_action[..., -1])
|
||||
|
||||
return normalized_action
|
||||
|
||||
|
||||
def invert_gripper_action(action: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Flip the sign of the gripper action (last dimension of action vector).
|
||||
|
||||
This is necessary for environments where -1 = open, +1 = close, since
|
||||
the RLDS dataloader aligns gripper actions such that 0 = close, 1 = open.
|
||||
|
||||
Args:
|
||||
action: Action array with gripper action in the last dimension
|
||||
|
||||
Returns:
|
||||
np.ndarray: Action array with inverted gripper action
|
||||
"""
|
||||
# Create a copy to avoid modifying the original
|
||||
inverted_action = action.copy()
|
||||
|
||||
# Invert the gripper action
|
||||
inverted_action[..., -1] *= -1.0
|
||||
|
||||
return inverted_action
|
||||
|
||||
|
||||
def move_to_cuda(obj, device="cuda"):
|
||||
if isinstance(obj, torch.Tensor):
|
||||
return obj.to(device)
|
||||
elif isinstance(obj, (dict, BatchFeature)):
|
||||
return {k: move_to_cuda(v, device) for k, v in obj.items()}
|
||||
elif isinstance(obj, list):
|
||||
return [move_to_cuda(v, device) for v in obj]
|
||||
elif isinstance(obj, tuple):
|
||||
return tuple(move_to_cuda(v, device) for v in obj)
|
||||
else:
|
||||
return obj
|
||||
|
||||
|
||||
def get_libero_env(task, model_family, resolution=256, seed=7):
|
||||
"""Initializes and returns the LIBERO environment, along with the task description."""
|
||||
task_description = task.language
|
||||
task_bddl_file = os.path.join(
|
||||
get_libero_path("bddl_files"), task.problem_folder, task.bddl_file
|
||||
)
|
||||
env_args = {
|
||||
"bddl_file_name": task_bddl_file,
|
||||
"camera_heights": resolution,
|
||||
"camera_widths": resolution,
|
||||
}
|
||||
env = OffScreenRenderEnv(**env_args)
|
||||
env.seed(
|
||||
seed
|
||||
) # IMPORTANT: seed seems to affect object positions even when using fixed initial state
|
||||
return env, task_description
|
||||
|
||||
|
||||
def get_libero_dummy_action(model_family: str):
|
||||
"""Get dummy/no-op action, used to roll out the simulation while the robot does nothing."""
|
||||
return [0, 0, 0, 0, 0, 0, -1]
|
||||
|
||||
|
||||
def get_libero_image(obs):
|
||||
"""Extracts third-person image from observations and preprocesses it."""
|
||||
img = obs["agentview_image"]
|
||||
img = img[::-1, ::-1] # IMPORTANT: rotate 180 degrees to match train preprocessing
|
||||
return img
|
||||
|
||||
|
||||
def get_libero_wrist_image(obs):
|
||||
"""Extracts wrist camera image from observations and preprocesses it."""
|
||||
img = obs["robot0_eye_in_hand_image"]
|
||||
img = img[::-1, ::-1] # IMPORTANT: rotate 180 degrees to match train preprocessing
|
||||
return img
|
||||
|
||||
|
||||
def save_rollout_video(
|
||||
rollout_dir,
|
||||
rollout_images,
|
||||
idx,
|
||||
success,
|
||||
task_description,
|
||||
log_file=None,
|
||||
model_family="openvla_oft",
|
||||
):
|
||||
"""Saves an MP4 replay of an episode."""
|
||||
processed_task_description = (
|
||||
task_description.lower()
|
||||
.replace(" ", "_")
|
||||
.replace("\n", "_")
|
||||
.replace(".", "_")[:50]
|
||||
)
|
||||
mp4_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}.mp4"
|
||||
video_writer = imageio.get_writer(mp4_path, fps=30)
|
||||
for img in rollout_images:
|
||||
video_writer.append_data(img)
|
||||
video_writer.close()
|
||||
print(f"Saved rollout MP4 at path {mp4_path}")
|
||||
if log_file is not None:
|
||||
log_file.write(f"Saved rollout MP4 at path {mp4_path}\n")
|
||||
return mp4_path
|
||||
|
||||
|
||||
def save_rollout_data(
|
||||
rollout_dir,
|
||||
rollout_data,
|
||||
idx,
|
||||
success,
|
||||
task_description,
|
||||
log_file=None,
|
||||
model_family="openvla_oft",
|
||||
):
|
||||
"""
|
||||
Saves an NPY file of the rollout data.
|
||||
"""
|
||||
|
||||
# Process task description to make it suitable for filename
|
||||
processed_task_description = (
|
||||
task_description.lower()
|
||||
.replace(" ", "_")
|
||||
.replace("\n", "_")
|
||||
.replace(".", "_")[:50]
|
||||
)
|
||||
|
||||
# Build .npy file path
|
||||
npy_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--action.npy"
|
||||
|
||||
# Save rollout_data as .npy file
|
||||
np.save(npy_path, rollout_data)
|
||||
print(f"Saved rollout data at path {npy_path}")
|
||||
|
||||
fig, axes = plt.subplots(nrows=1, ncols=7, figsize=(20, 3))
|
||||
|
||||
titles = ["x", "y", "z", "roll", "pitch", "yaw", "grasp"]
|
||||
|
||||
for i in range(rollout_data.shape[1]):
|
||||
ax = axes[i] # Select the i-th subplot
|
||||
ax.plot(rollout_data[:, i], label=f"Feature {i+1}") # Plot line chart
|
||||
ax.set_title(titles[i]) # Set subplot title
|
||||
ax.set_xlabel("Time in one episode") # Set x-axis label
|
||||
|
||||
axes[-1].legend(["predicted action"], loc="upper right")
|
||||
|
||||
plt.tight_layout()
|
||||
png_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--action.png"
|
||||
plt.savefig(png_path, dpi=300) # Save image as PNG file
|
||||
|
||||
# If log file is provided, record the save path
|
||||
if log_file is not None:
|
||||
log_file.write(f"Saved rollout data at path {npy_path}\n")
|
||||
|
||||
return npy_path
|
||||
|
||||
|
||||
def save_rollout_observation(
|
||||
rollout_dir,
|
||||
rollout_data,
|
||||
idx,
|
||||
success,
|
||||
task_description,
|
||||
log_file=None,
|
||||
model_family="openvla_oft",
|
||||
):
|
||||
"""
|
||||
Saves an NPY file of the rollout data.
|
||||
"""
|
||||
|
||||
# Process task description to make it suitable for filename
|
||||
processed_task_description = (
|
||||
task_description.lower()
|
||||
.replace(" ", "_")
|
||||
.replace("\n", "_")
|
||||
.replace(".", "_")[:50]
|
||||
)
|
||||
|
||||
# Build .npy file path
|
||||
npy_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--observation.npy"
|
||||
|
||||
# Save rollout_data as .npy file
|
||||
np.save(npy_path, rollout_data)
|
||||
print(f"Saved rollout data at path {npy_path}")
|
||||
|
||||
if rollout_data.shape[1] == 7:
|
||||
fig, axes = plt.subplots(nrows=1, ncols=7, figsize=(20, 3))
|
||||
titles = ["x", "y", "z", "roll", "pitch", "yaw", "grasp"]
|
||||
else:
|
||||
fig, axes = plt.subplots(nrows=1, ncols=8, figsize=(20, 3))
|
||||
titles = ["x", "y", "z", "roll", "pitch", "yaw", "-", "grasp"]
|
||||
|
||||
for i in range(rollout_data.shape[1]):
|
||||
ax = axes[i] # Select the i-th subplot
|
||||
ax.plot(rollout_data[:, i], label=f"Feature {i+1}") # Plot line chart
|
||||
ax.set_title(titles[i]) # Set subplot title
|
||||
ax.set_xlabel("Time in one episode") # Set x-axis label
|
||||
|
||||
axes[-1].legend(["predicted action"], loc="upper right")
|
||||
|
||||
plt.tight_layout()
|
||||
png_path = f"{rollout_dir}/episode={idx}--success={success}--task={processed_task_description}--observation.png"
|
||||
plt.savefig(png_path, dpi=300) # Save image as PNG file
|
||||
|
||||
# If log file is provided, record the save path
|
||||
if log_file is not None:
|
||||
log_file.write(f"Saved rollout data at path {npy_path}\n")
|
||||
|
||||
return npy_path
|
||||
|
||||
|
||||
def quat2axisangle(quat):
|
||||
"""
|
||||
Copied from robosuite: https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55
|
||||
|
||||
Converts quaternion to axis-angle format.
|
||||
Returns a unit vector direction scaled by its angle in radians.
|
||||
|
||||
Args:
|
||||
quat (np.array): (x,y,z,w) vec4 float angles
|
||||
|
||||
Returns:
|
||||
np.array: (ax,ay,az) axis-angle exponential coordinates
|
||||
"""
|
||||
# clip quaternion
|
||||
if quat[3] > 1.0:
|
||||
quat[3] = 1.0
|
||||
elif quat[3] < -1.0:
|
||||
quat[3] = -1.0
|
||||
|
||||
den = np.sqrt(1.0 - quat[3] * quat[3])
|
||||
if math.isclose(den, 0.0):
|
||||
# This is (close to) a zero degree rotation, immediately return
|
||||
return np.zeros(3)
|
||||
|
||||
return (quat[:3] * 2.0 * math.acos(quat[3])) / den
|
||||
Reference in New Issue
Block a user