"""Build fig1_overview.svg with a 3D volume cloud replacing the 2D heatmap.
Runs the PARA model, subsamples the 32×64×64 heatmap to 32×16×16 via max pooling,
projects each voxel isometrically, and emits SVG circles for the voxel cloud,
plus a wireframe cube outline.
"""
import os
import sys
import math
import base64
import time
_t_start = time.time()
import numpy as np
import torch
import torch.nn.functional as F
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
print(f"[{time.time()-_t_start:5.2f}s] imports complete")
sys.path.insert(0, "/data/cameron/para_normalized_losses/libero")
import model as model_module
from model import TrajectoryHeatmapPredictor, N_HEIGHT_BINS, PRED_SIZE
from libero.libero.envs import OffScreenRenderEnv
from libero.libero import benchmark as bm, get_libero_path
from robosuite.utils.camera_utils import (
get_camera_transform_matrix, project_points_from_world_to_camera,
)
import h5py
# ═══════════════════════════════════════════════════════════════════════════
# Step 1: Run PARA to get the 3D heatmap
# ═══════════════════════════════════════════════════════════════════════════
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
IMAGE_SIZE = 448
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
print(f"[{time.time()-_t_start:5.2f}s] loading PARA model...")
ckpt = torch.load(
"/data/cameron/para_normalized_losses/libero/checkpoints/para_v2_exp4_n64/best.pth",
map_location="cpu")
model_module.MIN_HEIGHT = float(ckpt["min_height"])
model_module.MAX_HEIGHT = float(ckpt["max_height"])
n_window = ckpt["model_state_dict"]["volume_head.weight"].shape[0] // N_HEIGHT_BINS
model = TrajectoryHeatmapPredictor(n_window=n_window)
model.load_state_dict(ckpt["model_state_dict"], strict=False)
model = model.to(device).eval()
print(f"[{time.time()-_t_start:5.2f}s] model loaded; initializing LIBERO...")
benchmark = bm.get_benchmark_dict()["libero_spatial"]()
task = benchmark.get_task(0)
bddl_file = os.path.join(get_libero_path("bddl_files"), task.problem_folder, task.bddl_file)
env = OffScreenRenderEnv(
bddl_file_name=bddl_file,
camera_heights=IMAGE_SIZE, camera_widths=IMAGE_SIZE,
camera_names=["agentview"])
env.seed(0); env.reset()
sim = env.env.sim
for name in ["wooden_cabinet_1_main", "flat_stove_1_main"]:
try: sim.model.body_pos[sim.model.body_name2id(name)] = [0, 0, -5]
except: pass
for name in ["akita_black_bowl_2_main", "cookies_1_main", "glazed_rim_porcelain_ramekin_1_main"]:
try:
dbid = sim.model.body_name2id(name)
for gid in range(sim.model.ngeom):
if sim.model.geom_bodyid[gid] == dbid: sim.model.geom_rgba[gid][3] = 0.0
except: pass
sim.forward()
with h5py.File(os.path.join(get_libero_path("datasets"), benchmark.get_task_demonstration(0)), "r") as f:
state = np.array(f["data/demo_0/states"][32])
obs = env.set_init_state(state); sim.forward()
obs = env.env._get_observations()
rgb = np.flipud(np.asarray(obs["agentview_image"]).copy())
eef = np.array(obs["robot0_eef_pos"])
img = (rgb.astype(np.float32)/255 - MEAN) / STD
img_t = torch.from_numpy(img.transpose(2,0,1)).float().unsqueeze(0).to(device)
w2c = get_camera_transform_matrix(sim, "agentview", IMAGE_SIZE, IMAGE_SIZE)
pix_rc = project_points_from_world_to_camera(eef.reshape(1,3), w2c, IMAGE_SIZE, IMAGE_SIZE)[0]
start_kp = torch.tensor([float(pix_rc[1]), float(pix_rc[0])], dtype=torch.float32).to(device)
print(f"[{time.time()-_t_start:5.2f}s] LIBERO ready; running PARA inference...")
with torch.no_grad():
vol_logits, _, _, _ = model(img_t, start_kp)
print(f"[{time.time()-_t_start:5.2f}s] inference done; subsampling + projecting voxels...")
vol_t = vol_logits[0, -1] # last window timestep
vol_probs = F.softmax(vol_t.reshape(-1), dim=0).reshape(vol_t.shape)
vol_np = vol_probs.cpu().numpy() # (32, 64, 64)
print(f"Raw volume: shape={vol_np.shape}, sum={vol_np.sum():.4f}, max={vol_np.max():.4f}")
env.close()
# ═══════════════════════════════════════════════════════════════════════════
# Step 2: Subsample and transform
# ═══════════════════════════════════════════════════════════════════════════
# Max pool 4x4 over spatial dims: 64 → 16 per axis
# vol_np shape: (32_h, 64_v, 64_u)
vol_sub = vol_np.reshape(32, 16, 4, 16, 4).max(axis=(2, 4)) # (32, 16, 16)
argmax_flat = vol_sub.argmax()
ah = argmax_flat // (16*16); av = (argmax_flat // 16) % 16; au = argmax_flat % 16
print(f"Subsampled: shape={vol_sub.shape}, max={vol_sub.max():.4f}, min={vol_sub.min():.3e}")
print(f"Argmax voxel: h={ah}, v={av}, u={au}")
# Normalize to [0, 1]
vol_norm = vol_sub / vol_sub.max()
# Log-scale visibility to spread values across the full [0, 1] range
# Without this, the peak dominates and everything else is invisible
eps = 1e-10
log_norm = np.log10(vol_norm + eps) # range roughly [-10, 0]
log_min, log_max = log_norm.min(), log_norm.max()
vol_vis = (log_norm - log_min) / (log_max - log_min + 1e-8) # [0, 1]
print(f"Log-vis: min={vol_vis.min():.3f}, max={vol_vis.max():.3f}, mean={vol_vis.mean():.3f}")
# ═══════════════════════════════════════════════════════════════════════════
# Step 3: Isometric projection
# ═══════════════════════════════════════════════════════════════════════════
# True isometric projection:
# - u axis (pixel col) → down-right at 30°
# - v axis (pixel row) → down-left at 30°
# - h axis (height bin) → straight up (screen y decreases)
#
# sx = (u - v) * cos30 * STEP_UV
# sy = (u + v) * sin30 * STEP_UV - h * STEP_H
STEP_UV = 4.2
STEP_H = 2.3
COS30 = math.cos(math.radians(30))
SIN30 = math.sin(math.radians(30))
def project(u, v, h):
sx = (u - v) * COS30 * STEP_UV
sy = (u + v) * SIN30 * STEP_UV - h * STEP_H
return sx, sy
# Compute the bounding box of the projected cube
all_corners = [project(u, v, h) for u in (0, 15) for v in (0, 15) for h in (0, 31)]
xs, ys = zip(*all_corners)
min_x, max_x = min(xs), max(xs)
min_y, max_y = min(ys), max(ys)
cube_w = max_x - min_x
cube_h = max_y - min_y
print(f"Cube projection bounds: {cube_w:.1f} x {cube_h:.1f}")
# Target box in the main SVG — Heatmap Volume new position
BOX_X = 240
BOX_Y = 80
BOX_W = 180
BOX_H = 160
# Center the cube in the box
offset_x = BOX_X + (BOX_W - cube_w) / 2 - min_x
offset_y = BOX_Y + (BOX_H - cube_h) / 2 - min_y
# ═══════════════════════════════════════════════════════════════════════════
# Step 4: Generate voxel cloud + wireframe
# ═══════════════════════════════════════════════════════════════════════════
plasma = plt.cm.plasma
# Render ALL voxels; alpha from log-scaled visibility
# Sort back-to-front: in isometric with camera at (+u, -v, +h) roughly,
# voxels at lower u, higher v, lower h are farthest.
# Depth measure: larger value = closer to camera (drawn last).
# camera direction roughly along (+u, -v, +h); depth = u - v + h
voxels = []
for h in range(32):
for v in range(16):
for u in range(16):
p_vis = float(vol_vis[h, v, u])
sx, sy = project(u, v, h)
sx += offset_x
sy += offset_y
# Standard isometric: camera at (+u, +v, +h). Farther voxels have
# lower (u+v+h) — draw them first. Higher = closer = drawn last.
depth = u + v + h
voxels.append((sx, sy, p_vis, depth))
# Sort ascending by depth: lowest depth drawn first (goes behind)
voxels.sort(key=lambda x: x[3])
print(f"Voxels to render: {len(voxels)} / {32*16*16}")
# Generate SVG — color by plasma, variable radius, alpha by vis
# - Dim voxels: tiny, faint, dark-purple (fog)
# - Bright voxels: large, opaque, yellow (peak)
voxel_svg_parts = []
for sx, sy, p_vis, _ in voxels:
# Color mapping: shift up so dim voxels aren't pure dark navy
color_val = 0.2 + 0.8 * p_vis
rgba = plasma(color_val)
r, g, b = int(rgba[0]*255), int(rgba[1]*255), int(rgba[2]*255)
# Variable radius: dim = small fog, bright = prominent peak
radius = 1.4 + 5.0 * (p_vis ** 1.0)
# Alpha: floor 0.08 for fog, peak gets 1.0
alpha = max(0.08, p_vis ** 1.5)
voxel_svg_parts.append(
f'')
voxel_svg = '\n '.join(voxel_svg_parts)
print(f"[{time.time()-_t_start:5.2f}s] {len(voxel_svg_parts)} voxel SVG shapes generated")
# Wireframe cube outline (8 corners, 12 edges)
def corner_svg(u, v, h):
sx, sy = project(u, v, h)
return (sx + offset_x, sy + offset_y)
corners = {(u, v, h): corner_svg(u, v, h)
for u in (0, 15) for v in (0, 15) for h in (0, 31)}
edges = [
# bottom face (h=0)
((0,0,0), (15,0,0)), ((15,0,0), (15,15,0)),
((15,15,0), (0,15,0)), ((0,15,0), (0,0,0)),
# top face (h=31)
((0,0,31), (15,0,31)), ((15,0,31), (15,15,31)),
((15,15,31), (0,15,31)), ((0,15,31), (0,0,31)),
# vertical (h-axis) edges
((0,0,0), (0,0,31)), ((15,0,0), (15,0,31)),
((15,15,0), (15,15,31)), ((0,15,0), (0,15,31)),
]
wireframe_parts = []
for a, b in edges:
x1, y1 = corners[a]
x2, y2 = corners[b]
wireframe_parts.append(
f'')
wireframe_svg = '\n '.join(wireframe_parts)
# ═══════════════════════════════════════════════════════════════════════════
# Step 5: Assemble the full fig1_overview.svg with the new volume block
# ═══════════════════════════════════════════════════════════════════════════
def b64(p): return base64.b64encode(open(p, "rb").read()).decode()
rgb_b64 = b64("/data/cameron/penpot/figures/extracted/fig1a_rgb_hires.png")
pca_b64 = b64("/data/cameron/penpot/figures/extracted/fig1a_pca_hires.png")
hm_b64 = b64("/data/cameron/penpot/figures/extracted/fig1a_heatmap_sq.png")
# Panel (a) left: 4 row illustrations
r1_towel_b64 = b64("/data/cameron/penpot/figures/extracted/fig1v2/row1_towel.png")
r2_train_b64 = b64("/data/cameron/penpot/figures/extracted/fig1v2/row2_train.png")
r2_test_b64 = b64("/data/cameron/penpot/figures/extracted/fig1v2/row2_test_v3.png")
r3_svd_b64 = b64("/data/cameron/penpot/figures/extracted/fig1v2/row3_svd_clean.png")
r4_hand_b64 = b64("/data/cameron/penpot/figures/extracted/fig6_arm_deleted_0.png")
r4_libero_b64 = b64("/data/cameron/penpot/figures/extracted/fig1a_rgb_hires.png")
# fig1 v3 (advisor-meeting rebuild, 2026-05-12): 3 stacked OOD-axis rows
v3_r1_train_b64 = b64("/data/cameron/penpot/figures/extracted/fig1v3/row1_train.png")
v3_r1_test_b64 = b64("/data/cameron/penpot/figures/extracted/fig1v3/row1_test.png")
v3_r2_train_b64 = b64("/data/cameron/penpot/figures/extracted/fig1v3/row2_train.png")
v3_r2_test_b64 = b64("/data/cameron/penpot/figures/extracted/fig1v3/row2_test.png")
v3_r3_test_b64 = b64("/data/cameron/penpot/figures/extracted/fig1v3/row3_test.png")
v3_r3_ta_b64 = b64("/data/cameron/penpot/figures/extracted/fig1v3/row3_train_a.png")
v3_r3_tb_b64 = b64("/data/cameron/penpot/figures/extracted/fig1v3/row3_train_b.png")
v3_r3_tc_b64 = b64("/data/cameron/penpot/figures/extracted/fig1v3/row3_train_c.png")
svg = f'''
'''
with open("/data/cameron/para/paper/figs/svg/fig1_overview.svg", "w") as f:
f.write(svg)
print(f"[{time.time()-_t_start:5.2f}s] wrote fig1_overview.svg ({len(svg)} bytes)")