mt50_evo1_client_prompt.py

# mt50_evo1_client.py
import asyncio
import json
import os
from typing import List, Optional, Dict, Set

import cv2
import gymnasium as gym
import metaworld  # noqa: F401
import numpy as np
import websockets
import random

import datetime

# ===================== Logging =====================
LOG_DIR = "logs"
os.makedirs(LOG_DIR, exist_ok=True)

def make_log_path(prefix="eval"):
    ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    return os.path.join(LOG_DIR, f"{prefix}_{ts}.txt")

LOG_PATH = make_log_path("mt50")
# ====================================================

SHOW_WINDOW = True
SAVE_IMAGE = False
SAVE_VIDEO = True  # save the video of each episode to disk

# ===================== Debug image saving =====================
INSPECT_SAMPLE_PER_EPISODE = True        
INSPECT_DIR = "inspect_frames"           
APPLY_ROT_180 = True                     
APPLY_CENTER_CROP = True                 
CROP_KEEP_RATIO = 2/3                    
INSPECT_SAVE_STEP_TAG = True             
# =============================================================

# ===================== Debug video saving ====================
VIDEO_SAVE_DIR = "episode_videos"
VIDEO_FPS = 10  # Original writing frame rate (used to control playback speed; the smaller the value, the slower the playback).
VIDEO_DUP_FRAMES = 1  # Number of times to duplicate each frame when writing video (used to control playback speed; the larger the value, the slower the playback).
# =============================================================


# ===================== User Config (edit here) =====================
SERVER_URL = "ws://127.0.0.1:9000"

# Camera & image settings
CAMERA_NAME = "corner2"        
IMG_SIZE = (448, 448)          

# Evo1 & rollout settings
STATE_TAKE = 8                
HORIZON = 15                  
EPISODES = 10                  
EPISODE_HORIZON = 400          
SEED = 4042

TARGET_LEVEL = "all"   # one of "all", "easy", "medium", "hard", "very_hard"

# Order source
ORDER_JSON_PATH = "mt50_order.json"      

FALLBACK_USE_FIRST_N: Optional[int] = 5
FALLBACK_IDX_LIST: Optional[List[int]] = None

# Prompt source
TASKS_JSONL_PATH = "tasks.jsonl"         
# ==================================================================

# Headless GL by default; switch to 'glfw' on a desktop if you want
os.environ.setdefault("MUJOCO_GL", "egl")
gym.logger.min_level = gym.logger.ERROR


# ---------------- Utils ----------------
def encode_image_uint8_list(img_bgr: np.ndarray):
    return img_bgr.astype(np.uint8).tolist()

def obs_to_state(obs, take: int = STATE_TAKE) -> List[float]:
    if isinstance(obs, dict):
        if "observation" in obs:
            arr = np.asarray(obs["observation"], dtype=np.float32).ravel()
        else:
            parts = [np.asarray(v).ravel() for v in obs.values()]
            arr = np.concatenate(parts).astype(np.float32)
    else:
        arr = np.asarray(obs, dtype=np.float32).ravel()
    return arr[:min(take, arr.shape[0])].tolist()

def fix_camera_angle(rgb: np.ndarray) -> np.ndarray:
    
    return cv2.rotate(rgb, cv2.ROTATE_180)

def center_crop_keep_ratio(rgb: np.ndarray, keep_ratio: float) -> np.ndarray:
    
    h, w = rgb.shape[:2]
    keep_ratio = float(keep_ratio)
    keep_ratio = max(1e-6, min(1.0, keep_ratio))  
    new_h = max(1, int(round(h * keep_ratio)))
    new_w = max(1, int(round(w * keep_ratio)))
    y0 = (h - new_h) // 2
    x0 = (w - new_w) // 2
    return rgb[y0:y0 + new_h, x0:x0 + new_w, :]

def render_single_bgr(env) -> np.ndarray:
  
    rgb = env.render()                               
    rgb = np.ascontiguousarray(rgb, dtype=np.uint8)   

   
    if APPLY_ROT_180:
        rgb = cv2.rotate(rgb, cv2.ROTATE_180)
        rgb = np.ascontiguousarray(rgb)

    
    if APPLY_CENTER_CROP and (0.0 < CROP_KEEP_RATIO < 1.0):
        h, w = rgb.shape[:2]
        keep = float(CROP_KEEP_RATIO)
        new_h = max(1, int(round(h * keep)))
        new_w = max(1, int(round(w * keep)))
        y0 = (h - new_h) // 2
        x0 = (w - new_w) // 2
        rgb = rgb[y0:y0 + new_h, x0:x0 + new_w, :].copy()
        rgb = np.ascontiguousarray(rgb)

   
    if IMG_SIZE is not None:
        rgb = cv2.resize(rgb, IMG_SIZE, interpolation=cv2.INTER_LINEAR)
        rgb = np.ascontiguousarray(rgb)

    
    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    bgr = np.ascontiguousarray(bgr, dtype=np.uint8)

    
    if 'SHOW_WINDOW' in globals() and SHOW_WINDOW:
        try:
            cv2.imshow("MetaWorld", bgr)
            cv2.waitKey(1)   
        except Exception:
           
            pass

    return bgr

