add wog and continual overlay

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:

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}

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,
     )
 

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

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)