robot-folding / app /src /content /embeds /folding /subtask-heatmap.html
pepijn223's picture
pepijn223 HF Staff
Improve DAgger explainer, add conclusion and expand references
f0f3d44 unverified
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8"/>
<meta name="viewport" content="width=device-width, initial-scale=1.0"/>
<style>
:root { --bg: transparent; --text: #e8eaf0; --subtext: #8b8fa8; --grid: #2a2d3a; --border: #2a2d3a; }
* { box-sizing: border-box; margin: 0; padding: 0; }
body { background: var(--bg); font-family: system-ui, sans-serif; color: var(--text); }
.tooltip {
position: absolute; background: #1a1d27; border: 1px solid var(--border);
border-radius: 8px; padding: 10px 14px; font-size: 13px; pointer-events: none;
opacity: 0; transition: opacity .15s; z-index: 10; min-width: 220px;
box-shadow: 0 4px 16px rgba(0,0,0,.4);
}
.tooltip strong { display: block; margin-bottom: 5px; }
.tooltip-row { display: flex; justify-content: space-between; gap: 12px; margin-top: 3px; font-size: 12px; color: var(--subtext); }
.tooltip-row span:last-child { color: var(--text); font-weight: 600; }
.legend-bar { display: flex; align-items: center; gap: 8px; margin-top: 10px; font-size: 13px; color: var(--subtext); justify-content: center; }
.legend-gradient { height: 10px; width: 180px; border-radius: 5px; flex-shrink: 0; }
.hm-toggle { display: flex; gap: 0; justify-content: center; margin-bottom: 6px; }
.hm-btn {
padding: 7px 20px; font-size: 14px; font-weight: 600; cursor: pointer; border: none;
background: #1a1d2e; color: #8b8fa8; transition: all .2s;
font-family: system-ui, sans-serif; border-bottom: 2px solid transparent;
}
.hm-btn:first-child { border-radius: 6px 0 0 0; }
.hm-btn:last-child { border-radius: 0 6px 0 0; }
.hm-btn.active { background: #252540; color: #e8eaf0; border-bottom-color: #6366f1; }
.hm-btn:hover:not(.active) { background: #1f2238; color: #bbb; }
.filter-toggle {
font-size: 12px; color: #8b8fa8; cursor: pointer; border: 1px solid #2a2d3a;
background: #1a1d27; padding: 4px 12px; border-radius: 14px; transition: all .15s;
user-select: none; display: inline-flex; align-items: center; gap: 4px;
}
.filter-toggle:hover { border-color: #555; color: #e8eaf0; }
.filter-toggle.showing-all { background: rgba(248,147,79,0.12); border-color: #f7934f; color: #f7934f; }
.hm-controls { display: flex; justify-content: center; gap: 8px; margin-bottom: 6px; align-items: center; }
</style>
</head>
<body>
<div class="hm-controls">
<button class="filter-toggle" id="hm-filter-btn" onclick="window._hmToggleFilter()">Show all 11</button>
</div>
<div style="position:relative">
<svg id="hm-chart"></svg>
<div class="tooltip" id="hm-tooltip"></div>
</div>
<div class="legend-bar" id="lgd-bar">
<span id="lgd-lo">Fast (0s)</span>
<canvas id="lgd" class="legend-gradient" width="180" height="10"></canvas>
<span id="lgd-hi">Slow (≥120s)</span>
</div>
<script>
function _initSubtaskHeatmap() {
const rawData = [
{label:"1.1 π0",series:"1",total_sr:40, times:[null, 19.2, 42.22, 14.33, 19.88, 27.25], l1time:121.5, quality:2.70},
{label:"1.2 π0.5",series:"1",total_sr:20, times:[50, 39.27, 41.5, 12.3, 13.75, 10.75], l1time:90.75, quality:2.50},
{label:"1.3 Relative",series:"1",total_sr:35, times:[null, 19.5, 44.2, 14.8, 30.33, 22.14], l1time:113.86,quality:2.80},
{label:"1.4 RABC low",series:"1",total_sr:15, times:[null, 20.8, 36.62, 10.0, 18.8, 12.67], l1time:78.33, quality:2.20},
{label:"1.5 RABC high",series:"1",total_sr:0, times:[240, 21.4, 100.0, null, null, null ], l1time:null, quality:1.00},
{label:"1.7 Rel+RABC",series:"1",total_sr:40, times:[157.5,19.33, 32.64, 8.9, 11.0, 23.38], l1time:99.5, quality:2.30},
{label:"2.1 HQ",series:"2",total_sr:40, times:[77.5, 11.08, 21.09, 5.45, 5.5, 11.5 ], l1time:57.57, quality:2.80},
{label:"2.2 HQ+RABC+Rel",series:"2",total_sr:75, times:[34.33,6.25, 12.31, 3.75, 5.31, 8.93 ], l1time:43.2, quality:3.30},
{label:"2.3 HQ+mirror",series:"2",total_sr:5, times:[49, 14.0, 23.71, 17.5, 11.0, 4.0 ], l1time:null, quality:1.00},
{label:"2.4 HQ chunk45",series:"2",total_sr:20, times:[120, 10.09, 41.18, 7.89, 7.33, 10.0 ], l1time:72.5, quality:1.80},
{label:"2.5 HQ+RABC+Rel★",series:"2",total_sr:90, times:[62.25,8.28, 12.0, 5.28, 5.22, 6.83 ], l1time:40.8, quality:4.10},
];
const HIDDEN_BY_DEFAULT = new Set(["1.4 RABC low","1.5 RABC high","2.3 HQ+mirror","2.4 HQ chunk45"]);
let hmShowAll = false;
window._hmToggleFilter = function() {
hmShowAll = !hmShowAll;
const btn = document.getElementById("hm-filter-btn");
btn.textContent = hmShowAll ? "Key experiments" : "Show all 11";
btn.classList.toggle("showing-all", hmShowAll);
render();
};
function getFilteredData() {
const src = hmShowAll ? rawData : rawData.filter(d => !HIDDEN_BY_DEFAULT.has(d.label));
return [...src].sort((a,b) => {
if (a.l1time === null && b.l1time === null) return 0;
if (a.l1time === null) return 1;
if (b.l1time === null) return -1;
return a.l1time - b.l1time;
});
}
const subtasks = ["Unfold","Fold 1","Fold 2","Fold 3","Fold 4","Rotation"];
const seriesColor = s => s === "2" ? "#f7934f" : "#4f8ef7";
let mode = "rel";
// Compute per-subtask means over ALL experiments (ignoring nulls)
const subtaskMeans = subtasks.map((_, ci) => {
const vals = rawData.map(d => d.times[ci]).filter(v => v !== null);
return vals.length ? vals.reduce((a,b) => a+b, 0) / vals.length : 1;
});
function lerpColor(stops, t) {
for (let i=0; i<stops.length-1; i++) {
const [t0,c0]=stops[i],[t1,c1]=stops[i+1];
if (t<=t1) {
const f=(t-t0)/(t1-t0);
return `rgb(${Math.round(c0[0]+(c1[0]-c0[0])*f)},${Math.round(c0[1]+(c1[1]-c0[1])*f)},${Math.round(c0[2]+(c1[2]-c0[2])*f)})`;
}
}
return `rgb(${stops[stops.length-1][1].join(",")})`;
}
// Absolute: green → yellow → red, capped at 120s
const absColor = v => lerpColor([[0,[46,200,138]],[0.4,[247,211,79]],[1,[220,60,60]]], Math.min(v,120)/120);
// Relative: green (0.2x) → neutral (1x) → red (3x), diverging around mean
const relColor = ratio => {
const t = Math.max(0, Math.min(1, (ratio - 0.2) / (3 - 0.2)));
return lerpColor([[0,[46,200,138]],[0.286,[230,230,210]],[1,[220,60,60]]], t);
};
// Quality: red (1.0) → yellow (2.5) → green (4.5)
const qualColor = v => {
const t = Math.max(0, Math.min(1, (v - 1.0) / (4.5 - 1.0)));
return lerpColor([[0,[220,60,60]],[0.43,[247,211,79]],[1,[46,200,138]]], t);
};
function drawLegend() {
const canvas = document.getElementById("lgd");
const ctx = canvas.getContext("2d");
for (let i=0; i<180; i++) {
const t = i/179;
ctx.fillStyle = mode === "abs" ? absColor(t*120) : relColor(0.2 + t*2.8);
ctx.fillRect(i,0,1,10);
}
document.getElementById("lgd-lo").textContent = mode === "abs" ? "Fast (0s)" : "Fast (0.2×)";
document.getElementById("lgd-hi").textContent = mode === "abs" ? "Slow (≥120s)" : "Slow (≥3×)";
}
drawLegend();
const margin = {top:12, right:16, bottom:36, left:120};
const svg = d3.select("#hm-chart");
const container = svg.node().parentElement;
const tooltip = d3.select("#hm-tooltip");
const EXPERIMENTS = {
"1.1 π0": { note:"Base pi0 policy trained from scratch on the full dataset." },
"1.2 π0.5": { note:"Upgraded to pi0.5 architecture, same data and steps." },
"1.3 Relative": { note:"Adds Relative Actions, expressing actions relative to current state." },
"1.4 RABC low": { note:"SARM with low κ (≈ mean threshold, not very selective)." },
"1.5 RABC high": { note:"SARM with κ = mean + ½ std, more selective filtering than 1.4." },
"1.7 Rel+RABC": { note:"Best of initial training. Base checkpoint for 2.5." },
"2.1 HQ": { note:"Fine-tunes 1.3 on curated high-quality data only." },
"2.2 HQ+RABC+Rel": { note:"Adds RABC on high-quality fine-tune from 1.3." },
"2.3 HQ+mirror": { note:"Augments the high-quality dataset with mirrored trajectories." },
"2.4 HQ chunk45": { note:"Larger action chunk size (45 vs default 30)." },
"2.5 HQ+RABC+Rel★": { note:"Top performer. Best overall result." },
};
const GAP = 10;
const sumCols = ["L1 Time","Quality"];
function render() {
svg.selectAll("*").remove();
const data = getFilteredData();
const W = container.clientWidth;
const totalCols = subtasks.length + sumCols.length;
const cellW = Math.floor((W - margin.left - margin.right - GAP) / totalCols);
const cellH = Math.max(34, Math.min(44, cellW * 0.7));
const H = data.length * cellH + margin.top + margin.bottom;
svg.attr("width",W).attr("height",H);
const g = svg.append("g").attr("transform",`translate(${margin.left},${margin.top})`);
const gridW = cellW * subtasks.length;
const gridH = cellH * data.length;
const sumX = gridW + GAP;
// Subtask column labels
g.selectAll(".col-lbl").data(subtasks).join("text")
.attr("x",(_,i)=>i*cellW+cellW/2).attr("y",gridH+22)
.attr("text-anchor","middle").attr("fill","#8b8fa8").attr("font-size",13)
.text(d=>d);
// Summary column labels (italic, slightly muted)
sumCols.forEach((label, i) => {
g.append("text")
.attr("x", sumX + i * cellW + cellW / 2).attr("y", gridH + 22)
.attr("text-anchor","middle").attr("fill","#6b6f88").attr("font-size",12)
.attr("font-style","italic")
.text(label);
});
// Vertical divider between subtask and summary columns
g.append("line")
.attr("x1", gridW + GAP / 2).attr("x2", gridW + GAP / 2)
.attr("y1", -2).attr("y2", gridH)
.attr("stroke","#2a2d3a").attr("stroke-width",1);
// Row labels + series stripe
data.forEach((d, ri) => {
g.append("rect")
.attr("x",-margin.left+2).attr("y",ri*cellH+2)
.attr("width",4).attr("height",cellH-4).attr("rx",2)
.attr("fill",seriesColor(d.series)).attr("opacity",0.9);
g.append("text")
.attr("x",-8).attr("y",ri*cellH+cellH/2)
.attr("text-anchor","end").attr("fill","#e8eaf0").attr("font-size",12).attr("font-weight","500")
.text(d.label);
g.append("text")
.attr("x",-8).attr("y",ri*cellH+cellH/2+12)
.attr("text-anchor","end").attr("fill","#8b8fa8").attr("font-size",10)
.text(d.total_sr+"% SR");
});
// Helper: draw a single heatmap cell
function drawCell(parent, x, y, w, h, val, col, tooltipHtml) {
const cg = parent.append("g").attr("transform",`translate(${x},${y})`);
if (val === null) return;
const lum = d3.hsl(col).l;
cg.append("rect").attr("width",w-3).attr("height",h-3).attr("rx",4)
.attr("fill",col).style("cursor","pointer")
.on("mousemove",function(event){
tooltip.style("opacity",1).html(tooltipHtml);
const bx=container.getBoundingClientRect();
const ex=event.clientX-bx.left, ey=event.clientY-bx.top;
tooltip.style("left",Math.min(ex+12,W-220)+"px").style("top",Math.max(ey-70,0)+"px");
})
.on("mouseleave",()=>tooltip.style("opacity",0));
const fs = Math.max(10, Math.min(14, w * 0.24));
cg.append("text")
.attr("x",w/2-1.5).attr("y",h/2+4)
.attr("text-anchor","middle")
.attr("fill", lum > 0.52 ? "#1a1d27" : "#e8eaf0")
.attr("font-size", fs).attr("font-weight","600")
.text(typeof val === "string" ? val : (val >= 100 ? Math.round(val)+"s" : val.toFixed(1)+"s"));
}
// Subtask cells
data.forEach((d, ri) => {
d.times.forEach((val, ci) => {
if (val === null) {
// skip — leave empty
} else {
const ratio = val / subtaskMeans[ci];
const col = mode === "abs" ? absColor(val) : relColor(ratio);
const cellLabel = mode === "abs"
? (val>=100 ? Math.round(val)+"s" : val.toFixed(1)+"s")
: ratio.toFixed(1)+"×";
const ttHtml = `
<strong>Experiment ${d.label}: ${subtasks[ci]}</strong>
<div class="tooltip-row"><span>Avg time</span><span>${val.toFixed(1)}s</span></div>
<div class="tooltip-row"><span>Subtask mean</span><span>${subtaskMeans[ci].toFixed(1)}s</span></div>
<div class="tooltip-row"><span>Relative</span><span>${ratio.toFixed(2)}×</span></div>
<div class="tooltip-row"><span>Total SR</span><span>${d.total_sr}%</span></div>
`;
const cg = g.append("g").attr("transform",`translate(${ci*cellW},${ri*cellH})`);
const lum = d3.hsl(col).l;
cg.append("rect").attr("width",cellW-3).attr("height",cellH-3).attr("rx",4)
.attr("fill",col).style("cursor","pointer")
.on("mousemove",function(event){
tooltip.style("opacity",1).html(ttHtml);
const bx=container.getBoundingClientRect();
const ex=event.clientX-bx.left, ey=event.clientY-bx.top;
tooltip.style("left",Math.min(ex+12,W-220)+"px").style("top",Math.max(ey-70,0)+"px");
})
.on("mouseleave",()=>tooltip.style("opacity",0));
const fs = Math.max(10, Math.min(14, cellW * 0.24));
cg.append("text")
.attr("x",cellW/2-1.5).attr("y",cellH/2+4)
.attr("text-anchor","middle")
.attr("fill", lum > 0.52 ? "#1a1d27" : "#e8eaf0")
.attr("font-size", fs).attr("font-weight","600")
.text(cellLabel);
}
});
});
// Summary columns: L1 Time and Quality
data.forEach((d, ri) => {
const note = (EXPERIMENTS[d.label] || {}).note || "";
// L1 Time cell
const l1 = d.l1time;
const l1Col = l1 !== null ? absColor(l1) : null;
const l1Label = l1 !== null ? (l1 >= 100 ? Math.round(l1) + "s" : l1.toFixed(1) + "s") : null;
const l1Tt = l1 !== null ? `
<strong>${d.label}: L1 Completion Time</strong>
<div class="tooltip-row"><span>L1 time</span><span>${l1.toFixed(1)}s</span></div>
<div class="tooltip-row"><span>Total SR</span><span>${d.total_sr}%</span></div>
${note ? `<div style="margin-top:5px;font-size:10px;color:#555;border-top:1px solid #2a2d3a;padding-top:5px">${note}</div>` : ""}
` : "";
drawCell(g, sumX, ri * cellH, cellW, cellH, l1 !== null ? l1Label : null, l1Col, l1Tt);
// Quality cell
const q = d.quality;
const qCol = qualColor(q);
const qLabel = q.toFixed(1);
const qTt = `
<strong>${d.label}: Fold Quality</strong>
<div class="tooltip-row"><span>Quality</span><span>${q.toFixed(2)} / 5</span></div>
<div class="tooltip-row"><span>Total SR</span><span>${d.total_sr}%</span></div>
${note ? `<div style="margin-top:5px;font-size:10px;color:#555;border-top:1px solid #2a2d3a;padding-top:5px">${note}</div>` : ""}
`;
drawCell(g, sumX + cellW, ri * cellH, cellW, cellH, qLabel, qCol, qTt);
});
// Sorted annotation
const totalW = gridW + GAP + cellW * sumCols.length;
g.append("text").attr("x",totalW).attr("y",-4).attr("text-anchor","end")
.attr("fill","#8b8fa8").attr("font-size",11)
.text("↓ sorted: fastest → slowest L1 time");
}
render();
window.addEventListener("resize", render);
}
if (typeof d3 !== "undefined") {
_initSubtaskHeatmap();
} else {
var s = document.createElement("script");
s.src = "https://cdnjs.cloudflare.com/ajax/libs/d3/7.9.0/d3.min.js";
s.onload = _initSubtaskHeatmap;
document.head.appendChild(s);
}
</script>
</body>
</html>