def create_video_writer(env, video_name: str):
    """
    create and return a cv2.VideoWriter object for saving episode videos.
    """
    os.makedirs(VIDEO_SAVE_DIR, exist_ok=True)
    probe_frame = render_single_bgr(env)  # Render one frame first to get the dimensions.
    h0, w0 = probe_frame.shape[:2]
    frame_size = (w0, h0)
    video_path = os.path.join(VIDEO_SAVE_DIR, video_name)
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    video_writer = cv2.VideoWriter(video_path, fourcc, VIDEO_FPS, frame_size)
    # Write the detection frame as the first frame.
    for _ in range(VIDEO_DUP_FRAMES):
        video_writer.write(probe_frame)
    return video_writer

def write_video(video_writer, img_bgr: np.ndarray):
    """
    write a frame to the given cv2.VideoWriter object.
    """
    try:
        if video_writer is not None:
            for _ in range(VIDEO_DUP_FRAMES):
                video_writer.write(img_bgr)
    except Exception as e:
        log_write(f"[video][ERROR] writer.write failed: {e}")

def save_episode_video(writer, video_name: str, task_idx: int, slug: str, ep_num: int):
    """save the video to disk and close video writer."""
    if writer is None:
        return
    try:
        video_path = os.path.join(VIDEO_SAVE_DIR, video_name)
        writer.release()
        log_write(f"[video] task={task_idx} slug={slug} ep={ep_num} saved video frames {video_path}")
    except Exception as e:
        log_write(f"[video][ERROR] closing writer failed: {e}")


async def evo1_infer(ws, img_bgr: np.ndarray, state_vec: List[float], prompt: Optional[str] = None) -> np.ndarray:
    assert prompt is not None and len(prompt) > 0, "prompt should be non-empty"
    dummy_img = np.zeros((448, 448, 3), dtype=np.uint8)
    payload = {
        "image": [encode_image_uint8_list(img_bgr),
                  encode_image_uint8_list(dummy_img),
                  encode_image_uint8_list(dummy_img)],
        "state": state_vec,
        "prompt": prompt,              
        "image_mask": [1, 0, 0],
        "action_mask": [1, 1, 1, 1] + [0]*20,
    }
    await ws.send(json.dumps(payload))
    data = json.loads(await ws.recv())
    return np.asarray(data, dtype=np.float32)


def save_sent_bgr_frame(img_bgr: np.ndarray, ep_num: int, idx: int, slug: str, step: Optional[int] = None):

    os.makedirs(INSPECT_DIR, exist_ok=True)
    tag = f"step{step:04d}" if (INSPECT_SAVE_STEP_TAG and step is not None) else "stepNA"
    out = os.path.join(INSPECT_DIR, f"ep{ep_num:03d}_idx{idx}_{slug}_{tag}.png")
    img_bgr_safe = np.ascontiguousarray(img_bgr)  
    cv2.imwrite(out, img_bgr_safe)
    h, w = img_bgr_safe.shape[:2]
    print(f"[inspect] saved {out}  size={w}x{h}  (identical to VLA input)")

def log_write(text: str):
    
    print(text)
    with open(LOG_PATH, "a", encoding="utf-8") as f:
        f.write(text + "\n")

# ---------------- Prompt loader ----------------
class PromptBook:

    def __init__(self, jsonl_path: str):
        self.by_idx: Dict[int, str] = {}
        self.by_slug: Dict[str, str] = {}
        self.seq: List[str] = []

        if not os.path.exists(jsonl_path):
            print(f"[WARN] {jsonl_path} not found; prompts will be empty.")
            return

        with open(jsonl_path, "r", encoding="utf-8") as f:
            lines = [json.loads(line) for line in f if line.strip()]

        for i, obj in enumerate(lines):
            task_txt = str(obj.get("task", "")).strip()
            if "idx" in obj:
                try:
                    self.by_idx[int(obj["idx"])] = task_txt
                except Exception:
                    pass
            if "slug" in obj:
                try:
                    self.by_slug[str(obj["slug"])] = task_txt
                except Exception:
                    pass
            self.seq.append(task_txt)

    def get(self, idx: int, slug: Optional[str] = None) -> str:
        if idx in self.by_idx:
            return self.by_idx[idx]
        if slug is not None and slug in self.by_slug:
            return self.by_slug[slug]
        if 0 <= idx < len(self.seq):
            return self.seq[idx]
        return ""


