Logic
67 lines
2.5 KiB
Python
67 lines
2.5 KiB
Python
"""Capture first and last (success) frames from front camera for both sim tasks."""
|
|
import os
|
|
import time
|
|
import numpy as np
|
|
import cv2
|
|
from roboimi.envs.double_pos_ctrl_env import make_sim_env
|
|
from roboimi.demos.diana_policy import TestPickAndTransferPolicy
|
|
from roboimi.demos.diana_air_insert_policy import TestAirInsertPolicy
|
|
from roboimi.utils.act_ex_utils import sample_transfer_pose, sample_air_insert_socket_peg_state
|
|
|
|
OUTPUT_DIR = "sim_frames"
|
|
IMG_SIZE = (256, 256)
|
|
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
|
|
|
|
|
def capture_task(task_name, policy, task_state, episode_len):
|
|
env = make_sim_env(task_name, headless=True)
|
|
env.reset(task_state)
|
|
time.sleep(1)
|
|
|
|
env._update_camera_images_sync()
|
|
first_frame = cv2.resize(env.front, IMG_SIZE)
|
|
|
|
last_success_frame = None
|
|
max_reward = 0
|
|
|
|
for step in range(episode_len):
|
|
action = policy.predict(task_state, step)
|
|
env.step(action)
|
|
env._update_camera_images_sync()
|
|
if env.rew is not None and env.rew >= max_reward:
|
|
max_reward = env.rew
|
|
if env.rew == env.max_reward:
|
|
last_success_frame = cv2.resize(env.front, IMG_SIZE)
|
|
|
|
if last_success_frame is None:
|
|
last_success_frame = cv2.resize(env.front, IMG_SIZE)
|
|
print(f" [WARN] {task_name}: max_reward={max_reward}/{env.max_reward}, using last frame")
|
|
else:
|
|
print(f" {task_name}: success! max_reward={max_reward}")
|
|
|
|
return first_frame, last_success_frame
|
|
|
|
|
|
def main():
|
|
# Transfer task (episode_len=700, policy covers full range)
|
|
print("Running sim_transfer...")
|
|
task_state_transfer = sample_transfer_pose()
|
|
policy_transfer = TestPickAndTransferPolicy(inject_noise=False)
|
|
first_t, last_t = capture_task('sim_transfer', policy_transfer, task_state_transfer, 700)
|
|
cv2.imwrite(os.path.join(OUTPUT_DIR, "transfer_front_first.png"), first_t)
|
|
cv2.imwrite(os.path.join(OUTPUT_DIR, "transfer_front_last.png"), last_t)
|
|
|
|
# Air insert / peg-in-socket (policy EPISODE_END_T=600)
|
|
print("Running sim_air_insert_socket_peg...")
|
|
task_state_insert = sample_air_insert_socket_peg_state()
|
|
policy_insert = TestAirInsertPolicy(inject_noise=False)
|
|
first_i, last_i = capture_task('sim_air_insert_socket_peg', policy_insert, task_state_insert, 600)
|
|
cv2.imwrite(os.path.join(OUTPUT_DIR, "air_insert_front_first.png"), first_i)
|
|
cv2.imwrite(os.path.join(OUTPUT_DIR, "air_insert_front_last.png"), last_i)
|
|
|
|
print(f"\nDone! Images saved to {OUTPUT_DIR}/")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|