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,744 @@
|
||||
import os
|
||||
import json
|
||||
import numpy as np
|
||||
from typing import Dict, Any, Tuple, List
|
||||
from libero.libero import benchmark
|
||||
from wall_x.infer.env import BaseEnv, InferConfig
|
||||
from wall_x.serving.policy.wall_x_policy import WallXPolicy
|
||||
|
||||
from wall_x.infer.base_dataclass import RobotStateActionData
|
||||
from wall_x.infer.utils_libero import (
|
||||
get_libero_env,
|
||||
get_libero_dummy_action,
|
||||
get_libero_image,
|
||||
get_libero_wrist_image,
|
||||
quat2axisangle,
|
||||
TaskSuite,
|
||||
save_rollout_video,
|
||||
)
|
||||
from robosuite.wrappers import VisualizationWrapper
|
||||
|
||||
|
||||
def _create_libero_env_standalone(
|
||||
task_id: int,
|
||||
task_suite_name: str,
|
||||
model_family: str = "wallx",
|
||||
resolution: int = 256,
|
||||
seed: int = 7,
|
||||
) -> Any:
|
||||
"""
|
||||
Standalone function to create a Libero environment, independent of LiberoRobotEnv instance.
|
||||
Used for creating environments in subprocess during multi-batch inference,
|
||||
avoiding serialization of large objects containing the model.
|
||||
|
||||
Args:
|
||||
task_id: Task ID
|
||||
task_suite_name: Task suite name
|
||||
model_family: Model family
|
||||
resolution: Resolution
|
||||
seed: Random seed
|
||||
|
||||
Returns:
|
||||
Environment instance
|
||||
"""
|
||||
from libero.libero import benchmark
|
||||
|
||||
# Get task suite and task
|
||||
benchmark_dict = benchmark.get_benchmark_dict()
|
||||
task_suite = benchmark_dict[task_suite_name]()
|
||||
task = task_suite.get_task(task_id)
|
||||
|
||||
# Create environment
|
||||
env, _ = get_libero_env(
|
||||
task,
|
||||
model_family=model_family,
|
||||
resolution=resolution,
|
||||
seed=seed,
|
||||
)
|
||||
|
||||
# Wrap environment
|
||||
env.env = VisualizationWrapper(env.env)
|
||||
env.env.set_visualization_setting(setting="grippers", visible=False)
|
||||
|
||||
return env
|
||||
|
||||
|
||||
class LiberoRobotEnv(BaseEnv):
|
||||
def __init__(
|
||||
self,
|
||||
config: InferConfig,
|
||||
task_suite_name: str = TaskSuite.LIBERO_SPATIAL,
|
||||
initial_states_path: str = "DEFAULT",
|
||||
rollout_dir: str = "./rollouts",
|
||||
model_family: str = "wallx",
|
||||
resolution: int = 256,
|
||||
seed: int = 7,
|
||||
):
|
||||
|
||||
super().__init__(config)
|
||||
self.logger.info(
|
||||
f"Initializing LiberoRobotEnv (Stateless), task suite: {task_suite_name}"
|
||||
)
|
||||
|
||||
self.model = self._register_model()
|
||||
self.model_family = model_family
|
||||
self.resolution = resolution
|
||||
self.seed = seed
|
||||
|
||||
self.logger.info("Importing Libero and related utils...")
|
||||
self.RobotStateActionData = RobotStateActionData
|
||||
|
||||
self.rollout_dir = os.path.join(rollout_dir, task_suite_name)
|
||||
os.makedirs(self.rollout_dir, exist_ok=True)
|
||||
|
||||
if save_rollout_video is not None:
|
||||
self.save_rollout_video = save_rollout_video
|
||||
else:
|
||||
self.save_rollout_video = None
|
||||
self.logger.warning("save_rollout_video not found, video saving disabled.")
|
||||
|
||||
self.task_suite_name = task_suite_name
|
||||
benchmark_dict = benchmark.get_benchmark_dict()
|
||||
self.task_suite = benchmark_dict[self.task_suite_name]()
|
||||
self.num_tasks = self.task_suite.n_tasks
|
||||
|
||||
self.initial_states_path = initial_states_path
|
||||
self.all_initial_states = None
|
||||
if self.initial_states_path != "DEFAULT":
|
||||
try:
|
||||
with open(self.initial_states_path, "r") as f:
|
||||
self.all_initial_states = json.load(f)
|
||||
self.logger.info(
|
||||
f"Loaded custom initial states from {self.initial_states_path}"
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to load initial states file: {e}")
|
||||
raise
|
||||
|
||||
def _register_model(self) -> WallXPolicy:
|
||||
|
||||
return WallXPolicy(
|
||||
model_path=self.config.model_path,
|
||||
train_config=self.config.train_config,
|
||||
action_tokenizer_path=self.config.action_tokenizer_path,
|
||||
action_dim=self.config.action_dim,
|
||||
agent_pos_dim=self.config.action_dim,
|
||||
pred_horizon=self.config.pred_horizon,
|
||||
camera_key=self.config.cam_names,
|
||||
predict_mode=self.config.predict_mode,
|
||||
)
|
||||
|
||||
def get_instruction(self, task_desc: str) -> str:
|
||||
return task_desc
|
||||
|
||||
def get_observation(self, raw_obs: Dict[str, Any]) -> Dict[str, Any]:
|
||||
if raw_obs is None:
|
||||
raise ValueError("Raw observation is None")
|
||||
|
||||
data_obj = self.RobotStateActionData(config=self.config)
|
||||
|
||||
pos = raw_obs["robot0_eef_pos"]
|
||||
rot = quat2axisangle(raw_obs["robot0_eef_quat"])
|
||||
grip = raw_obs["robot0_gripper_qpos"][0:1]
|
||||
|
||||
data_obj.save_state_data_with_key(pos[None], "follow_right_ee_cartesian_pos")
|
||||
data_obj.save_state_data_with_key(rot[None], "follow_right_ee_rotation")
|
||||
data_obj.save_state_data_with_key(grip[None], "follow_right_gripper")
|
||||
data_obj.dof_mask = self._get_dof_mask()
|
||||
|
||||
face_view = get_libero_image(raw_obs)
|
||||
right_wrist_view = get_libero_wrist_image(raw_obs)
|
||||
|
||||
return {
|
||||
"robot_state_action_data": data_obj,
|
||||
"face_view": face_view,
|
||||
"right_wrist_view": right_wrist_view,
|
||||
}
|
||||
|
||||
def apply_action(
|
||||
self, input_data: Dict[str, Any], env: Any = None, replay_images: list = None
|
||||
) -> bool:
|
||||
if env is None:
|
||||
raise ValueError(
|
||||
"In Stateless mode, apply_action must be called with explicit 'env' parameter"
|
||||
)
|
||||
|
||||
action_data = input_data["robot_state_action_data"]
|
||||
right_arm_traj = self._get_right_arm_action(action_data)
|
||||
while (
|
||||
right_arm_traj is not None
|
||||
and right_arm_traj.ndim > 2
|
||||
and right_arm_traj.shape[0] == 1
|
||||
):
|
||||
right_arm_traj = right_arm_traj.squeeze(0)
|
||||
|
||||
done = False
|
||||
t = 0
|
||||
|
||||
try:
|
||||
for i in range(len(right_arm_traj)):
|
||||
if done:
|
||||
break
|
||||
|
||||
action_7d = right_arm_traj[i]
|
||||
obs, reward, done, info = env.step(action_7d)
|
||||
t += 1
|
||||
|
||||
if obs is not None and replay_images is not None:
|
||||
replay_images.append(get_libero_image(obs))
|
||||
|
||||
input_data["_last_obs"] = obs
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Env step error: {e}")
|
||||
return False # Error treated as failure
|
||||
|
||||
return done, t
|
||||
|
||||
def apply_action_batch(
|
||||
self,
|
||||
vec_env: Any,
|
||||
trajectories: List[np.ndarray],
|
||||
active_indices: List[int],
|
||||
status_list: List[Dict[str, Any]],
|
||||
model_outputs: List[Dict[str, Any]],
|
||||
) -> None:
|
||||
"""
|
||||
Execute action trajectories in parallel batch.
|
||||
|
||||
Uses SubprocVectorEnv to execute actions in parallel for all active environments.
|
||||
Integrates vec_env.step(batch_actions, id=still_active) in this function.
|
||||
"""
|
||||
if not trajectories:
|
||||
return
|
||||
|
||||
max_traj_len = max(len(traj) for traj in trajectories)
|
||||
if max_traj_len == 0:
|
||||
return
|
||||
|
||||
for step_idx in range(max_traj_len):
|
||||
# Check if there are still active environments
|
||||
still_active = [
|
||||
idx
|
||||
for idx in active_indices
|
||||
if (not status_list[idx]["done"])
|
||||
and status_list[idx]["count"] > 0
|
||||
and step_idx < len(trajectories[active_indices.index(idx)])
|
||||
]
|
||||
if not still_active:
|
||||
break
|
||||
|
||||
# Build batch actions (only includes actions for still_active environments)
|
||||
batch_actions = []
|
||||
for idx in still_active:
|
||||
traj_idx = active_indices.index(idx)
|
||||
action_7d = trajectories[traj_idx][step_idx]
|
||||
# Ensure action_7d is numpy array or list
|
||||
if isinstance(action_7d, np.ndarray):
|
||||
batch_actions.append(action_7d)
|
||||
else:
|
||||
batch_actions.append(np.array(action_7d))
|
||||
|
||||
# Convert to numpy array with shape (batch_size, action_dim)
|
||||
batch_actions = np.array(batch_actions)
|
||||
|
||||
# Execute step in parallel (only for still_active environments)
|
||||
obs_list, reward_list, done_list, info_list = vec_env.step(
|
||||
batch_actions, id=still_active
|
||||
)
|
||||
|
||||
# Process returned results
|
||||
if obs_list.dtype == object:
|
||||
obs_list = list(obs_list)
|
||||
else:
|
||||
obs_list = [obs_list[i] for i in range(len(obs_list))]
|
||||
done_list = [bool(done_list[i]) for i in range(len(done_list))]
|
||||
|
||||
# Update each environment's status
|
||||
for i, idx in enumerate(still_active):
|
||||
st = status_list[idx]
|
||||
obs = obs_list[i]
|
||||
done = done_list[i]
|
||||
|
||||
if obs is not None:
|
||||
st["current_obs"] = obs
|
||||
if st["replay_images"] is not None:
|
||||
st["replay_images"].append(get_libero_image(obs))
|
||||
|
||||
st["count"] -= 1
|
||||
st["success"] = done
|
||||
st["done"] = done or st["count"] <= 0
|
||||
|
||||
# Update model_output's _last_obs
|
||||
model_outputs[active_indices.index(idx)]["_last_obs"] = obs
|
||||
|
||||
def get_task_info(self, task_id: int) -> Tuple[str, Any]:
|
||||
"""
|
||||
Get task information (task description and initial states) without creating environment.
|
||||
|
||||
Returns:
|
||||
Tuple[str, Any]: (task_desc, default_initial_states)
|
||||
"""
|
||||
if task_id < 0 or task_id >= self.num_tasks:
|
||||
raise ValueError(f"Invalid task ID: {task_id}")
|
||||
|
||||
task = self.task_suite.get_task(task_id)
|
||||
task_desc = task.language
|
||||
default_initial_states = self.task_suite.get_task_init_states(task_id)
|
||||
|
||||
return task_desc, default_initial_states
|
||||
|
||||
def create_env_for_task(self, task_id: int) -> Tuple[Any, str, Any]:
|
||||
if task_id < 0 or task_id >= self.num_tasks:
|
||||
raise ValueError(f"Invalid task ID: {task_id}")
|
||||
|
||||
task = self.task_suite.get_task(task_id)
|
||||
default_initial_states = self.task_suite.get_task_init_states(task_id)
|
||||
|
||||
env, task_desc = get_libero_env(
|
||||
task,
|
||||
model_family=self.model_family,
|
||||
resolution=self.resolution,
|
||||
seed=self.seed,
|
||||
)
|
||||
|
||||
env.env = VisualizationWrapper(env.env)
|
||||
env.env.set_visualization_setting(setting="grippers", visible=False)
|
||||
|
||||
return env, task_desc, default_initial_states
|
||||
|
||||
def _get_initial_state_for_episode(
|
||||
self, task_desc: str, default_states: Any, episode_idx: int
|
||||
) -> np.ndarray:
|
||||
if self.initial_states_path == "DEFAULT":
|
||||
if default_states is None:
|
||||
raise ValueError("Default states missing")
|
||||
return default_states[episode_idx]
|
||||
else:
|
||||
if self.all_initial_states is None:
|
||||
raise ValueError("Custom states not loaded")
|
||||
initial_states_task_key = task_desc.replace(" ", "_")
|
||||
episode_key = f"demo_{episode_idx}"
|
||||
if not self.all_initial_states[initial_states_task_key][episode_key][
|
||||
"success"
|
||||
]:
|
||||
raise ValueError(f"Expert demo failed for {episode_key}")
|
||||
return np.array(
|
||||
self.all_initial_states[initial_states_task_key][episode_key][
|
||||
"initial_state"
|
||||
]
|
||||
)
|
||||
|
||||
def reset_env(
|
||||
self, env: Any, task_desc: str, default_states: Any, episode_idx: int
|
||||
) -> Any:
|
||||
try:
|
||||
if episode_idx >= 0:
|
||||
state = self._get_initial_state_for_episode(
|
||||
task_desc, default_states, episode_idx
|
||||
)
|
||||
obs = env.set_init_state(state)
|
||||
return obs
|
||||
else:
|
||||
return env.reset()
|
||||
except Exception as e:
|
||||
self.logger.error(f"Reset failed: {e}, falling back to default reset")
|
||||
return env.reset()
|
||||
|
||||
def _get_dof_mask(self):
|
||||
dof_config = self.config.train_config["dof_config"]
|
||||
total_dof = sum(dof_config.values())
|
||||
dof_mask = np.ones((1, self.config.action_horizon, total_dof))
|
||||
mask_keys = [
|
||||
"follow_left_ee_cartesian_pos",
|
||||
"follow_left_ee_rotation",
|
||||
"follow_left_gripper",
|
||||
"head_actions",
|
||||
"height",
|
||||
"velocity_decomposed",
|
||||
]
|
||||
start_idx = 0
|
||||
for key, dof_size in dof_config.items():
|
||||
if key in mask_keys:
|
||||
dof_mask[:, :, start_idx : start_idx + dof_size] = 0
|
||||
start_idx += dof_size
|
||||
return dof_mask
|
||||
|
||||
def _get_right_arm_action(
|
||||
self, robot_state_action_data: RobotStateActionData
|
||||
) -> np.ndarray:
|
||||
right_ee_cartesian_pos = robot_state_action_data.data[
|
||||
"action_right_ee_cartesian_pos"
|
||||
]
|
||||
right_ee_rotation = robot_state_action_data.data["action_right_ee_rotation"]
|
||||
right_gripper = robot_state_action_data.data["action_right_gripper"]
|
||||
return np.concatenate(
|
||||
[right_ee_cartesian_pos, right_ee_rotation, right_gripper], axis=1
|
||||
)
|
||||
|
||||
def _get_left_arm_action(
|
||||
self, robot_state_action_data: RobotStateActionData
|
||||
) -> np.ndarray:
|
||||
left_ee_cartesian_pos = robot_state_action_data.data[
|
||||
"action_left_ee_cartesian_pos"
|
||||
]
|
||||
left_ee_rotation = robot_state_action_data.data["action_left_ee_rotation"]
|
||||
left_gripper = robot_state_action_data.data["action_left_gripper"]
|
||||
return np.concatenate(
|
||||
[left_ee_cartesian_pos, left_ee_rotation, left_gripper], axis=1
|
||||
)
|
||||
|
||||
def _save_rollout(
|
||||
self,
|
||||
replay_images: List[np.ndarray],
|
||||
success: bool,
|
||||
task_id: int,
|
||||
task_desc: str,
|
||||
episode_idx: int,
|
||||
):
|
||||
if not self.save_rollout_video or not replay_images:
|
||||
return
|
||||
try:
|
||||
task_name_safe = task_desc.replace(" ", "_").replace(".", "")
|
||||
filename = f"{episode_idx}{'_SUCCESS' if success else '_FAILURE'}--_{task_name_safe}.mp4"
|
||||
self.save_rollout_video(
|
||||
self.rollout_dir,
|
||||
replay_images,
|
||||
filename,
|
||||
success=success,
|
||||
task_description=task_desc,
|
||||
log_file=None,
|
||||
model_family=self.model_family,
|
||||
)
|
||||
self.logger.info(f"Saved video: {filename}")
|
||||
except Exception as e:
|
||||
self.logger.error(f"Save video failed: {e}")
|
||||
|
||||
def run_infer_flow_action(
|
||||
self,
|
||||
env: Any,
|
||||
task_id: int,
|
||||
task_desc: str,
|
||||
default_initial_states: Any,
|
||||
episode_idx: int,
|
||||
max_infer_times: int = 5,
|
||||
num_steps_wait: int = 10,
|
||||
) -> bool:
|
||||
replay_images = []
|
||||
num_steps = 0
|
||||
done = False
|
||||
count = max_infer_times
|
||||
|
||||
current_obs = self.reset_env(
|
||||
env, task_desc, default_initial_states, episode_idx
|
||||
)
|
||||
if current_obs is None:
|
||||
return False
|
||||
|
||||
while num_steps < num_steps_wait:
|
||||
obs, reward, done, info = env.step(
|
||||
get_libero_dummy_action(self.model_family)
|
||||
)
|
||||
num_steps += 1
|
||||
if obs is not None:
|
||||
current_obs = obs
|
||||
|
||||
while not done and count > 0:
|
||||
try:
|
||||
model_input = self.get_observation(current_obs)
|
||||
instruction = self.get_instruction(task_desc)
|
||||
model_input["prompt"] = instruction
|
||||
model_input["dataset_names"] = "libero_all"
|
||||
|
||||
state = np.concatenate(
|
||||
[
|
||||
model_input["robot_state_action_data"].data[
|
||||
"state_right_ee_cartesian_pos"
|
||||
],
|
||||
model_input["robot_state_action_data"].data[
|
||||
"state_right_ee_rotation"
|
||||
],
|
||||
model_input["robot_state_action_data"].data[
|
||||
"state_right_gripper"
|
||||
],
|
||||
],
|
||||
axis=-1,
|
||||
)
|
||||
|
||||
model_input["state"] = state
|
||||
model_output = self.model.infer(model_input)
|
||||
|
||||
model_output["robot_state_action_data"] = model_input[
|
||||
"robot_state_action_data"
|
||||
]
|
||||
model_output["robot_state_action_data"].save_action_data(
|
||||
model_output["predict_action"]
|
||||
)
|
||||
|
||||
model_output["_last_obs"] = None
|
||||
|
||||
done, delta_t = self.apply_action(
|
||||
model_output, env=env, replay_images=replay_images
|
||||
)
|
||||
|
||||
if model_output.get("_last_obs") is not None:
|
||||
current_obs = model_output["_last_obs"]
|
||||
|
||||
count -= delta_t
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Episode Error: {e}")
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
break
|
||||
|
||||
success = done
|
||||
if count <= 0 and not done:
|
||||
self.logger.warning(
|
||||
f"Timeout: reached {max_infer_times} steps without success."
|
||||
)
|
||||
success = False
|
||||
|
||||
self._save_rollout(replay_images, success, task_id, task_desc, episode_idx)
|
||||
return success
|
||||
|
||||
def run_infer_flow_action_batch(
|
||||
self,
|
||||
vec_env: Any,
|
||||
task_ids: List[int] = None,
|
||||
task_descs: List[str] = None,
|
||||
default_initial_states_list: List[Any] = None,
|
||||
episode_indices: List[int] = None,
|
||||
max_infer_times: int = 5,
|
||||
num_steps_wait: int = 10,
|
||||
) -> List[bool]:
|
||||
"""
|
||||
Support batch inference: model inference in parallel (batch), environment execution in parallel (SubprocVectorEnv).
|
||||
|
||||
Uses SubprocVectorEnv to run environments in subprocess during multi-batch inference, all environments execute actions in parallel.
|
||||
|
||||
Returns a list of success flags for each sample.
|
||||
"""
|
||||
if vec_env is None:
|
||||
raise ValueError("vec_env must be specified")
|
||||
batch_size = len(vec_env)
|
||||
if task_ids is not None:
|
||||
assert len(task_ids) == batch_size, "task_ids length must match envs"
|
||||
if episode_indices is not None:
|
||||
assert (
|
||||
len(episode_indices) == batch_size
|
||||
), "episode_indices length must match envs"
|
||||
|
||||
status_list = []
|
||||
for i in range(batch_size):
|
||||
status_list.append(
|
||||
{
|
||||
"vec_env": vec_env,
|
||||
"env_id": i, # Index in vec_env
|
||||
"task_desc": task_descs[i],
|
||||
"replay_images": [],
|
||||
"num_steps": 0,
|
||||
"done": False, # Whether episode has ended
|
||||
"success": False, # Whether successfully completed
|
||||
"count": max_infer_times,
|
||||
"current_obs": None,
|
||||
"default_states": None,
|
||||
}
|
||||
)
|
||||
|
||||
# Initialize/reset: Use SubprocVectorEnv to batch set initial states
|
||||
init_states_to_set = []
|
||||
for i in range(batch_size):
|
||||
task_desc = task_descs[i]
|
||||
default_states = default_initial_states_list[i]
|
||||
status_list[i]["default_states"] = default_states
|
||||
ep_i = episode_indices[i]
|
||||
init_state = self._get_initial_state_for_episode(
|
||||
task_desc, default_states, ep_i
|
||||
)
|
||||
init_states_to_set.append(init_state)
|
||||
|
||||
# Batch set initial states
|
||||
try:
|
||||
obs_list = vec_env.set_init_state(init_states_to_set)
|
||||
if obs_list.dtype == object:
|
||||
obs_list = list(obs_list)
|
||||
else:
|
||||
obs_list = [obs_list[i] for i in range(len(obs_list))]
|
||||
|
||||
for i, obs in enumerate(obs_list):
|
||||
if obs is None:
|
||||
raise ValueError(
|
||||
f"Reset environment returned None, task_id: {task_ids[i]}, episode_idx: {episode_indices[i]}"
|
||||
)
|
||||
status_list[i]["current_obs"] = obs
|
||||
status_list[i]["done"] = False
|
||||
status_list[i]["success"] = False
|
||||
status_list[i]["count"] = max_infer_times
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to batch set initial states: {e}")
|
||||
raise
|
||||
|
||||
# Warmup steps (batch execution)
|
||||
dummy_action = get_libero_dummy_action(self.model_family)
|
||||
dummy_actions = np.array([dummy_action] * batch_size)
|
||||
for _ in range(num_steps_wait):
|
||||
obs_list, _, done_list, _ = vec_env.step(dummy_actions)
|
||||
# Update current_obs
|
||||
if obs_list.dtype == object:
|
||||
obs_list = list(obs_list)
|
||||
else:
|
||||
obs_list = [obs_list[i] for i in range(len(obs_list))]
|
||||
for i, obs in enumerate(obs_list):
|
||||
if obs is not None:
|
||||
status_list[i]["current_obs"] = obs
|
||||
|
||||
# Main loop: model parallel inference, environment parallel execution (SubprocVectorEnv)
|
||||
while any((not st["done"]) and st["count"] > 0 for st in status_list):
|
||||
active_indices = [
|
||||
idx
|
||||
for idx, st in enumerate(status_list)
|
||||
if (not st["done"]) and st["count"] > 0
|
||||
]
|
||||
if not active_indices:
|
||||
break
|
||||
print(f"Batch infer loop, active indices: {active_indices}")
|
||||
|
||||
observations = []
|
||||
instructions = []
|
||||
for idx in active_indices:
|
||||
st = status_list[idx]
|
||||
observations.append(self.get_observation(st["current_obs"]))
|
||||
instructions.append(self.get_instruction(st["task_desc"]))
|
||||
|
||||
# Model batch inference
|
||||
model_outputs = self.model.infer_flow_action_batch(
|
||||
observations, instructions
|
||||
)
|
||||
# Extract action trajectories for all active environments
|
||||
trajectories = []
|
||||
for out in model_outputs:
|
||||
action_data = out["robot_state_action_data"]
|
||||
right_arm_traj = self._get_right_arm_action(action_data)
|
||||
while (
|
||||
right_arm_traj is not None
|
||||
and right_arm_traj.ndim > 2
|
||||
and right_arm_traj.shape[0] == 1
|
||||
):
|
||||
right_arm_traj = right_arm_traj.squeeze(0)
|
||||
if right_arm_traj is None or len(right_arm_traj) == 0:
|
||||
# If trajectory is empty, create an empty trajectory
|
||||
right_arm_traj = np.array([]).reshape(0, 7)
|
||||
trajectories.append(right_arm_traj)
|
||||
|
||||
# Use apply_action_batch to execute action trajectories in parallel
|
||||
try:
|
||||
self.apply_action_batch(
|
||||
vec_env=vec_env,
|
||||
trajectories=trajectories,
|
||||
active_indices=active_indices,
|
||||
status_list=status_list,
|
||||
model_outputs=model_outputs,
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Batch parallel action error: {e}")
|
||||
# Mark all active environments as failed
|
||||
for idx in active_indices:
|
||||
status_list[idx]["done"] = True
|
||||
status_list[idx]["success"] = False
|
||||
|
||||
# Save replay and results
|
||||
success_list = []
|
||||
for i, st in enumerate(status_list):
|
||||
success = st.get("success", False)
|
||||
if st["count"] <= 0 and not st["success"]:
|
||||
self.logger.warning(
|
||||
f"Batch timeout: reached {max_infer_times} steps without success (idx {i})."
|
||||
)
|
||||
st["replay_images"] = st.get("replay_images", [])
|
||||
tid_i = task_ids[i]
|
||||
epi_i = episode_indices[i]
|
||||
self._save_rollout(
|
||||
st["replay_images"],
|
||||
success,
|
||||
tid_i,
|
||||
st["task_desc"],
|
||||
epi_i,
|
||||
)
|
||||
success_list.append(success)
|
||||
|
||||
return success_list
|
||||
|
||||
def run_infer_ar_action(
|
||||
self,
|
||||
env: Any,
|
||||
task_id: int,
|
||||
task_desc: str,
|
||||
default_initial_states: Any,
|
||||
episode_idx: int,
|
||||
max_infer_times: int = 10,
|
||||
num_steps_wait: int = 10,
|
||||
) -> bool:
|
||||
replay_images = []
|
||||
num_steps = 0
|
||||
done = False
|
||||
count = max_infer_times
|
||||
|
||||
current_obs = self.reset_env(
|
||||
env, task_desc, default_initial_states, episode_idx
|
||||
)
|
||||
if current_obs is None:
|
||||
self.logger.error("Environment reset returned None.")
|
||||
return False
|
||||
|
||||
while num_steps < num_steps_wait:
|
||||
obs, reward, done, info = env.step(
|
||||
get_libero_dummy_action(self.model_family)
|
||||
)
|
||||
num_steps += 1
|
||||
if obs is not None:
|
||||
current_obs = obs
|
||||
|
||||
while not done and count > 0:
|
||||
try:
|
||||
model_input = self.get_observation(current_obs)
|
||||
instruction = self.get_instruction(task_desc)
|
||||
|
||||
model_output = self.model.infer_ar_action(model_input, instruction)
|
||||
|
||||
model_output["_last_obs"] = None
|
||||
|
||||
done, delta_t = self.apply_action(
|
||||
model_output, env=env, replay_images=replay_images
|
||||
)
|
||||
|
||||
if model_output.get("_last_obs") is not None:
|
||||
current_obs = model_output["_last_obs"]
|
||||
else:
|
||||
if not done:
|
||||
self.logger.warning(
|
||||
"Did not receive new observation after apply_action, but episode is not done."
|
||||
)
|
||||
|
||||
count -= delta_t
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
f"AR Episode Run Error (Task {task_id}, Ep {episode_idx}): {e}"
|
||||
)
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
break
|
||||
|
||||
success = done
|
||||
if count <= 0 and not done:
|
||||
self.logger.warning(
|
||||
f"Timeout: AR policy reached {max_infer_times} steps without success."
|
||||
)
|
||||
success = False
|
||||
|
||||
self._save_rollout(replay_images, success, task_id, task_desc, episode_idx)
|
||||
|
||||
return success
|
||||
Reference in New Issue
Block a user