training/scripts/planned_comparison_common.py

"""
Shared helpers for the proof-gated planned comparison benchmark.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Dict, Iterable, Tuple

import mujoco
import numpy as np


IMAGE_WIDTH = 96
IMAGE_HEIGHT = 96
OBJECT_HALF_EXTENTS = np.array([0.03, 0.03, 0.03], dtype=np.float32)
OBJECT_CENTER_Z = 0.43
TRAIN_CAMERAS = ("front",)
OOD_CAMERAS = ("top", "angle1")
QUERY_FORCE = 7.0
CALIBRATION_FORCES = (4.0, 6.0, 8.0)


@dataclass(frozen=True)
class CameraSpec:
    name: str
    pos: Tuple[float, float, float]
    xyaxes: Tuple[float, float, float, float, float, float]
    fovy: float


@dataclass(frozen=True)
class MaterialSpec:
    name: str
    mass_range: Tuple[float, float]
    damping_range: Tuple[float, float]


CAMERAS: Dict[str, CameraSpec] = {
    "front": CameraSpec("front", (0.0, -0.55, 0.72), (1.0, 0.0, 0.0, 0.0, 0.58, 0.81), 42.0),
    "top": CameraSpec("top", (0.0, 0.0, 1.1), (1.0, 0.0, 0.0, 0.0, 1.0, 0.0), 38.0),
    "angle1": CameraSpec("angle1", (0.33, -0.42, 0.66), (0.78, 0.62, 0.0, -0.25, 0.31, 0.92), 44.0),
}

HIDDEN_MATERIALS: Dict[str, MaterialSpec] = {
    "steel_core": MaterialSpec("steel_core", (1.6, 1.9), (8.5, 10.0)),
    "wood_core": MaterialSpec("wood_core", (1.1, 1.4), (5.0, 6.2)),
    "plastic_core": MaterialSpec("plastic_core", (0.7, 0.95), (2.8, 4.0)),
    "foam_core": MaterialSpec("foam_core", (0.35, 0.55), (0.9, 1.8)),
}


def _normalize(vec: Iterable[float]) -> np.ndarray:
    arr = np.asarray(vec, dtype=np.float64)
    norm = np.linalg.norm(arr)
    if norm == 0:
        raise ValueError("Cannot normalize zero vector.")
    return arr / norm


def camera_axes(spec: CameraSpec) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    x_axis = _normalize(spec.xyaxes[:3])
    y_axis = _normalize(spec.xyaxes[3:])
    z_axis = _normalize(np.cross(x_axis, y_axis))
    return x_axis, y_axis, z_axis


def camera_intrinsics(spec: CameraSpec, width: int = IMAGE_WIDTH, height: int = IMAGE_HEIGHT) -> np.ndarray:
    focal = 0.5 * height / math.tan(math.radians(spec.fovy) / 2.0)
    cx = (width - 1) / 2.0
    cy = (height - 1) / 2.0
    return np.array(
        [
            [focal, 0.0, cx],
            [0.0, focal, cy],
            [0.0, 0.0, 1.0],
        ],
        dtype=np.float32,
    )


def camera_extrinsics(spec: CameraSpec) -> np.ndarray:
    x_axis, y_axis, z_axis = camera_axes(spec)
    rotation_c2w = np.stack([x_axis, y_axis, z_axis], axis=1)
    rotation_w2c = rotation_c2w.T
    translation_w2c = -rotation_w2c @ np.asarray(spec.pos, dtype=np.float64)
    extrinsics = np.eye(4, dtype=np.float32)
    extrinsics[:3, :3] = rotation_w2c.astype(np.float32)
    extrinsics[:3, 3] = translation_w2c.astype(np.float32)
    return extrinsics


def world_to_pixel(world_point: np.ndarray, spec: CameraSpec, width: int = IMAGE_WIDTH, height: int = IMAGE_HEIGHT) -> np.ndarray:
    point = np.asarray(world_point, dtype=np.float64)
    cam_pos = np.asarray(spec.pos, dtype=np.float64)
    x_axis, y_axis, z_axis = camera_axes(spec)
    rotation_w2c = np.stack([x_axis, y_axis, z_axis], axis=1).T
    point_cam = rotation_w2c @ (point - cam_pos)
    depth = -point_cam[2]
    if depth <= 1e-8:
        raise ValueError("Point projects behind the camera.")
    intrinsics = camera_intrinsics(spec, width=width, height=height)
    u = intrinsics[0, 0] * (point_cam[0] / depth) + intrinsics[0, 2]
    v = intrinsics[1, 2] - intrinsics[1, 1] * (point_cam[1] / depth)
    return np.array([u, v], dtype=np.float32)


def pixel_to_world_on_plane(
    pixel_uv: Tuple[float, float],
    spec: CameraSpec,
    plane_z: float = OBJECT_CENTER_Z,
    width: int = IMAGE_WIDTH,
    height: int = IMAGE_HEIGHT,
) -> np.ndarray:
    intrinsics = camera_intrinsics(spec, width=width, height=height)
    cx = intrinsics[0, 2]
    cy = intrinsics[1, 2]
    fx = intrinsics[0, 0]
    fy = intrinsics[1, 1]
    u, v = pixel_uv
    ray_cam = np.array([(u - cx) / fx, -(v - cy) / fy, -1.0], dtype=np.float64)
    ray_cam /= np.linalg.norm(ray_cam)
    x_axis, y_axis, z_axis = camera_axes(spec)
    rotation_c2w = np.stack([x_axis, y_axis, z_axis], axis=1)
    ray_world = rotation_c2w @ ray_cam
    cam_pos = np.asarray(spec.pos, dtype=np.float64)
    denom = ray_world[2]
    if abs(denom) < 1e-8:
        raise ValueError("Ray is parallel to the plane.")
    scale = (plane_z - cam_pos[2]) / denom
    if scale <= 0:
        raise ValueError("Ray-plane intersection lies behind the camera.")
    return (cam_pos + scale * ray_world).astype(np.float32)


def detect_object_centroid(image: np.ndarray) -> Tuple[float, float]:
    red = image[..., 0].astype(np.int16)
    green = image[..., 1].astype(np.int16)
    blue = image[..., 2].astype(np.int16)
    mask = (red > 120) & ((red - green) > 55) & ((red - blue) > 55)
    if int(mask.sum()) < 8:
        score = red - green - blue
        flat_index = int(np.argmax(score))
        y_idx, x_idx = np.unravel_index(flat_index, score.shape)
        return float(x_idx), float(y_idx)
    y_coords, x_coords = np.nonzero(mask)
    return float(x_coords.mean()), float(y_coords.mean())


def extract_structured_state(
    rgb_pre: np.ndarray,
    rgb_post: np.ndarray,
    camera_name: str,
    calibration_force: float,
    width: int = IMAGE_WIDTH,
    height: int = IMAGE_HEIGHT,
) -> Dict[str, np.ndarray]:
    spec = CAMERAS[camera_name]
    pre_uv = detect_object_centroid(rgb_pre)
    post_uv = detect_object_centroid(rgb_post)
    p_pre = pixel_to_world_on_plane(pre_uv, spec, plane_z=OBJECT_CENTER_Z, width=width, height=height)
    p_post = pixel_to_world_on_plane(post_uv, spec, plane_z=OBJECT_CENTER_Z, width=width, height=height)
    delta_x = float(p_post[0] - p_pre[0])
    mobility = delta_x / float(calibration_force)
    return {
        "pre_uv": np.asarray(pre_uv, dtype=np.float32),
        "post_uv": np.asarray(post_uv, dtype=np.float32),
        "p_pre": p_pre.astype(np.float32),
        "p_post": p_post.astype(np.float32),
        "delta_x": np.array(delta_x, dtype=np.float32),
        "mobility": np.array(mobility, dtype=np.float32),
    }


def build_scene_xml(
    mass: float,
    damping: float,
    initial_x: float,
    table_rgba: Tuple[float, float, float],
    floor_rgba: Tuple[float, float, float],
    light_1: Tuple[float, float, float],
    light_2: Tuple[float, float, float],
) -> str:
    camera_xml = []
    for spec in CAMERAS.values():
        camera_xml.append(
            '<camera name="{name}" pos="{px:.3f} {py:.3f} {pz:.3f}" '
            'xyaxes="{x0:.3f} {x1:.3f} {x2:.3f} {y0:.3f} {y1:.3f} {y2:.3f}" '
            'fovy="{fovy:.1f}"/>'.format(
                name=spec.name,
                px=spec.pos[0],
                py=spec.pos[1],
                pz=spec.pos[2],
                x0=spec.xyaxes[0],
                x1=spec.xyaxes[1],
                x2=spec.xyaxes[2],
                y0=spec.xyaxes[3],
                y1=spec.xyaxes[4],
                y2=spec.xyaxes[5],
                fovy=spec.fovy,
            )
        )
    return f"""