PROMPTS = PromptBook(TASKS_JSONL_PATH)


# ---------------- Order & groups loader ----------------
def load_order_and_groups(total_envs: int):
   
    if os.path.exists(ORDER_JSON_PATH):
        with open(ORDER_JSON_PATH, "r") as f:
            data = json.load(f)
        ordered_indices = list(map(int, data["ordered_indices"]))
     
        groups = {k: set(v) for k, v in data["groups"].items()}
        idx_to_slug = {int(k): v for k, v in data["idx_to_slug"].items()}
        print(f"[INFO] Loaded order from {ORDER_JSON_PATH} (len={len(ordered_indices)})")
        log_write(f"[INFO] Metaworld Evaluation Begins ...")
        return ordered_indices, groups, idx_to_slug

  
    if FALLBACK_IDX_LIST:
        idx_list = [i for i in FALLBACK_IDX_LIST if 0 <= i < total_envs]
    elif FALLBACK_USE_FIRST_N:
        idx_list = list(range(min(FALLBACK_USE_FIRST_N, total_envs)))
    else:
        idx_list = list(range(total_envs))
    print("[WARN] mt50_order.json not found; falling back to:", idx_list)
    
    idx_to_slug = {i: f"task-{i}" for i in idx_list}
    groups = {"easy": set(), "medium": set(), "hard": set(), "very_hard": set()}
    return idx_list, groups, idx_to_slug


# ---------------- Core eval (MT50 only, ordered by mt50_order.json) ----------------
async def eval_mt50_with_groups(server_url: str,
                                num_eval_episodes: int = EPISODES,
                                episode_horizon: int = EPISODE_HORIZON,
                                seed: int = SEED):
  
    # 1) Build MT50 with fixed camera
    envs = gym.make_vec(
        "Meta-World/MT50",
        vector_strategy="sync",
        seed=seed,
        render_mode="rgb_array",
        camera_name=CAMERA_NAME,
    )
    total_envs = len(envs.envs)

    # 2) Load ordered idx list & groups
    ordered_indices, groups, idx_to_slug = load_order_and_groups(total_envs)
    ordered_indices = [i for i in ordered_indices if 0 <= i < total_envs]

    
    if TARGET_LEVEL.lower() != "all":
        allowed_slugs = groups.get(TARGET_LEVEL.lower(), set())
        before = len(ordered_indices)
        ordered_indices = [i for i in ordered_indices if idx_to_slug.get(i, "") in allowed_slugs]
        print(f"[INFO] Filtered tasks: keep only {TARGET_LEVEL} ({len(ordered_indices)}/{before})")


    # 3) Accumulators
    success_counts: Dict[int, int] = {i: 0 for i in ordered_indices}
    trials_counts: Dict[int, int] = {i: 0 for i in ordered_indices}
    group_success = {k: 0 for k in ["easy", "medium", "hard", "very_hard"]}
    group_trials  = {k: 0 for k in ["easy", "medium", "hard", "very_hard"]}

    # 4) Main loop
    async with websockets.connect(server_url, max_size=100_000_000) as ws:
        for idx in ordered_indices:
            sub = envs.envs[idx]
            slug = idx_to_slug.get(idx, f"task-{idx}")

            
            task_prompt = PROMPTS.get(idx, slug=slug)
            

            gname_for_task = None
            for gname in group_trials.keys():
                if slug in groups.get(gname, set()):
                    gname_for_task = gname
                    break

            for ep in range(num_eval_episodes):
                for obj in (sub, getattr(sub, "unwrapped", None)):
                    fn = getattr(obj, "iterate_goal_position", None)
                    if callable(fn):
                        try: fn()
                        except Exception: pass
                        break

                inspect_choice = INSPECT_SAMPLE_PER_EPISODE
                saved_this_episode = False

               
                obs, _ = sub.reset(seed=seed + ep)  
                trials_counts[idx] += 1
                if gname_for_task is not None:
                    group_trials[gname_for_task] += 1

                steps = 0
                done = False
                video_name = f"task{idx:02d}_{slug}_ep{ep+1:03d}.mp4"
                video_writer = None if not SAVE_VIDEO else create_video_writer(sub, video_name)

                try:
                    a0 = np.zeros(sub.action_space.shape, dtype=np.float32)
                    a0 = np.clip(a0, sub.action_space.low, sub.action_space.high)
                    obs, _, _, _, _ = sub.step(a0)
                except Exception:
                    pass

                while steps < episode_horizon and not done:
                    img_bgr = render_single_bgr(sub)
                    
                    if SAVE_VIDEO:
                        write_video(video_writer, img_bgr)

                    if SAVE_IMAGE and inspect_choice and (not saved_this_episode):
                        save_sent_bgr_frame(
                            img_bgr, ep_num=ep + 1, idx=idx, slug=slug,
                            step=steps if INSPECT_SAVE_STEP_TAG else None
                        )
                        saved_this_episode = True

                    state_vec = obs_to_state(obs)

                 
                    actions = await evo1_infer(ws, img_bgr, state_vec, prompt=task_prompt)

                    for i in range(HORIZON):
                        a4 = np.asarray(actions[i][:4], dtype=np.float32)
                        a4 = np.clip(a4, sub.action_space.low, sub.action_space.high)
                        obs, _, terminated, truncated, info = sub.step(a4)
                        steps += 1

                        if isinstance(info, dict) and info.get("success", 0) == 1:
                            success_counts[idx] += 1
                            if gname_for_task is not None:
                                group_success[gname_for_task] += 1
                            done = True
                            break

                        if terminated or truncated or steps >= episode_horizon:
                            done = True
                            break
                
                # close video writer
                if done and SAVE_VIDEO:
                    final_frame = render_single_bgr(sub)
                    write_video(video_writer, final_frame)
                    save_episode_video(video_writer, video_name, idx, slug, ep + 1)
                
          
            s = success_counts[idx]
            t = trials_counts[idx]
            task_rate = s / max(1, t)
            msg = (f"[Task {idx} {slug}] {task_prompt} finished {num_eval_episodes} episodes -> "
                  f"success_rate={task_rate:.3f}  (s={s}, t={t})")
            log_write(msg)

    envs.close()

    # 5) Build metrics
    per_task: Dict[str, float] = {}
    for idx in ordered_indices:
        slug = idx_to_slug.get(idx, f"task-{idx}")
        s, t = success_counts[idx], trials_counts[idx]
        per_task[slug] = (s / t) if t > 0 else 0.0

    per_group: Dict[str, float] = {}
    for gname in ["easy", "medium", "hard", "very_hard"]:
        s, t = group_success[gname], group_trials[gname]
        per_group[gname] = (s / t) if t > 0 else 0.0

    overall = (sum(success_counts.values()) /
               max(1, sum(trials_counts.values())))

    return per_task, per_group, overall


