feat(vla): add SmolVLA conditioning and experiment artifacts

capture_sim_frames.py

@@ -0,0 +1,66 @@
+"""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()