Buckets:
| import argparse | |
| import json | |
| from pathlib import Path | |
| import cv2 | |
| import mujoco | |
| import numpy as np | |
| from real2sim import video | |
| from real2sim.mask_loss import MaskLoss | |
| from real2sim.optimization import search | |
| ACTORS = ("left", "right", "object") | |
| INVALID_LOSS = 1000.0 | |
| SIMULATION_ERRORS = ( | |
| RuntimeError, | |
| FloatingPointError, | |
| np.linalg.LinAlgError, | |
| mujoco.FatalError, | |
| ) | |
| def serializable(value): | |
| if isinstance(value, dict): | |
| return {str(key): serializable(item) for key, item in value.items()} | |
| if isinstance(value, (list, tuple)): | |
| return [serializable(item) for item in value] | |
| if isinstance(value, np.ndarray): | |
| return serializable(value.tolist()) | |
| if isinstance(value, np.generic): | |
| return serializable(value.item()) | |
| if isinstance(value, float) and not np.isfinite(value): | |
| return None | |
| if isinstance(value, Path): | |
| return str(value) | |
| return value | |
| def write_json(path, value): | |
| path.write_text(json.dumps(serializable(value), indent=2, allow_nan=False) + "\n") | |
| def resize_masks(masks, size): | |
| return np.asarray( | |
| [ | |
| [ | |
| cv2.resize(mask.astype(np.uint8), size, interpolation=cv2.INTER_NEAREST) | |
| for mask in track | |
| ] | |
| for track in masks | |
| ], | |
| dtype=bool, | |
| ) | |
| def regularization(parameters, mean, std): | |
| normalized = (parameters[:24] - mean[:24]) / std[:24] | |
| keys = ((parameters[6:24] - mean[6:24]) / std[6:24]).reshape(2, 3, 3) | |
| reference = 0.002 * np.mean(normalized**2) | |
| smoothness = 0.001 * np.mean(np.diff(keys, axis=1) ** 2) | |
| return float(reference + smoothness) | |
| def contact_diagnostics(result): | |
| for key in ("contact_diagnostics", "diagnostics", "contacts"): | |
| if key in result: | |
| return result[key] | |
| return {} | |
| def validate_rollout(result, frame_count, size): | |
| expected = (len(ACTORS), frame_count, size[1], size[0]) | |
| if np.asarray(result["masks"]).shape != expected: | |
| raise ValueError(f"PhysicsScene masks must have shape {expected}") | |
| for name in ("object_qpos", "object_positions"): | |
| values = np.asarray(result[name]) | |
| if values.shape[0] != frame_count or not np.isfinite(values).all(): | |
| raise FloatingPointError(f"Invalid {name} in physics rollout") | |
| if np.asarray(result["object_positions"]).shape != (frame_count, 3): | |
| raise ValueError("PhysicsScene object_positions must have shape (frames, 3)") | |
| diagnostics = contact_diagnostics(result) | |
| stable = diagnostics.get("stable", True) if isinstance(diagnostics, dict) else True | |
| if not bool(np.asarray(stable).all()): | |
| raise FloatingPointError("PhysicsScene reported an unstable rollout") | |
| def flattened_arrays(value, prefix=""): | |
| arrays = {} | |
| for key, item in value.items(): | |
| name = f"{prefix}{key}" | |
| if isinstance(item, dict): | |
| arrays.update(flattened_arrays(item, f"{name}__")) | |
| elif item is not None: | |
| array = np.asarray(item) | |
| if array.dtype.kind in "biufcSU" and name != "rgb": | |
| arrays[name] = array | |
| return arrays | |
| def lift_metrics(result): | |
| positions = np.asarray(result["object_positions"]) | |
| heights = positions[:, 2] - positions[0, 2] | |
| diagnostics = contact_diagnostics(result) | |
| hand_contact = None | |
| if isinstance(diagnostics, dict): | |
| evidence = [] | |
| for key, value in diagnostics.items(): | |
| name = key.lower() | |
| if ( | |
| "hand" in name | |
| and "object" in name | |
| and ("contact" in name or "force" in name) | |
| ): | |
| values = np.asarray(value) | |
| if values.dtype.kind in "biuf": | |
| evidence.append(bool(np.any(values > 0))) | |
| if evidence: | |
| hand_contact = any(evidence) | |
| lifted = bool(np.max(heights) >= 0.02) | |
| return { | |
| "stable": True, | |
| "maximum_lift_m": float(np.max(heights)), | |
| "final_height_change_m": float(heights[-1]), | |
| "lift_threshold_m": 0.02, | |
| "lifted": lifted, | |
| "hand_object_contact_detected": hand_contact, | |
| "lift_and_hand_contact_observed": lifted and hand_contact | |
| if hand_contact is not None | |
| else None, | |
| "contact_diagnostics": diagnostics, | |
| } | |
| def overlay_masks(frames, masks): | |
| palette = np.asarray(((80, 170, 255), (255, 100, 100), (70, 230, 120))) | |
| output = np.asarray(frames, dtype=np.uint8).copy() | |
| for actor, color in enumerate(palette): | |
| selected = masks[actor] | |
| output[selected] = ( | |
| 0.55 * output[selected].astype(float) + 0.45 * color | |
| ).astype(np.uint8) | |
| return output | |
| def labeled_frame(frame, label): | |
| output = frame.copy() | |
| cv2.rectangle(output, (0, 0), (output.shape[1], 25), (15, 15, 15), -1) | |
| cv2.putText( | |
| output, label, (8, 18), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1 | |
| ) | |
| return output | |
| def load_real_frames(folder, count, size): | |
| path = folder / "mujoco_replay.mp4" | |
| capture = cv2.VideoCapture(str(path)) | |
| if not capture.isOpened(): | |
| raise FileNotFoundError(path) | |
| frames = [] | |
| try: | |
| for _ in range(count): | |
| ok, frame = capture.read() | |
| if not ok: | |
| raise ValueError("Replay video has fewer frames than PhysicsScene") | |
| panel = frame[:, : frame.shape[1] // 3] | |
| frames.append(cv2.resize(cv2.cvtColor(panel, cv2.COLOR_BGR2RGB), size)) | |
| finally: | |
| capture.release() | |
| return np.asarray(frames) | |
| def plot_loss(path, results, initial, population): | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| figure, axis = plt.subplots(figsize=(6, 4)) | |
| axis.axhline(initial, color="gray", linestyle="--", label="Initial") | |
| for method, result in results.items(): | |
| axis.plot( | |
| 1 + np.arange(len(result.history)) * population, | |
| result.history, | |
| marker="o", | |
| label=method.upper(), | |
| ) | |
| axis.set(xlabel="Objective evaluations per method", ylabel="Best objective") | |
| axis.legend() | |
| axis.grid(alpha=0.2) | |
| figure.tight_layout() | |
| figure.savefig(path, dpi=160) | |
| plt.close(figure) | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("folder", type=Path) | |
| parser.add_argument("--asset", type=Path, required=True) | |
| parser.add_argument("--output", type=Path) | |
| parser.add_argument("--population", type=int, default=16) | |
| parser.add_argument("--iterations", type=int, default=6) | |
| parser.add_argument("--stride", type=int, default=5) | |
| parser.add_argument("--seed", type=int, default=42) | |
| args = parser.parse_args() | |
| if args.stride < 1 or args.population < 2 or args.iterations < 1: | |
| parser.error("stride and iterations must be positive; population must be >= 2") | |
| folder = args.folder.resolve() | |
| asset = args.asset.resolve() | |
| output = (args.output or folder / "physics").resolve() | |
| if output == folder: | |
| parser.error("output must differ from the input folder") | |
| if not asset.is_file(): | |
| parser.error(f"asset does not exist: {asset}") | |
| with np.load(folder / "real_masks.npz", allow_pickle=False) as saved: | |
| if "model" not in saved: | |
| parser.error("real_masks.npz must identify its SAM3.1 model") | |
| model_name = str(saved["model"].item()) | |
| if "sam3.1" not in model_name.lower(): | |
| parser.error("real_masks.npz must contain SAM3.1 masks") | |
| names = saved["names"].tolist() | |
| order = [names.index(name) for name in ACTORS] | |
| masks = saved["masks"][order] | |
| confidence = np.asarray(saved["confidence"][order], dtype=float) | |
| observed = np.asarray(saved["observed"][order], dtype=bool) | |
| mask_timestamps = np.asarray(saved["timestamps"], dtype=float) | |
| source_indices = saved["frame_indices"].copy() | |
| from real2sim.sim.physics import PhysicsScene | |
| scene = PhysicsScene(folder, asset) | |
| try: | |
| timestamps = np.asarray(scene.timestamps, dtype=float) | |
| if ( | |
| len(timestamps) < 2 | |
| or not np.isfinite(timestamps).all() | |
| or np.any(np.diff(timestamps) <= 0) | |
| ): | |
| raise ValueError("PhysicsScene timestamps must be finite and increasing") | |
| if masks.shape[1] != len(timestamps) or not np.allclose( | |
| mask_timestamps, timestamps, atol=1e-6, rtol=0 | |
| ): | |
| raise ValueError("SAM3.1 masks and PhysicsScene timestamps must align") | |
| mean = np.asarray(scene.parameter_mean, dtype=float) | |
| std = np.asarray(scene.parameter_std, dtype=float) | |
| bounds = np.asarray(scene.parameter_bounds, dtype=float) | |
| if mean.shape != (25,) or std.shape != (25,) or bounds.shape != (25, 2): | |
| raise ValueError("PhysicsScene must expose 25 parameters and their bounds") | |
| if not np.isfinite(std).all() or np.any(std <= 0): | |
| raise ValueError("PhysicsScene parameter_std must be finite and positive") | |
| mean = np.clip(mean, bounds[:, 0], bounds[:, 1]) | |
| selected = np.unique( | |
| np.append(np.arange(0, len(timestamps), args.stride), len(timestamps) - 1) | |
| ) | |
| size = (256, 144) | |
| targets = resize_masks(masks[:, selected], size) | |
| weights = confidence[:, selected] * observed[:, selected] | |
| mask_loss = MaskLoss(targets, weights) | |
| output.mkdir(parents=True, exist_ok=True) | |
| progress_path = output / "progress.jsonl" | |
| progress_path.write_text("") | |
| write_json( | |
| output / "search_config.json", | |
| { | |
| "input_folder": folder, | |
| "asset": asset, | |
| "parameter_mean": mean, | |
| "parameter_std": std, | |
| "parameter_bounds": bounds, | |
| "parameter_order": { | |
| "0:6": "initial_object_pose_correction", | |
| "6:24": "left_right_wrist_xyz_at_three_keyframes", | |
| "24": "friction", | |
| }, | |
| "key_times": np.linspace(timestamps[0], timestamps[-1], 3), | |
| "population": args.population, | |
| "iterations": args.iterations, | |
| "seed": args.seed, | |
| "sampled_frame_indices": selected, | |
| "render_size": size, | |
| }, | |
| ) | |
| failures = {"initial": 0, "random": 0, "cem": 0} | |
| current_method = "initial" | |
| 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 objective(parameters): | |
| penalty = regularization(parameters, mean, std) | |
| try: | |
| result = scene.rollout(parameters, selected, render_size=size) | |
| validate_rollout(result, len(selected), size) | |
| value = mask_loss(result["masks"]) + penalty | |
| if not np.isfinite(value): | |
| raise FloatingPointError("Nonfinite mask objective") | |
| return float(value) | |
| except SIMULATION_ERRORS as error: | |
| failures[current_method] += 1 | |
| if failures[current_method] <= 3: | |
| progress( | |
| { | |
| "method": current_method, | |
| "event": "unstable_candidate", | |
| "error": str(error), | |
| } | |
| ) | |
| return INVALID_LOSS + penalty | |
| initial_loss = objective(mean) | |
| progress({"method": "initial", "loss": initial_loss}) | |
| results = {} | |
| for method in ("random", "cem"): | |
| current_method = method | |
| result = search( | |
| objective, | |
| mean, | |
| std, | |
| bounds, | |
| population=args.population, | |
| iterations=args.iterations, | |
| seed=args.seed, | |
| method=method, | |
| callback=lambda iteration, loss, name=method: progress( | |
| {"method": name, "iteration": iteration, "loss": loss} | |
| ), | |
| ) | |
| results[method] = result | |
| write_json( | |
| output / f"{method}_search.json", | |
| { | |
| "parameters": result.parameters, | |
| "loss": result.loss, | |
| "history": result.history, | |
| "evaluations": result.evaluations, | |
| "unstable_candidates": failures[method], | |
| }, | |
| ) | |
| candidates = {"initial": (mean, initial_loss)} | |
| candidates.update( | |
| {name: (result.parameters, result.loss) for name, result in results.items()} | |
| ) | |
| best_method = min(candidates, key=lambda name: candidates[name][1]) | |
| best_parameters, best_loss = candidates[best_method] | |
| full_indices = np.arange(len(timestamps)) | |
| full_targets = resize_masks(masks, size) | |
| full_loss = MaskLoss(full_targets, confidence * observed) | |
| fps = 1 / float(np.median(np.diff(timestamps))) | |
| real_frames = load_real_frames(folder, len(timestamps), size) | |
| exported = {} | |
| renderings = {} | |
| for name, (parameters, sampled_loss) in candidates.items(): | |
| destination = output / name | |
| destination.mkdir(exist_ok=True) | |
| np.savez_compressed( | |
| destination / "parameters.npz", | |
| parameters=parameters, | |
| parameter_mean=mean, | |
| parameter_std=std, | |
| parameter_bounds=bounds, | |
| frame_indices=source_indices, | |
| rollout_frame_indices=full_indices, | |
| sampled_frame_indices=selected, | |
| timestamps=timestamps, | |
| seed=args.seed, | |
| ) | |
| write_json(destination / "parameters.json", parameters) | |
| try: | |
| result = scene.rollout( | |
| parameters, full_indices, render_size=size, rgb=True | |
| ) | |
| validate_rollout(result, len(full_indices), size) | |
| except SIMULATION_ERRORS as error: | |
| exported[name] = {"stable": False, "error": str(error)} | |
| write_json(destination / "metrics.json", exported[name]) | |
| continue | |
| scene.save(destination) | |
| arrays = flattened_arrays(result) | |
| arrays.update( | |
| parameters=parameters, | |
| timestamps=timestamps, | |
| frame_indices=source_indices, | |
| names=np.asarray(ACTORS), | |
| ) | |
| np.savez_compressed(destination / "rollout.npz", **arrays) | |
| np.savez_compressed( | |
| destination / "masks.npz", | |
| masks=np.asarray(result["masks"], dtype=bool), | |
| names=np.asarray(ACTORS), | |
| frame_indices=source_indices, | |
| timestamps=timestamps, | |
| model="mujoco_renderer", | |
| ) | |
| measured = lift_metrics(result) | |
| measured.update( | |
| sampled_objective=sampled_loss, | |
| full_mask_loss=full_loss(result["masks"]), | |
| regularization=regularization(parameters, mean, std), | |
| ) | |
| exported[name] = measured | |
| write_json(destination / "metrics.json", measured) | |
| rgb = np.asarray(result["rgb"], dtype=np.uint8) | |
| if rgb.shape != (len(timestamps), size[1], size[0], 3): | |
| raise ValueError( | |
| "PhysicsScene rgb must have shape (frames, height, width, 3)" | |
| ) | |
| renderings[name] = rgb | |
| video.write(destination / "replay.mp4", rgb, fps) | |
| video.write( | |
| destination / "mask_overlay.mp4", | |
| overlay_masks(rgb, np.asarray(result["masks"], dtype=bool)), | |
| fps, | |
| ) | |
| progress( | |
| { | |
| "event": "exported", | |
| "method": name, | |
| "sampled_objective": sampled_loss, | |
| "full_mask_loss": measured["full_mask_loss"], | |
| "maximum_lift_m": measured["maximum_lift_m"], | |
| "lifted": measured["lifted"], | |
| "hand_object_contact_detected": measured[ | |
| "hand_object_contact_detected" | |
| ], | |
| } | |
| ) | |
| metrics = { | |
| "stage": "contact_physics_inverse_simulation", | |
| "simulation_mode": "dynamic_object_contact_physics", | |
| "mask_model": model_name, | |
| "mask_actor_order": ACTORS, | |
| "objective": "balanced_actor_iou_boundary_reference_smoothness", | |
| "input_folder": folder, | |
| "asset": asset, | |
| "population": args.population, | |
| "iterations": args.iterations, | |
| "seed": args.seed, | |
| "frames": len(timestamps), | |
| "fps": fps, | |
| "sampled_frame_indices": selected, | |
| "initial_loss": initial_loss, | |
| "best_method": best_method, | |
| "best_loss": best_loss, | |
| "initial_baseline_evaluations": 1, | |
| "evaluations_per_method": args.population * args.iterations + 1, | |
| "export_rollouts": len(candidates), | |
| "methods": { | |
| name: { | |
| "loss": result.loss, | |
| "history": result.history, | |
| "evaluations": result.evaluations, | |
| "unstable_candidates": failures[name], | |
| } | |
| for name, result in results.items() | |
| }, | |
| "rollouts": exported, | |
| } | |
| write_json(output / "metrics.json", metrics) | |
| write_json( | |
| output / "best_parameters.json", | |
| {"method": best_method, "parameters": best_parameters, "loss": best_loss}, | |
| ) | |
| plot_loss(output / "loss.png", results, initial_loss, args.population) | |
| if "initial" in renderings and best_method in renderings: | |
| comparisons = np.asarray( | |
| [ | |
| np.concatenate( | |
| ( | |
| labeled_frame(real, "Real"), | |
| labeled_frame(initial, "Initial physics"), | |
| labeled_frame(best, f"Best: {best_method.upper()}"), | |
| ), | |
| axis=1, | |
| ) | |
| for real, initial, best in zip( | |
| real_frames, renderings["initial"], renderings[best_method] | |
| ) | |
| ] | |
| ) | |
| video.write(output / "comparison.mp4", comparisons, fps) | |
| cv2.imwrite( | |
| str(output / "preview.png"), | |
| cv2.cvtColor(comparisons[len(comparisons) // 2], cv2.COLOR_RGB2BGR), | |
| ) | |
| overlays = np.concatenate( | |
| ( | |
| overlay_masks(real_frames, full_targets), | |
| video.read(output / best_method / "mask_overlay.mp4"), | |
| ), | |
| axis=2, | |
| ) | |
| video.write(output / "mask_comparison.mp4", overlays, fps) | |
| progress( | |
| { | |
| "event": "completed", | |
| "output": output, | |
| "best_method": best_method, | |
| "best_loss": best_loss, | |
| "best_rollout_stable": exported[best_method]["stable"], | |
| "maximum_lift_m": exported[best_method].get("maximum_lift_m"), | |
| } | |
| ) | |
| finally: | |
| scene.close() | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 19.8 kB
- Xet hash:
- 8e9825fcbe88d6fa3d5b813f13bad7a2e663d0b712025d7a7382af2c0e9fa8c4
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.