<mujoco model="planned_comparison">
  <option timestep="0.002" gravity="0 0 0"/>
  <visual>
    <global offwidth="{IMAGE_WIDTH}" offheight="{IMAGE_HEIGHT}"/>
    <headlight ambient="0.65 0.65 0.65" diffuse="0.35 0.35 0.35"/>
  </visual>
  <asset>
    <material name="object_mat" rgba="0.88 0.18 0.18 1" specular="0.02" shininess="0.05"/>
  </asset>
  <worldbody>
    <light pos="{light_1[0]:.3f} {light_1[1]:.3f} {light_1[2]:.3f}" diffuse="0.55 0.55 0.55" specular="0 0 0"/>
    <light pos="{light_2[0]:.3f} {light_2[1]:.3f} {light_2[2]:.3f}" diffuse="0.30 0.30 0.30" specular="0 0 0"/>
    <geom name="floor" type="plane" size="2 2 0.1" rgba="{floor_rgba[0]:.3f} {floor_rgba[1]:.3f} {floor_rgba[2]:.3f} 1"/>
    <body name="table" pos="0 0 0.38">
      <geom type="box" size="0.42 0.26 0.02" rgba="{table_rgba[0]:.3f} {table_rgba[1]:.3f} {table_rgba[2]:.3f} 1"/>
    </body>
    <body name="object_body" pos="0 0 {OBJECT_CENTER_Z:.3f}">
      <joint name="slide_x" type="slide" axis="1 0 0" damping="{damping:.6f}"/>
      <geom name="object_geom" type="box" size="{OBJECT_HALF_EXTENTS[0]:.3f} {OBJECT_HALF_EXTENTS[1]:.3f} {OBJECT_HALF_EXTENTS[2]:.3f}"
            mass="{mass:.6f}" material="object_mat"/>
    </body>
    {' '.join(camera_xml)}
  </worldbody>
  <actuator>
    <motor name="push_motor" joint="slide_x" gear="1"/>
  </actuator>