# ---------------- Entrypoint ----------------
async def _amain():
    per_task, per_group, overall = await eval_mt50_with_groups(
        server_url=SERVER_URL,
        num_eval_episodes=EPISODES,
        episode_horizon=EPISODE_HORIZON,
        seed=SEED,
    )

    # Pretty print
    # print("\n==== Per-task success rate ====")
    # for slug, rate in per_task.items():
    #     print(f"{slug:24s}  {rate:.3f}")

    # print("\n==== Difficulty buckets ====")
    # print(f"easy      : {per_group.get('easy', 0.0):.3f}")
    # print(f"medium    : {per_group.get('medium', 0.0):.3f}")
    # print(f"hard      : {per_group.get('hard', 0.0):.3f}")
    # print(f"very_hard : {per_group.get('very_hard', 0.0):.3f}")

    avg = (per_group.get('easy', 0.0) + per_group.get('medium', 0.0) + per_group.get('hard', 0.0) + per_group.get('very_hard', 0.0)) / 4
    # print(f"\n==== Overall Average as Success Rate ====\n{avg:.3f}")

    # log
    log_write(f"\n==== Evaluation Log ====\nLog file: {LOG_PATH}")
    log_write(f"Target difficulty: {TARGET_LEVEL}")
    log_write(f"Server URL: {SERVER_URL}")
    log_write(f"Episodes per task: {EPISODES}")
    log_write(f"Episode horizon: {EPISODE_HORIZON}")
    log_write(f"HORIZON: {HORIZON}")
    log_write(f"Seed: {SEED}\n")
    

    log_write("==== Per-task success rate ====")
    for slug, rate in per_task.items():
        log_write(f"{slug:24s}  {rate:.3f}")

    log_write("\n==== Difficulty buckets ====")
    log_write(f"easy      : {per_group.get('easy', 0.0):.3f}")
    log_write(f"medium    : {per_group.get('medium', 0.0):.3f}")
    log_write(f"hard      : {per_group.get('hard', 0.0):.3f}")
    log_write(f"very_hard : {per_group.get('very_hard', 0.0):.3f}")

    log_write(f"\n==== Overall Average as Success Rate ====\n{avg:.3f}")

if __name__ == "__main__":
    asyncio.run(_amain())



# if __name__ == "__main__":
#     N_REPEAT = 1
#     for run_id in range(N_REPEAT):
#         print(f"\n\n===== 🌟 Run {run_id + 1}/{N_REPEAT} =====")
#         asyncio.run(_amain())