Buckets:

twanghcmut
/

ego-real2sim

Files

xet

twanghcmut/ego-real2sim / basic_pick_place /code /examples /refine_physics.py

twanghcmut

12 days ago

download

raw

27.1 kB

import argparse
import json
from pathlib import Path

import cv2
import numpy as np
from fit_physics import (
    ACTORS,
    INVALID_LOSS,
    SIMULATION_ERRORS,
    contact_diagnostics,
    flattened_arrays,
    labeled_frame,
    lift_metrics,
    load_real_frames,
    plot_loss,
    regularization,
    resize_masks,
    serializable,
    validate_rollout,
    write_json,
)

from real2sim import video
from real2sim.error_trace import MaskErrorTrace
from real2sim.mask_loss import MaskLoss
from real2sim.optimization import search

BLOCKS = {
    "object_pose": np.arange(0, 6),
    "left_wrist": np.arange(6, 15),
    "right_wrist": np.arange(15, 24),
    "friction": np.array([24]),
}
ACTOR_COVERAGE = {
    "object_pose": ("object",),
    "left_wrist": ("left", "object"),
    "right_wrist": ("right", "object"),
    "friction": ("object",),
}


def load_parameters(path):
    if path.is_dir():
        options = ("best_parameters.json", "best_parameters.npz", "parameters.npz")
        path = next((path / name for name in options if (path / name).is_file()), path)
    if path.suffix == ".npz":
        with np.load(path, allow_pickle=False) as saved:
            values = saved["parameters"].copy()
    else:
        value = json.loads(path.read_text())
        values = value["parameters"] if isinstance(value, dict) else value
    values = np.asarray(values, dtype=float)
    if values.shape != (25,) or not np.isfinite(values).all():
        raise ValueError("Baseline must provide 25 finite physics parameters")
    return values


def geom_mapping(scene, result):
    mapping = result.get("geom_mapping")
    if mapping is not None:
        return serializable(mapping)
    ids = result.get("geom_actor_ids", getattr(scene, "geom_actor_ids", None))
    names = result.get("actor_names", ACTORS)
    values = {}
    for index in range(scene.model.ngeom):
        actor_id = int(ids[index]) if ids is not None else -1
        values[str(index)] = {
            "name": scene.model.geom(index).name,
            "actor": str(names[actor_id])
            if 0 <= actor_id < len(names)
            else "background",
            "actor_id": actor_id + 1 if actor_id >= 0 else 0,
        }
    return values


def make_trace(scene, result, targets, confidence, observed, source_indices):
    return MaskErrorTrace(
        targets,
        result["masks"],
        confidence,
        observed,
        scene.timestamps,
        source_indices,
        ACTORS,
        result.get("geom_ids"),
        geom_mapping(scene, result),
    )


def trace_overlay(frames, trace):
    output = np.asarray(frames, dtype=np.uint8).copy()
    positive = trace.false_positive.any(axis=0)
    negative = trace.false_negative.any(axis=0)
    for selected, color in (
        (positive & ~negative, (255, 70, 70)),
        (negative & ~positive, (60, 210, 255)),
        (positive & negative, (255, 175, 45)),
    ):
        output[selected] = (0.45 * output[selected] + 0.55 * np.asarray(color)).astype(
            np.uint8
        )
    untrusted = trace.disagreement & ~trace.trusted_disagreement
    output[untrusted] = (
        0.5 * output[untrusted] + 0.5 * np.array([150, 150, 150])
    ).astype(np.uint8)
    return output


def save_trace(folder, trace, frames, fps, prefix=""):
    trace.save(folder / f"{prefix}trace.npz")
    write_json(folder / f"{prefix}trace.json", trace.summary())
    heat = trace.heatmap()
    heat = np.clip(255 * heat / max(float(heat.max()), 1e-12), 0, 255).astype(np.uint8)
    colored = cv2.applyColorMap(heat, cv2.COLORMAP_INFERNO)
    cv2.imwrite(str(folder / f"{prefix}error_heatmap.png"), colored)
    video.write(
        folder / f"{prefix}trace_overlay.mp4", trace_overlay(frames, trace), fps
    )


def block_is_admissible(name, entries, baseline, target_actor):
    if ACTORS[target_actor] in ACTOR_COVERAGE[name]:
        return True
    reference = baseline["actor_losses"][target_actor]
    return any(
        item["stable"]
        and item["actor_losses"][target_actor] is not None
        and abs(item["actor_losses"][target_actor] - reference) > 1e-8
        for item in entries
    )


def candidate_is_acceptable(candidate, baseline, target_actor):
    target_loss = candidate["actor_losses"][target_actor]
    return (
        candidate["stable"]
        and target_loss is not None
        and target_loss <= baseline["actor_losses"][target_actor] + 1e-9
        and candidate["loss"] <= baseline["loss"]
    )


def guarded_objective(candidate, baseline, target_actor):
    if not candidate["stable"]:
        return candidate["loss"]
    change = (
        candidate["actor_losses"][target_actor] - baseline["actor_losses"][target_actor]
    )
    if change > 1e-9:
        return INVALID_LOSS + candidate["loss"] + change
    return candidate["loss"]


def include_query_frame(selected, source_indices, frame):
    matches = np.flatnonzero(np.asarray(source_indices) == frame)
    if len(matches) != 1:
        raise ValueError(f"Source frame {frame} is absent or ambiguous")
    return np.unique(np.append(selected, int(matches[0])))


