From 5334e96b7e9917bdd332395721b0d794da541f61 Mon Sep 17 00:00:00 2001 From: selen-suyue Date: Sun, 1 Feb 2026 18:21:26 +0800 Subject: [PATCH] add wog and continual overlay --- README.md | 21 +++ configs/policy/wog.yaml | 31 +++++ main.py | 36 +++++ .../customize/policy_model/wog_model.py | 126 ++++++++++++++++++ maniunicon/utils/overlay_viewer.py | 80 +++++++++++ 5 files changed, 294 insertions(+) create mode 100644 configs/policy/wog.yaml create mode 100644 maniunicon/customize/policy_model/wog_model.py create mode 100644 maniunicon/utils/overlay_viewer.py diff --git a/README.md b/README.md index 6483168..98aed2a 100644 --- a/README.md +++ b/README.md @@ -296,6 +296,27 @@ class MyCustomPolicy(mp.Process): 3. Add configuration in `configs/policy/my_policy.yaml` +4. Deploy your policy/vla model: +``` +# Quick Deploy +python main.py robot=ur5 policy=openvla_oft + +# Example to deploy your policy with overlay viewer: +python main.py robot=ur5 policy=wog policy.record_dir="your_record_dir" ++overlay.enabled=true ++overlay.first_frames_dir="your_target_frames_to_be_overlayed" ++overlay.camera_name="camera_0" ++overlay.ref_index=1 ++overlay.alpha=0.7 + +## Example first frame directory structure: + +target_frames/ +├── 0.jpg +├── 1.jpg +└── ... + +``` ## 🛠️ Available Tools The `tools/` directory contains utility scripts for: diff --git a/configs/policy/wog.yaml b/configs/policy/wog.yaml new file mode 100644 index 0000000..9f0563b --- /dev/null +++ b/configs/policy/wog.yaml @@ -0,0 +1,31 @@ +_target_: maniunicon.policies.torch_model_vla.VLAPolicy + +frame_latency: ${eval:'1 / 30'} +infer_latency: ${eval:'1 / 20'} +steps_per_inference: 16 # expected openloop steps +use_real_time_chunking: false +enable_recording: true +record_dir: null # path to recorded files +device: ${device} +synchronized: ${synchronized} + +model: + _target_: maniunicon.customize.policy_model.wog_model.WoGpolicy + saved_model_path: "your_ckpt_path" + # device: ${device} + unnorm_key: "fold" + task_name: "fold the blue towel" + +obs_wrapper: + _target_: maniunicon.customize.obs_wrapper.spatialvla_rgb_wrapper.SpatialVLAImageWrapper + camera_config: ${data.camera_config} + state_keys: ["tcp_position", "tcp_orientation", "joint_positions", "joint_velocities", "gripper_state"] + num_cams: 1 # Set to 1 for single camera, 2+ for multi-camera setup + observation_horizon: 1 # Should match the window size in VLA model + +act_wrapper: + _target_: maniunicon.customize.act_wrapper.spatialvla_eepose_wrapper.SpatialVLAEEPoseWrapper + hist_action: 0 # Number of historical actions to consider + action_horizon: 8 + control_mode: cartesian + dt: ${policy.dt} \ No newline at end of file diff --git a/main.py b/main.py index 4628cad..eb36c7b 100644 --- a/main.py +++ b/main.py @@ -10,6 +10,7 @@ from omegaconf import OmegaConf from maniunicon.utils.shared_memory.shared_storage import SharedStorage +from maniunicon.utils.overlay_viewer import run_overlay_viewer class RobotControlSystem: @@ -21,6 +22,7 @@ def __init__( robot_cfg: Dict[str, Any], policy_cfg: Dict[str, Any], sensors_cfg: Dict[str, Any], + overlay_cfg: Dict[str, Any] | None = None, max_record_steps: int = 500, ): # Create shared memory @@ -56,6 +58,10 @@ def __init__( for name, sensor_cfg in sensors_cfg.items() } + # Overlay viewer config + self.overlay_cfg = overlay_cfg + self.overlay_proc = None + # Set up signal handlers signal.signal(signal.SIGINT, self._signal_handler) signal.signal(signal.SIGTERM, self._signal_handler) @@ -82,6 +88,26 @@ def _on_press(key): print(f"{sensor_name} started.") print("All sensors are started.") + # Spawn shared-memory overlay viewer if enabled + try: + if self.overlay_cfg is not None and getattr(self.overlay_cfg, "enabled", False): + ref_dir = getattr(self.overlay_cfg, "first_frames_dir", None) + cam_name = getattr(self.overlay_cfg, "camera_name", "camera_0") + ref_index = int(getattr(self.overlay_cfg, "ref_index", 0)) + alpha = float(getattr(self.overlay_cfg, "alpha", 0.7)) + if ref_dir: + self.overlay_proc = mp.Process( + target=run_overlay_viewer, + args=(self.shared_storage, ref_dir, cam_name, ref_index, alpha), + daemon=True, + ) + self.overlay_proc.start() + print("Overlay viewer started.") + else: + print("Overlay enabled but first_frames_dir not set; skipping viewer.") + except Exception as e: + print(f"Failed to start overlay viewer: {e}") + self.policy.start() print("Policy started.") @@ -102,6 +128,15 @@ def stop(self): sensor.stop() self.shared_storage.is_running.value = False + # Stop overlay viewer if running + try: + if self.overlay_proc is not None and self.overlay_proc.is_alive(): + self.overlay_proc.terminate() + self.overlay_proc.join(timeout=2) + print("Overlay viewer stopped.") + except Exception as e: + print(f"Failed to stop overlay viewer: {e}") + self.shm_manager.shutdown() print("All processes stopped") @@ -133,6 +168,7 @@ def main(cfg): robot_cfg=cfg.robot, policy_cfg=cfg.policy, sensors_cfg=cfg.sensors, + overlay_cfg=getattr(cfg, "overlay", None), max_record_steps=cfg.max_record_steps, ) diff --git a/maniunicon/customize/policy_model/wog_model.py b/maniunicon/customize/policy_model/wog_model.py new file mode 100644 index 0000000..7f7c271 --- /dev/null +++ b/maniunicon/customize/policy_model/wog_model.py @@ -0,0 +1,126 @@ +import numpy as np +from PIL import Image +from typing import Optional, Tuple, Union +import os +import argparse +import json +import math +from flask import Flask, request, jsonify +import tempfile +import torch +from vla import load_vla +from sim_wog.adaptive_ensemble import AdaptiveEnsembler +from maniunicon.utils.vla_utils import euler_pose_to_quat + +app = Flask(__name__) + + +class WoGpolicy: + def __init__( + self, + saved_model_path: str = "", + unnorm_key: str = None, + image_size: list[int] = [224, 224], + action_model_type: str = "DiT-L", + future_action_window_size: int = 15, + use_bf16: bool = True, + action_dim: int = 7, + action_ensemble: bool = True, + adaptive_ensemble_alpha: float = 0.1, + action_ensemble_horizon: int = 2, + action_chunking: bool = False, + action_chunking_window: Optional[int] = None, + args=None, + task_name: str = None, + **kwargs, + ) -> None: + os.environ["TOKENIZERS_PARALLELISM"] = "false" + assert not (action_chunking and action_ensemble), "Now 'action_chunking' and 'action_ensemble' cannot both be True." + + self.unnorm_key = unnorm_key + self.task_name = task_name + + print(f"*** unnorm_key: {unnorm_key} ***") + self.vla = load_vla( + saved_model_path, + load_for_training=False, + action_model_type=action_model_type, + future_action_window_size=future_action_window_size, + action_dim=action_dim, + ) + if use_bf16: + self.vla.vlm = self.vla.vlm.to(torch.bfloat16) + self.vla = self.vla.to("cuda").eval() + + self.image_size = image_size + self.action_ensemble = action_ensemble + self.adaptive_ensemble_alpha = adaptive_ensemble_alpha + self.action_ensemble_horizon = action_ensemble_horizon + self.action_chunking = action_chunking + self.action_chunking_window = action_chunking_window + if self.action_ensemble: + self.action_ensembler = AdaptiveEnsembler(self.action_ensemble_horizon, self.adaptive_ensemble_alpha) + else: + self.action_ensembler = None + + self.args = args + self.reset() + + def reset(self) -> None: + if self.action_ensemble: + self.action_ensembler.reset() + + def __call__( + self, obs, + *args, **kwargs, + ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: + task_description = self.task_name + image_numpy=obs["image"]["camera_0"].squeeze().cpu().detach().numpy() + image_pil = Image.fromarray(image_numpy) + + resized_image = resize_image(image_pil, size=self.image_size) + unnormed_actions, normalized_actions = self.vla.predict_action( + image=resized_image, + instruction=task_description, + unnorm_key=self.unnorm_key, + do_sample=False, + ) + + N, A = unnormed_actions.shape + unnormed_actions = np.concatenate( + [ + euler_pose_to_quat(unnormed_actions[..., :-1].reshape(-1, A - 1)).reshape( + N, -1 + ), + unnormed_actions[..., -1:], + ], + axis=-1, + ) + + + print(f"Instruction: {task_description}") + return unnormed_actions[:8,:] # better for real-time inference + +def resize_image(image: Image, size=(224, 224), shift_to_left=0): + w, h = image.size + # assert h < w, "Height should be less than width" + left_margin = (w - h) // 2 - shift_to_left + left_margin = min(max(left_margin, 0), w - h) + image = image.crop((left_margin, 0, left_margin + h, h)) + + image = image.resize(size, resample=Image.LANCZOS) + + image = scale_and_resize(image, target_size=(224, 224), scale=0.9, margin_w_ratio=0.5, margin_h_ratio=0.5) + return image + +def scale_and_resize(image : Image, target_size=(224, 224), scale=0.9, margin_w_ratio=0.5, margin_h_ratio=0.5): + w, h = image.size + new_w = int(w * math.sqrt(scale)) + new_h = int(h * math.sqrt(scale)) + margin_w_max = w - new_w + margin_h_max = h - new_h + margin_w = int(margin_w_max * margin_w_ratio) + margin_h = int(margin_h_max * margin_h_ratio) + image = image.crop((margin_w, margin_h, margin_w + new_w, margin_h + new_h)) + image = image.resize(target_size, resample=Image.LANCZOS) + return image \ No newline at end of file diff --git a/maniunicon/utils/overlay_viewer.py b/maniunicon/utils/overlay_viewer.py new file mode 100644 index 0000000..bd0fad2 --- /dev/null +++ b/maniunicon/utils/overlay_viewer.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +import cv2 as cv +import numpy as np +from pathlib import Path +import time + +def _load_ref_image(ref_dir: Path, index: int): + img_path = ref_dir / f"{index}.jpg" + if not img_path.exists(): + print(f"No reference image path: {img_path}") + return None + img = cv.imread(str(img_path), cv.IMREAD_COLOR) + if img is None: + print(f"Failed to load reference image: {img_path}") + return img + +def run_overlay_viewer(shared_storage, first_frames_dir: str, camera_name: str = "camera_0", + ref_index: int = 0, alpha_init: float = 0.7): + ref_dir = Path(first_frames_dir) + if not ref_dir.is_dir(): + print(f"First frames directory does not exist: {ref_dir}") + return + + ref_img = _load_ref_image(ref_dir, ref_index) + if ref_img is None: + print("Failed to load initial reference image. Exiting overlay process.") + return + + win_name = f"Shared Camera Overlay ({camera_name})" + cv.namedWindow(win_name, cv.WINDOW_NORMAL) + cv.resizeWindow(win_name, 1280, 720) + + cv.createTrackbar("alpha", win_name, int(alpha_init * 100), 100, lambda x: None) + + ref_indices = sorted([int(p.stem) for p in ref_dir.glob("*.jpg") if p.stem.isdigit()]) + max_ref_index = ref_indices[-1] if ref_indices else ref_index + cv.createTrackbar("ref_index", win_name, min(ref_index, max_ref_index), max_ref_index, lambda x: None) + + print("Overlay process started. Interaction instructions:") + print(" q/ESC exit overlay process") + print(" trackbars adjust alpha and ref_index") + + try: + while True: + data = shared_storage.read_single_camera(camera_name) + if data is None: + time.sleep(0.01) + continue + + rgb = data.color # RGB + if rgb is None: + continue + bgr = cv.cvtColor(rgb, cv.COLOR_RGB2BGR) + + alpha_slider = cv.getTrackbarPos("alpha", win_name) + alpha = max(0.0, min(1.0, alpha_slider / 100.0)) + beta = 1.0 - alpha + + track_idx = cv.getTrackbarPos("ref_index", win_name) + if track_idx != ref_index: + new_ref = _load_ref_image(ref_dir, track_idx) + if new_ref is not None: + ref_img = new_ref + ref_index = track_idx + else: + cv.setTrackbarPos("ref_index", win_name, ref_index) + + h, w = bgr.shape[:2] + ref_resized = cv.resize(ref_img, (w, h), interpolation=cv.INTER_AREA) + + overlay = cv.addWeighted(bgr, alpha, ref_resized, beta, 0.0) + text = f"ref_index={ref_index} alpha={alpha:.2f}" + cv.putText(overlay, text, (10, 30), cv.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 0), 2) + + cv.imshow(win_name, overlay) + key = cv.waitKey(1) & 0xFF + if key in (ord('q'), 27): # q or ESC + break + finally: + cv.destroyWindow(win_name) \ No newline at end of file