</mujoco>
""".strip()


def sample_material_and_context(rng: np.random.Generator) -> Dict[str, object]:
    material_name = rng.choice(tuple(HIDDEN_MATERIALS.keys()))
    material = HIDDEN_MATERIALS[material_name]
    mass = float(rng.uniform(*material.mass_range))
    damping = float(rng.uniform(*material.damping_range))
    initial_x = float(rng.uniform(-0.10, 0.10))
    table_rgba = tuple(float(v) for v in rng.uniform([0.34, 0.24, 0.18], [0.52, 0.38, 0.28], size=3))
    floor_rgba = tuple(float(v) for v in rng.uniform([0.22, 0.22, 0.25], [0.42, 0.42, 0.45], size=3))
    light_1 = tuple(float(v) for v in rng.uniform([-0.5, -0.5, 1.0], [0.5, 0.2, 1.4], size=3))
    light_2 = tuple(float(v) for v in rng.uniform([-0.4, 0.0, 0.8], [0.4, 0.6, 1.2], size=3))
    calibration_force = float(rng.choice(CALIBRATION_FORCES))
    return {
        "material_name": material_name,
        "mass": mass,
        "damping": damping,
        "initial_x": initial_x,
        "table_rgba": table_rgba,
        "floor_rgba": floor_rgba,
        "light_1": light_1,
        "light_2": light_2,
        "calibration_force": calibration_force,
    }


def render_rgb(model: mujoco.MjModel, data: mujoco.MjData, camera_name: str) -> np.ndarray:
    camera_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, camera_name)
    renderer = mujoco.Renderer(model, height=IMAGE_HEIGHT, width=IMAGE_WIDTH)
    try:
        renderer.update_scene(data, camera=camera_id)
        return renderer.render().copy()
    finally:
        renderer.close()


def simulate_push_displacement(model: mujoco.MjModel, force: float, initial_x: float, duration: float = 0.35) -> float:
    data = mujoco.MjData(model)
    joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "slide_x")
    qpos_adr = model.jnt_qposadr[joint_id]
    data.qpos[qpos_adr] = initial_x
    mujoco.mj_forward(model, data)
    x_before = float(data.qpos[qpos_adr])
    total_steps = int(duration / model.opt.timestep)
    push_steps = max(1, int(0.12 / model.opt.timestep))
    for step in range(total_steps):
        data.ctrl[0] = force if step < push_steps else 0.0
        mujoco.mj_step(model, data)
    x_after = float(data.qpos[qpos_adr])
    return x_after - x_before


def generate_episode(
    rng: np.random.Generator,
    camera_name: str,
) -> Dict[str, object]:
    context = sample_material_and_context(rng)
    xml = build_scene_xml(
        mass=context["mass"],
        damping=context["damping"],
        initial_x=context["initial_x"],
        table_rgba=context["table_rgba"],
        floor_rgba=context["floor_rgba"],
        light_1=context["light_1"],
        light_2=context["light_2"],
    )
    model = mujoco.MjModel.from_xml_string(xml)
    data = mujoco.MjData(model)
    joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "slide_x")
    qpos_adr = model.jnt_qposadr[joint_id]
    data.qpos[qpos_adr] = context["initial_x"]
    mujoco.mj_forward(model, data)
    rgb_pre = render_rgb(model, data, camera_name)

    total_steps = int(0.35 / model.opt.timestep)
    push_steps = max(1, int(0.12 / model.opt.timestep))
    for step in range(total_steps):
        data.ctrl[0] = context["calibration_force"] if step < push_steps else 0.0
        mujoco.mj_step(model, data)
    rgb_post = render_rgb(model, data, camera_name)
    gt_pre = np.array([context["initial_x"], 0.0, OBJECT_CENTER_Z], dtype=np.float32)
    gt_post = np.array([float(data.qpos[qpos_adr]), 0.0, OBJECT_CENTER_Z], dtype=np.float32)
    gt_delta_x = float(gt_post[0] - gt_pre[0])
    target_query_dx = float(simulate_push_displacement(model, QUERY_FORCE, context["initial_x"]))
    state = extract_structured_state(rgb_pre, rgb_post, camera_name, context["calibration_force"])
    spec = CAMERAS[camera_name]
    return {
        "rgb_pre": rgb_pre,
        "rgb_post": rgb_post,
        "camera_name": camera_name,
        "camera_intrinsics": camera_intrinsics(spec),
        "camera_extrinsics": camera_extrinsics(spec),
        "calibration_force": np.array(context["calibration_force"], dtype=np.float32),
        "query_force": np.array(QUERY_FORCE, dtype=np.float32),
        "structured_pre": state["p_pre"],
        "structured_post": state["p_post"],
        "structured_delta_x": state["delta_x"],
        "structured_mobility": state["mobility"],
        "target_query_dx": np.array(target_query_dx, dtype=np.float32),
        "detected_pre_uv": state["pre_uv"],
        "detected_post_uv": state["post_uv"],
        "material_name": context["material_name"],
        "mass": np.array(context["mass"], dtype=np.float32),
        "damping": np.array(context["damping"], dtype=np.float32),
        "gt_pre_world": gt_pre,
        "gt_post_world": gt_post,
        "gt_delta_x": np.array(gt_delta_x, dtype=np.float32),
    }


def average_pool_rgb(image: np.ndarray, downsample_factor: int = 4) -> np.ndarray:
    height, width, channels = image.shape
    if height % downsample_factor != 0 or width % downsample_factor != 0:
        raise ValueError("Image size must be divisible by downsample_factor.")
    pooled = image.reshape(
        height // downsample_factor,
        downsample_factor,
        width // downsample_factor,
        downsample_factor,
        channels,
    ).mean(axis=(1, 3))
    return pooled.astype(np.float32)