def choose_block(sensitivity, baseline, target_actor, trace_scores):
    ranked = []
    for name, entries in sensitivity.items():
        if not block_is_admissible(name, entries, baseline, target_actor):
            continue
        stable = [item for item in entries if item["stable"]]
        if not stable:
            continue
        combined_gain = max(baseline["loss"] - item["loss"] for item in stable)
        target_gain = max(
            baseline["actor_losses"][target_actor] - item["actor_losses"][target_actor]
            for item in stable
        )
        coverage = sum(
            trace_scores[ACTORS.index(actor)] for actor in ACTOR_COVERAGE[name]
        )
        ranked.append((name, combined_gain, target_gain, float(coverage)))
    if not ranked:
        raise RuntimeError(
            "No stable admissible parameter block affects the target actor"
        )
    improving = [item for item in ranked if item[1] > 1e-9]
    if improving:
        selected = max(improving, key=lambda item: (item[1], item[2], item[3]))
        reason = "combined_objective_improvement"
    else:
        improving = [item for item in ranked if item[2] > 1e-9]
        if improving:
            selected = max(improving, key=lambda item: (item[2], item[1], item[3]))
            reason = "target_actor_mask_improvement"
        else:
            selected = max(ranked, key=lambda item: (item[3], item[1], item[2]))
            reason = "actor_coverage_after_inconclusive_perturbations"
    return selected[0], reason


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("folder", type=Path)
    parser.add_argument("--asset", type=Path, required=True)
    parser.add_argument("--physics", type=Path)
    parser.add_argument("--baseline", type=Path)
    parser.add_argument("--reference", type=Path)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--population", type=int, default=8)
    parser.add_argument("--iterations", type=int, default=4)
    parser.add_argument("--stride", type=int)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--perturbation", type=float, default=0.5)
    parser.add_argument("--std-scale", type=float, default=0.5)
    parser.add_argument("--pixel", type=int, nargs=2, metavar=("X", "Y"))
    parser.add_argument("--frame", type=int)
    parser.add_argument("--trace", type=Path)
    parser.add_argument("--query-only", action="store_true")
    parser.add_argument("--trace-only", action="store_true")
    args = parser.parse_args()
    if (args.pixel is None) != (args.frame is None):
        parser.error("pixel and frame must be supplied together")
    if (
        args.population < 2
        or args.iterations < 1
        or (args.stride is not None and args.stride < 1)
    ):
        parser.error("population must be >= 2; iterations and stride must be positive")
    if (
        not np.isfinite([args.perturbation, args.std_scale]).all()
        or min(args.perturbation, args.std_scale) <= 0
    ):
        parser.error("perturbation and std-scale must be finite and positive")
    folder = args.folder.resolve()
    physics = (args.physics or folder / "physics").resolve()
    output = (args.output or physics / "refinement").resolve()
    if args.query_only:
        if args.pixel is None:
            parser.error("query-only requires pixel and frame")
        path = args.trace or output / "trace.npz"
        if not path.is_file():
            parser.error(f"trace does not exist: {path}")
        try:
            query = MaskErrorTrace.load(path).query(args.frame, *args.pixel)
        except ValueError as error:
            parser.error(str(error))
        print(json.dumps(query, indent=2), flush=True)
        return
    baseline_path = (args.baseline or physics / "best_parameters.json").resolve()
    if not baseline_path.exists():
        parser.error(f"baseline parameters are not ready: {baseline_path}")
    if output in (folder, physics, baseline_path, baseline_path.parent):
        parser.error("output must be a separate refinement folder")
    if not args.asset.is_file():
        parser.error(f"asset does not exist: {args.asset}")
    baseline_parameters = load_parameters(baseline_path)
    config_path = physics / "search_config.json"
    config = json.loads(config_path.read_text()) if config_path.is_file() else {}
    with np.load(folder / "real_masks.npz", allow_pickle=False) as saved:
        if "model" not in saved or "sam3.1" not in str(saved["model"].item()).lower():
            parser.error("real_masks.npz must identify SAM3.1 masks")
        mask_model = str(saved["model"].item())
        order = [saved["names"].tolist().index(name) for name in ACTORS]
        masks = saved["masks"][order]
        confidence = saved["confidence"][order].astype(float)
        observed = saved["observed"][order].astype(bool)
        source_indices = saved["frame_indices"].copy()
        mask_times = saved["timestamps"].copy()
    from real2sim.sim.physics import PhysicsScene

    scene = PhysicsScene(folder, args.asset.resolve(), reference=args.reference)
    try:
        times = np.asarray(scene.timestamps)
        if len(times) != masks.shape[1] or not np.allclose(
            times, mask_times, atol=1e-6, rtol=0
        ):
            raise ValueError("Mask and physics timestamps must align")
        mean = np.asarray(
            config.get("parameter_mean", scene.parameter_mean), dtype=float
        )
        std = np.asarray(config.get("parameter_std", scene.parameter_std), dtype=float)
        bounds = np.asarray(scene.parameter_bounds, dtype=float)
        if (
            mean.shape != (25,)
            or std.shape != (25,)
            or not np.isfinite(std).all()
            or np.any(std <= 0)
        ):
            raise ValueError(
                "Baseline must provide a valid 25-parameter regularization reference"
            )
        if np.any(baseline_parameters < bounds[:, 0]) or np.any(
            baseline_parameters > bounds[:, 1]
        ):
            raise ValueError("Baseline parameters exceed current physics bounds")
        size = tuple(config.get("render_size", (256, 144)))
        if len(size) != 2 or min(size) < 1:
            raise ValueError(
                "Baseline render size must contain positive width and height"
            )
        if args.stride is None and "sampled_frame_indices" in config:
            selected = np.asarray(config["sampled_frame_indices"], dtype=int)
        else:
            selected = np.unique(
                np.append(np.arange(0, len(times), args.stride or 5), len(times) - 1)
            )
        if args.frame is not None:
            try:
                selected = include_query_frame(selected, source_indices, args.frame)
            except ValueError as error:
                parser.error(str(error))
        if (
            selected.ndim != 1
            or not len(selected)
            or np.any(np.diff(selected) <= 0)
            or selected[0] < 0
            or selected[-1] >= len(times)
        ):
            raise ValueError("Sampled frame indices must increase within the clip")
        targets = resize_masks(masks, size)
        weights = confidence[:, selected] * observed[:, selected]
        combined_loss = MaskLoss(targets[:, selected], weights)
        actor_losses = [
            MaskLoss(targets[actor : actor + 1, selected], weights[actor : actor + 1])
            if weights[actor].any()
            else None
            for actor in range(len(ACTORS))
        ]
        full_loss = MaskLoss(targets, confidence * observed)
        all_indices = np.arange(len(times))
        fps = 1 / float(np.median(np.diff(times)))
        output.mkdir(parents=True, exist_ok=True)
        progress_path = output / "progress.jsonl"
        progress_path.write_text("")
        evaluations = 0
        failures = 0

        def progress(event):
            line = json.dumps(serializable(event), allow_nan=False)
            with progress_path.open("a") as stream:
                stream.write(line + "\n")
            print(line, flush=True)

        def evaluate(parameters):
            nonlocal evaluations, failures
            evaluations += 1
            penalty = regularization(parameters, mean, std)
            try:
                result = scene.rollout(parameters, selected, render_size=size)
                validate_rollout(result, len(selected), size)
                if contact_diagnostics(result).get("unstable", False):
                    raise FloatingPointError(
                        "PhysicsScene reported an unstable rollout"
                    )
                loss = combined_loss(result["masks"])
                components = [
                    function(result["masks"][actor : actor + 1]) if function else None
                    for actor, function in enumerate(actor_losses)
                ]
                if not np.isfinite(loss):
                    raise FloatingPointError("Nonfinite refinement objective")
                return {
                    "parameters": parameters.copy(),
                    "loss": float(loss + penalty),
                    "mask_loss": float(loss),
                    "regularization": penalty,
                    "actor_losses": components,
                    "stable": True,
                }
            except SIMULATION_ERRORS as error:
                failures += 1
                return {
                    "parameters": parameters.copy(),
                    "loss": INVALID_LOSS + penalty,
                    "regularization": penalty,
                    "actor_losses": [None] * len(ACTORS),
                    "stable": False,
                    "error": str(error),
                }

        baseline_full = scene.rollout(
            baseline_parameters, all_indices, render_size=size, rgb=True
        )
        validate_rollout(baseline_full, len(times), size)
        baseline_trace = make_trace(
            scene, baseline_full, targets, confidence, observed, source_indices
        )
        real_frames = load_real_frames(folder, len(times), size)
        save_trace(output, baseline_trace, real_frames, fps)
        query = None
        if args.pixel is not None:
            try:
                query = baseline_trace.query(args.frame, *args.pixel)
            except ValueError as error:
                parser.error(str(error))
            write_json(output / "pixel_query.json", query)
            progress({"event": "pixel_query", **query})
        if args.trace_only:
            write_json(
                output / "metrics.json",
                {
                    "stage": "physics_pixel_error_trace",
                    "trace": baseline_trace.summary(),
                    "pixel_query": query,
                },
            )
            progress({"event": "trace_completed", "output": output})
            return
        baseline = evaluate(baseline_parameters)
        if not baseline["stable"]:
            raise RuntimeError("Baseline physics rollout is unstable")
        trace_scores, _ = baseline_trace.error_scores()
        eligible = np.asarray([function is not None for function in actor_losses])
        target_actor = int(np.argmax(np.where(eligible, trace_scores, -1)))
        target_selection = "video_worst_observed_actor"
        if query is not None:
            affected = np.asarray(
                [
                    eligible[actor]
                    and query["observed"][name]
                    and any(query["actor_errors"][name].values())
                    for actor, name in enumerate(ACTORS)
                ]
            )
            if affected.any():
                target_actor = int(np.argmax(np.where(affected, trace_scores, -1)))
                target_selection = "pixel_affected_actor_with_largest_video_error"
        progress(
            {
                "event": "baseline",
                "loss": baseline["loss"],
                "target_actor": ACTORS[target_actor],
                "target_selection": target_selection,
                "actor_losses": baseline["actor_losses"],
            }
        )
        sensitivity = {}
        candidates = [{"source": "baseline", **baseline}]
        random = np.random.default_rng(args.seed)
        for name, indices in BLOCKS.items():
            direction = random.normal(size=len(indices))
            direction *= np.sqrt(len(indices)) / max(
                float(np.linalg.norm(direction)), 1e-12
            )
            delta = args.perturbation * std[indices] * direction
            sensitivity[name] = []
            for sign in (-1, 1):
                candidate = baseline_parameters.copy()
                candidate[indices] = np.clip(
                    candidate[indices] + sign * delta,
                    bounds[indices, 0],
                    bounds[indices, 1],
                )
                measured = {"block": name, "sign": sign, **evaluate(candidate)}
                measured["combined_objective_change"] = (
                    measured["loss"] - baseline["loss"]
                )
                measured["target_actor_loss_change"] = (
                    measured["actor_losses"][target_actor]
                    - baseline["actor_losses"][target_actor]
                    if measured["stable"]
                    else None
                )
                sensitivity[name].append(measured)
                candidates.append({"source": f"sensitivity_{name}_{sign}", **measured})
                progress(
                    {
                        "event": "sensitivity",
                        "block": name,
                        "sign": sign,
                        "loss": measured["loss"],
                        "actor_losses": measured["actor_losses"],
                        "stable": measured["stable"],
                    }
                )
        admissible_blocks = [
            name
            for name, entries in sensitivity.items()
            if block_is_admissible(name, entries, baseline, target_actor)
        ]
        block, reason = choose_block(sensitivity, baseline, target_actor, trace_scores)
        indices = BLOCKS[block]
        selected_candidates = [
            baseline,
            *[
                item
                for item in sensitivity[block]
                if candidate_is_acceptable(item, baseline, target_actor)
            ],
        ]
        center = min(selected_candidates, key=lambda item: item["loss"])[
            "parameters"
        ].copy()
        progress(
            {
                "event": "block_selected",
                "block": block,
                "reason": reason,
                "target_actor": ACTORS[target_actor],
            }
        )
        write_json(
            output / "sensitivity.json",
            {
                "target_actor": ACTORS[target_actor],
                "selected_block": block,
                "admissible_blocks": admissible_blocks,
                "target_selection": target_selection,
                "selection_reason": reason,
                "blocks": sensitivity,
                "parameter_indices": BLOCKS,
                "actor_coverage": ACTOR_COVERAGE,
                "physical_cause_uniquely_identified": False,
            },
        )

        def objective(values):
            parameters = baseline_parameters.copy()
            parameters[indices] = values
            return guarded_objective(evaluate(parameters), baseline, target_actor)

        result = search(
            objective,
            center[indices],
            std[indices] * args.std_scale,
            bounds[indices],
            population=args.population,
            iterations=args.iterations,
            seed=args.seed,
            method="cem",
            callback=lambda iteration, loss: progress(
                {
                    "event": "block_cem",
                    "block": block,
                    "iteration": iteration,
                    "loss": loss,
                }
            ),
        )
        parameters = baseline_parameters.copy()
        parameters[indices] = result.parameters
        candidates.append({"source": "block_cem", **evaluate(parameters)})
        best = min(
            (
                item
                for item in candidates
                if candidate_is_acceptable(item, baseline, target_actor)
                and (
                    item["source"] in ("baseline", "block_cem")
                    or item.get("block") in admissible_blocks
                )
            ),
            key=lambda item: item["loss"],
        )
        final = scene.rollout(
            best["parameters"], all_indices, render_size=size, rgb=True
        )
        validate_rollout(final, len(times), size)
        scene.save(output)
        arrays = flattened_arrays(final)
        arrays.update(
            parameters=best["parameters"],
            names=np.asarray(ACTORS),
            timestamps=times,
            frame_indices=source_indices,
        )
        np.savez_compressed(output / "rollout.npz", **arrays)
        np.savez_compressed(
            output / "best_parameters.npz",
            parameters=best["parameters"],
            parameter_mean=mean,
            parameter_std=std,
            parameter_bounds=bounds,
            timestamps=times,
            frame_indices=source_indices,
            sampled_frame_indices=selected,
            selected_block=block,
            seed=args.seed,
        )
        write_json(
            output / "best_parameters.json",
            {
                "parameters": best["parameters"],
                "loss": best["loss"],
                "source": best["source"],
            },
        )
        final_trace = make_trace(
            scene, final, targets, confidence, observed, source_indices
        )
        save_trace(output, final_trace, real_frames, fps, prefix="refined_")
        video.write(output / "replay.mp4", final["rgb"], fps)
        comparisons = [
            np.concatenate(
                (
                    labeled_frame(real, "Real"),
                    labeled_frame(before, "Baseline physics"),
                    labeled_frame(after, "Refined physics"),
                ),
                axis=1,
            )
            for real, before, after in zip(
                real_frames, baseline_full["rgb"], final["rgb"]
            )
        ]
        video.write(output / "comparison.mp4", comparisons, fps)
        cv2.imwrite(
            str(output / "preview.png"),
            cv2.cvtColor(comparisons[len(comparisons) // 2], cv2.COLOR_RGB2BGR),
        )
        plot_loss(
            output / "loss.png",
            {f"{block}_cem": result},
            baseline["loss"],
            args.population,
        )
        measurements = lift_metrics(final)
        metrics = {
            "stage": "pixel_trace_physics_block_refinement",
            "simulation_mode": "dynamic_object_contact_physics",
            "mask_model": mask_model,
            "mask_actor_order": ACTORS,
            "baseline_parameters": baseline_parameters,
            "baseline_file": baseline_path,
            "asset": args.asset.resolve(),
            "reference": args.reference,
            "parameter_reference_mean": mean,
            "parameter_reference_std": std,
            "sampled_frame_indices": selected,
            "queried_frame_in_objective": args.frame is not None,
            "render_size": size,
            "target_actor": ACTORS[target_actor],
            "target_selection": target_selection,
            "selected_block": block,
            "admissible_blocks": admissible_blocks,
            "selection_reason": reason,
            "acceptance_constraints": {
                "target_actor_maximum_loss_increase": 1e-9,
                "combined_objective_must_not_increase": True,
                "fallback": "baseline",
            },
            "cem_objective": "combined_objective_with_target_actor_worsening_penalty",
            "physical_cause_uniquely_identified": False,
            "mass_kg": 0.1,
            "friction": float(best["parameters"][24]),
            "baseline_loss": baseline["loss"],
            "best_loss": best["loss"],
            "objective_improvement": baseline["loss"] - best["loss"],
            "best_source": best["source"],
            "baseline_actor_losses": baseline["actor_losses"],
            "best_actor_losses": best["actor_losses"],
            "target_actor_loss_change": best["actor_losses"][target_actor]
            - baseline["actor_losses"][target_actor],
            "baseline_full_mask_loss": full_loss(baseline_full["masks"]),
            "refined_full_mask_loss": full_loss(final["masks"]),
            "population": args.population,
            "iterations": args.iterations,
            "seed": args.seed,
            "sensitivity_evaluations": 2 * len(BLOCKS),
            "cem_evaluations": result.evaluations,
            "total_objective_evaluations": evaluations,
            "full_export_rollouts": 2,
            "unstable_candidates": failures,
            "cem_history": result.history,
            "pixel_query": query,
            "baseline_trace": baseline_trace.summary(),
            "refined_trace": final_trace.summary(),
            "baseline_contact_and_lift": lift_metrics(baseline_full),
            "refined_contact_and_lift": measurements,
        }
        write_json(output / "metrics.json", metrics)
        progress(
            {
                "event": "completed",
                "output": output,
                "baseline_loss": baseline["loss"],
                "best_loss": best["loss"],
                "best_source": best["source"],
                "selected_block": block,
                "maximum_lift_m": measurements["maximum_lift_m"],
            }
        )
    finally:
        scene.close()


if __name__ == "__main__":
    main()

Xet Storage Details

Size:: 27.1 kB
Xet hash:: aa863dbde76dcf30ae0a33fc99307d0bfa0e01354b05e1a8357370bcd62f7626

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	import argparse
	import json
	from pathlib import Path

	import cv2
	import numpy as np
	from fit_physics import (
	ACTORS,
	INVALID_LOSS,
	SIMULATION_ERRORS,
	contact_diagnostics,
	flattened_arrays,
	labeled_frame,
	lift_metrics,
	load_real_frames,
	plot_loss,
	regularization,
	resize_masks,
	serializable,
	validate_rollout,
	write_json,
	)

	from real2sim import video
	from real2sim.error_trace import MaskErrorTrace
	from real2sim.mask_loss import MaskLoss
	from real2sim.optimization import search

	BLOCKS = {
	"object_pose": np.arange(0, 6),
	"left_wrist": np.arange(6, 15),
	"right_wrist": np.arange(15, 24),
	"friction": np.array([24]),
	}
	ACTOR_COVERAGE = {
	"object_pose": ("object",),
	"left_wrist": ("left", "object"),
	"right_wrist": ("right", "object"),
	"friction": ("object",),
	}


	def load_parameters(path):
	if path.is_dir():
	options = ("best_parameters.json", "best_parameters.npz", "parameters.npz")
	path = next((path / name for name in options if (path / name).is_file()), path)
	if path.suffix == ".npz":
	with np.load(path, allow_pickle=False) as saved:
	values = saved["parameters"].copy()
	else:
	value = json.loads(path.read_text())
	values = value["parameters"] if isinstance(value, dict) else value
	values = np.asarray(values, dtype=float)
	if values.shape != (25,) or not np.isfinite(values).all():
	raise ValueError("Baseline must provide 25 finite physics parameters")
	return values


	def geom_mapping(scene, result):
	mapping = result.get("geom_mapping")
	if mapping is not None:
	return serializable(mapping)
	ids = result.get("geom_actor_ids", getattr(scene, "geom_actor_ids", None))
	names = result.get("actor_names", ACTORS)
	values = {}
	for index in range(scene.model.ngeom):
	actor_id = int(ids[index]) if ids is not None else -1
	values[str(index)] = {
	"name": scene.model.geom(index).name,
	"actor": str(names[actor_id])
	if 0 <= actor_id < len(names)
	else "background",
	"actor_id": actor_id + 1 if actor_id >= 0 else 0,
	}
	return values


	def make_trace(scene, result, targets, confidence, observed, source_indices):
	return MaskErrorTrace(
	targets,
	result["masks"],
	confidence,
	observed,
	scene.timestamps,
	source_indices,
	ACTORS,
	result.get("geom_ids"),
	geom_mapping(scene, result),
	)


	def trace_overlay(frames, trace):
	output = np.asarray(frames, dtype=np.uint8).copy()
	positive = trace.false_positive.any(axis=0)
	negative = trace.false_negative.any(axis=0)
	for selected, color in (
	(positive & ~negative, (255, 70, 70)),
	(negative & ~positive, (60, 210, 255)),
	(positive & negative, (255, 175, 45)),
	):
	output[selected] = (0.45 * output[selected] + 0.55 * np.asarray(color)).astype(
	np.uint8
	)
	untrusted = trace.disagreement & ~trace.trusted_disagreement
	output[untrusted] = (
	0.5 * output[untrusted] + 0.5 * np.array([150, 150, 150])
	).astype(np.uint8)
	return output


	def save_trace(folder, trace, frames, fps, prefix=""):
	trace.save(folder / f"{prefix}trace.npz")
	write_json(folder / f"{prefix}trace.json", trace.summary())
	heat = trace.heatmap()
	heat = np.clip(255 * heat / max(float(heat.max()), 1e-12), 0, 255).astype(np.uint8)
	colored = cv2.applyColorMap(heat, cv2.COLORMAP_INFERNO)
	cv2.imwrite(str(folder / f"{prefix}error_heatmap.png"), colored)
	video.write(
	folder / f"{prefix}trace_overlay.mp4", trace_overlay(frames, trace), fps
	)


	def block_is_admissible(name, entries, baseline, target_actor):
	if ACTORS[target_actor] in ACTOR_COVERAGE[name]:
	return True
	reference = baseline["actor_losses"][target_actor]
	return any(
	item["stable"]
	and item["actor_losses"][target_actor] is not None
	and abs(item["actor_losses"][target_actor] - reference) > 1e-8
	for item in entries
	)


	def candidate_is_acceptable(candidate, baseline, target_actor):
	target_loss = candidate["actor_losses"][target_actor]
	return (
	candidate["stable"]
	and target_loss is not None
	and target_loss <= baseline["actor_losses"][target_actor] + 1e-9
	and candidate["loss"] <= baseline["loss"]
	)


	def guarded_objective(candidate, baseline, target_actor):
	if not candidate["stable"]:
	return candidate["loss"]
	change = (
	candidate["actor_losses"][target_actor] - baseline["actor_losses"][target_actor]
	)
	if change > 1e-9:
	return INVALID_LOSS + candidate["loss"] + change
	return candidate["loss"]


	def include_query_frame(selected, source_indices, frame):
	matches = np.flatnonzero(np.asarray(source_indices) == frame)
	if len(matches) != 1:
	raise ValueError(f"Source frame {frame} is absent or ambiguous")
	return np.unique(np.append(selected, int(matches[0])))


	def choose_block(sensitivity, baseline, target_actor, trace_scores):
	ranked = []
	for name, entries in sensitivity.items():
	if not block_is_admissible(name, entries, baseline, target_actor):
	continue
	stable = [item for item in entries if item["stable"]]
	if not stable:
	continue
	combined_gain = max(baseline["loss"] - item["loss"] for item in stable)
	target_gain = max(
	baseline["actor_losses"][target_actor] - item["actor_losses"][target_actor]
	for item in stable
	)
	coverage = sum(
	trace_scores[ACTORS.index(actor)] for actor in ACTOR_COVERAGE[name]
	)
	ranked.append((name, combined_gain, target_gain, float(coverage)))
	if not ranked:
	raise RuntimeError(
	"No stable admissible parameter block affects the target actor"
	)
	improving = [item for item in ranked if item[1] > 1e-9]
	if improving:
	selected = max(improving, key=lambda item: (item[1], item[2], item[3]))
	reason = "combined_objective_improvement"
	else:
	improving = [item for item in ranked if item[2] > 1e-9]
	if improving:
	selected = max(improving, key=lambda item: (item[2], item[1], item[3]))
	reason = "target_actor_mask_improvement"
	else:
	selected = max(ranked, key=lambda item: (item[3], item[1], item[2]))
	reason = "actor_coverage_after_inconclusive_perturbations"
	return selected[0], reason


	def main():
	parser = argparse.ArgumentParser()
	parser.add_argument("folder", type=Path)
	parser.add_argument("--asset", type=Path, required=True)
	parser.add_argument("--physics", type=Path)
	parser.add_argument("--baseline", type=Path)
	parser.add_argument("--reference", type=Path)
	parser.add_argument("--output", type=Path)
	parser.add_argument("--population", type=int, default=8)
	parser.add_argument("--iterations", type=int, default=4)
	parser.add_argument("--stride", type=int)
	parser.add_argument("--seed", type=int, default=42)
	parser.add_argument("--perturbation", type=float, default=0.5)
	parser.add_argument("--std-scale", type=float, default=0.5)
	parser.add_argument("--pixel", type=int, nargs=2, metavar=("X", "Y"))
	parser.add_argument("--frame", type=int)
	parser.add_argument("--trace", type=Path)
	parser.add_argument("--query-only", action="store_true")
	parser.add_argument("--trace-only", action="store_true")
	args = parser.parse_args()
	if (args.pixel is None) != (args.frame is None):
	parser.error("pixel and frame must be supplied together")
	if (
	args.population < 2
	or args.iterations < 1
	or (args.stride is not None and args.stride < 1)
	):
	parser.error("population must be >= 2; iterations and stride must be positive")
	if (
	not np.isfinite([args.perturbation, args.std_scale]).all()
	or min(args.perturbation, args.std_scale) <= 0
	):
	parser.error("perturbation and std-scale must be finite and positive")
	folder = args.folder.resolve()
	physics = (args.physics or folder / "physics").resolve()
	output = (args.output or physics / "refinement").resolve()
	if args.query_only:
	if args.pixel is None:
	parser.error("query-only requires pixel and frame")
	path = args.trace or output / "trace.npz"
	if not path.is_file():
	parser.error(f"trace does not exist: {path}")
	try:
	query = MaskErrorTrace.load(path).query(args.frame, *args.pixel)
	except ValueError as error:
	parser.error(str(error))
	print(json.dumps(query, indent=2), flush=True)
	return
	baseline_path = (args.baseline or physics / "best_parameters.json").resolve()
	if not baseline_path.exists():
	parser.error(f"baseline parameters are not ready: {baseline_path}")
	if output in (folder, physics, baseline_path, baseline_path.parent):
	parser.error("output must be a separate refinement folder")
	if not args.asset.is_file():
	parser.error(f"asset does not exist: {args.asset}")
	baseline_parameters = load_parameters(baseline_path)
	config_path = physics / "search_config.json"
	config = json.loads(config_path.read_text()) if config_path.is_file() else {}
	with np.load(folder / "real_masks.npz", allow_pickle=False) as saved:
	if "model" not in saved or "sam3.1" not in str(saved["model"].item()).lower():
	parser.error("real_masks.npz must identify SAM3.1 masks")
	mask_model = str(saved["model"].item())
	order = [saved["names"].tolist().index(name) for name in ACTORS]
	masks = saved["masks"][order]
	confidence = saved["confidence"][order].astype(float)
	observed = saved["observed"][order].astype(bool)
	source_indices = saved["frame_indices"].copy()
	mask_times = saved["timestamps"].copy()
	from real2sim.sim.physics import PhysicsScene

	scene = PhysicsScene(folder, args.asset.resolve(), reference=args.reference)
	try:
	times = np.asarray(scene.timestamps)
	if len(times) != masks.shape[1] or not np.allclose(
	times, mask_times, atol=1e-6, rtol=0
	):
	raise ValueError("Mask and physics timestamps must align")
	mean = np.asarray(
	config.get("parameter_mean", scene.parameter_mean), dtype=float
	)
	std = np.asarray(config.get("parameter_std", scene.parameter_std), dtype=float)
	bounds = np.asarray(scene.parameter_bounds, dtype=float)
	if (
	mean.shape != (25,)
	or std.shape != (25,)
	or not np.isfinite(std).all()
	or np.any(std <= 0)
	):
	raise ValueError(
	"Baseline must provide a valid 25-parameter regularization reference"
	)
	if np.any(baseline_parameters < bounds[:, 0]) or np.any(
	baseline_parameters > bounds[:, 1]
	):
	raise ValueError("Baseline parameters exceed current physics bounds")
	size = tuple(config.get("render_size", (256, 144)))
	if len(size) != 2 or min(size) < 1:
	raise ValueError(
	"Baseline render size must contain positive width and height"
	)
	if args.stride is None and "sampled_frame_indices" in config:
	selected = np.asarray(config["sampled_frame_indices"], dtype=int)
	else:
	selected = np.unique(
	np.append(np.arange(0, len(times), args.stride or 5), len(times) - 1)
	)
	if args.frame is not None:
	try:
	selected = include_query_frame(selected, source_indices, args.frame)
	except ValueError as error:
	parser.error(str(error))
	if (
	selected.ndim != 1
	or not len(selected)
	or np.any(np.diff(selected) <= 0)
	or selected[0] < 0
	or selected[-1] >= len(times)
	):
	raise ValueError("Sampled frame indices must increase within the clip")
	targets = resize_masks(masks, size)
	weights = confidence[:, selected] * observed[:, selected]
	combined_loss = MaskLoss(targets[:, selected], weights)
	actor_losses = [
	MaskLoss(targets[actor : actor + 1, selected], weights[actor : actor + 1])
	if weights[actor].any()
	else None
	for actor in range(len(ACTORS))
	]
	full_loss = MaskLoss(targets, confidence * observed)
	all_indices = np.arange(len(times))
	fps = 1 / float(np.median(np.diff(times)))
	output.mkdir(parents=True, exist_ok=True)
	progress_path = output / "progress.jsonl"
	progress_path.write_text("")
	evaluations = 0
	failures = 0

	def progress(event):
	line = json.dumps(serializable(event), allow_nan=False)
	with progress_path.open("a") as stream:
	stream.write(line + "\n")
	print(line, flush=True)

	def evaluate(parameters):
	nonlocal evaluations, failures
	evaluations += 1
	penalty = regularization(parameters, mean, std)
	try:
	result = scene.rollout(parameters, selected, render_size=size)
	validate_rollout(result, len(selected), size)
	if contact_diagnostics(result).get("unstable", False):
	raise FloatingPointError(
	"PhysicsScene reported an unstable rollout"
	)
	loss = combined_loss(result["masks"])
	components = [
	function(result["masks"][actor : actor + 1]) if function else None
	for actor, function in enumerate(actor_losses)
	]
	if not np.isfinite(loss):
	raise FloatingPointError("Nonfinite refinement objective")
	return {
	"parameters": parameters.copy(),
	"loss": float(loss + penalty),
	"mask_loss": float(loss),
	"regularization": penalty,
	"actor_losses": components,
	"stable": True,
	}
	except SIMULATION_ERRORS as error:
	failures += 1
	return {
	"parameters": parameters.copy(),
	"loss": INVALID_LOSS + penalty,
	"regularization": penalty,
	"actor_losses": [None] * len(ACTORS),
	"stable": False,
	"error": str(error),
	}

	baseline_full = scene.rollout(
	baseline_parameters, all_indices, render_size=size, rgb=True
	)
	validate_rollout(baseline_full, len(times), size)
	baseline_trace = make_trace(
	scene, baseline_full, targets, confidence, observed, source_indices
	)
	real_frames = load_real_frames(folder, len(times), size)
	save_trace(output, baseline_trace, real_frames, fps)
	query = None
	if args.pixel is not None:
	try:
	query = baseline_trace.query(args.frame, *args.pixel)
	except ValueError as error:
	parser.error(str(error))
	write_json(output / "pixel_query.json", query)
	progress({"event": "pixel_query", **query})
	if args.trace_only:
	write_json(
	output / "metrics.json",
	{
	"stage": "physics_pixel_error_trace",
	"trace": baseline_trace.summary(),
	"pixel_query": query,
	},
	)
	progress({"event": "trace_completed", "output": output})
	return
	baseline = evaluate(baseline_parameters)
	if not baseline["stable"]:
	raise RuntimeError("Baseline physics rollout is unstable")
	trace_scores, _ = baseline_trace.error_scores()
	eligible = np.asarray([function is not None for function in actor_losses])
	target_actor = int(np.argmax(np.where(eligible, trace_scores, -1)))
	target_selection = "video_worst_observed_actor"
	if query is not None:
	affected = np.asarray(
	[
	eligible[actor]
	and query["observed"][name]
	and any(query["actor_errors"][name].values())
	for actor, name in enumerate(ACTORS)
	]
	)
	if affected.any():
	target_actor = int(np.argmax(np.where(affected, trace_scores, -1)))
	target_selection = "pixel_affected_actor_with_largest_video_error"
	progress(
	{
	"event": "baseline",
	"loss": baseline["loss"],
	"target_actor": ACTORS[target_actor],
	"target_selection": target_selection,
	"actor_losses": baseline["actor_losses"],
	}
	)
	sensitivity = {}
	candidates = [{"source": "baseline", **baseline}]
	random = np.random.default_rng(args.seed)
	for name, indices in BLOCKS.items():
	direction = random.normal(size=len(indices))
	direction *= np.sqrt(len(indices)) / max(
	float(np.linalg.norm(direction)), 1e-12
	)
	delta = args.perturbation * std[indices] * direction
	sensitivity[name] = []
	for sign in (-1, 1):
	candidate = baseline_parameters.copy()
	candidate[indices] = np.clip(
	candidate[indices] + sign * delta,
	bounds[indices, 0],
	bounds[indices, 1],
	)
	measured = {"block": name, "sign": sign, **evaluate(candidate)}
	measured["combined_objective_change"] = (
	measured["loss"] - baseline["loss"]
	)
	measured["target_actor_loss_change"] = (
	measured["actor_losses"][target_actor]
	- baseline["actor_losses"][target_actor]
	if measured["stable"]
	else None
	)
	sensitivity[name].append(measured)
	candidates.append({"source": f"sensitivity_{name}_{sign}", **measured})
	progress(
	{
	"event": "sensitivity",
	"block": name,
	"sign": sign,
	"loss": measured["loss"],
	"actor_losses": measured["actor_losses"],
	"stable": measured["stable"],
	}
	)
	admissible_blocks = [
	name
	for name, entries in sensitivity.items()
	if block_is_admissible(name, entries, baseline, target_actor)
	]
	block, reason = choose_block(sensitivity, baseline, target_actor, trace_scores)
	indices = BLOCKS[block]
	selected_candidates = [
	baseline,
	*[
	item
	for item in sensitivity[block]
	if candidate_is_acceptable(item, baseline, target_actor)
	],
	]
	center = min(selected_candidates, key=lambda item: item["loss"])[
	"parameters"
	].copy()
	progress(
	{
	"event": "block_selected",
	"block": block,
	"reason": reason,
	"target_actor": ACTORS[target_actor],
	}
	)
	write_json(
	output / "sensitivity.json",
	{
	"target_actor": ACTORS[target_actor],
	"selected_block": block,
	"admissible_blocks": admissible_blocks,
	"target_selection": target_selection,
	"selection_reason": reason,
	"blocks": sensitivity,
	"parameter_indices": BLOCKS,
	"actor_coverage": ACTOR_COVERAGE,
	"physical_cause_uniquely_identified": False,
	},
	)

	def objective(values):
	parameters = baseline_parameters.copy()
	parameters[indices] = values
	return guarded_objective(evaluate(parameters), baseline, target_actor)

	result = search(
	objective,
	center[indices],
	std[indices] * args.std_scale,
	bounds[indices],
	population=args.population,
	iterations=args.iterations,
	seed=args.seed,
	method="cem",
	callback=lambda iteration, loss: progress(
	{
	"event": "block_cem",
	"block": block,
	"iteration": iteration,
	"loss": loss,
	}
	),
	)
	parameters = baseline_parameters.copy()
	parameters[indices] = result.parameters
	candidates.append({"source": "block_cem", **evaluate(parameters)})
	best = min(
	(
	item
	for item in candidates
	if candidate_is_acceptable(item, baseline, target_actor)
	and (
	item["source"] in ("baseline", "block_cem")
	or item.get("block") in admissible_blocks
	)
	),
	key=lambda item: item["loss"],
	)
	final = scene.rollout(
	best["parameters"], all_indices, render_size=size, rgb=True
	)
	validate_rollout(final, len(times), size)
	scene.save(output)
	arrays = flattened_arrays(final)
	arrays.update(
	parameters=best["parameters"],
	names=np.asarray(ACTORS),
	timestamps=times,
	frame_indices=source_indices,
	)
	np.savez_compressed(output / "rollout.npz", **arrays)
	np.savez_compressed(
	output / "best_parameters.npz",
	parameters=best["parameters"],
	parameter_mean=mean,
	parameter_std=std,
	parameter_bounds=bounds,
	timestamps=times,
	frame_indices=source_indices,
	sampled_frame_indices=selected,
	selected_block=block,
	seed=args.seed,
	)
	write_json(
	output / "best_parameters.json",
	{
	"parameters": best["parameters"],
	"loss": best["loss"],
	"source": best["source"],
	},
	)
	final_trace = make_trace(
	scene, final, targets, confidence, observed, source_indices
	)
	save_trace(output, final_trace, real_frames, fps, prefix="refined_")
	video.write(output / "replay.mp4", final["rgb"], fps)
	comparisons = [
	np.concatenate(
	(
	labeled_frame(real, "Real"),
	labeled_frame(before, "Baseline physics"),
	labeled_frame(after, "Refined physics"),
	),
	axis=1,
	)
	for real, before, after in zip(
	real_frames, baseline_full["rgb"], final["rgb"]
	)
	]
	video.write(output / "comparison.mp4", comparisons, fps)
	cv2.imwrite(
	str(output / "preview.png"),
	cv2.cvtColor(comparisons[len(comparisons) // 2], cv2.COLOR_RGB2BGR),
	)
	plot_loss(
	output / "loss.png",
	{f"{block}_cem": result},
	baseline["loss"],
	args.population,
	)
	measurements = lift_metrics(final)
	metrics = {
	"stage": "pixel_trace_physics_block_refinement",
	"simulation_mode": "dynamic_object_contact_physics",
	"mask_model": mask_model,
	"mask_actor_order": ACTORS,
	"baseline_parameters": baseline_parameters,
	"baseline_file": baseline_path,
	"asset": args.asset.resolve(),
	"reference": args.reference,
	"parameter_reference_mean": mean,
	"parameter_reference_std": std,
	"sampled_frame_indices": selected,
	"queried_frame_in_objective": args.frame is not None,
	"render_size": size,
	"target_actor": ACTORS[target_actor],
	"target_selection": target_selection,
	"selected_block": block,
	"admissible_blocks": admissible_blocks,
	"selection_reason": reason,
	"acceptance_constraints": {
	"target_actor_maximum_loss_increase": 1e-9,
	"combined_objective_must_not_increase": True,
	"fallback": "baseline",
	},
	"cem_objective": "combined_objective_with_target_actor_worsening_penalty",
	"physical_cause_uniquely_identified": False,
	"mass_kg": 0.1,
	"friction": float(best["parameters"][24]),
	"baseline_loss": baseline["loss"],
	"best_loss": best["loss"],
	"objective_improvement": baseline["loss"] - best["loss"],
	"best_source": best["source"],
	"baseline_actor_losses": baseline["actor_losses"],
	"best_actor_losses": best["actor_losses"],
	"target_actor_loss_change": best["actor_losses"][target_actor]
	- baseline["actor_losses"][target_actor],
	"baseline_full_mask_loss": full_loss(baseline_full["masks"]),
	"refined_full_mask_loss": full_loss(final["masks"]),
	"population": args.population,
	"iterations": args.iterations,
	"seed": args.seed,
	"sensitivity_evaluations": 2 * len(BLOCKS),
	"cem_evaluations": result.evaluations,
	"total_objective_evaluations": evaluations,
	"full_export_rollouts": 2,
	"unstable_candidates": failures,
	"cem_history": result.history,
	"pixel_query": query,
	"baseline_trace": baseline_trace.summary(),
	"refined_trace": final_trace.summary(),
	"baseline_contact_and_lift": lift_metrics(baseline_full),
	"refined_contact_and_lift": measurements,
	}
	write_json(output / "metrics.json", metrics)
	progress(
	{
	"event": "completed",
	"output": output,
	"baseline_loss": baseline["loss"],
	"best_loss": best["loss"],
	"best_source": best["source"],
	"selected_block": block,
	"maximum_lift_m": measurements["maximum_lift_m"],
	}
	)
	finally:
	scene.close()


	if __name__ == "__main__":
	main()