跳转至

Reconstruction optimization simplification plan (2026-04-28)

Goal

把 pipeline/reconstruction/optimization 改成和 pipeline/reconstruction/preprocess 类似的代码风格：

  • 主流程平铺，能直接看到输入、初始化、优化、可选后处理、输出。
  • 不写 provider / registry / runner class / stage plan / callback。
  • 不把 fixed 和 free 拆成两套 pipeline。
  • 阶段间对象在内存里传递，debug 文件只用于检查，不作为 handoff。
  • 保留 params["..."]、sequence["camera"]["T_c2w"]、world["support_plane_offset"] 这类自然层级。

这次重构的核心是：只保留一个 optimize_sequence(...)，fixed/free 只表示 root pose 参数 shape 和 requires_grad mask 的差异。

Current problems

当前主要混乱不在 solver loss 公式，而在编排层：

  • engine/runner.py、engine/runtime.py、engine/stages.py 互相传 run_spec、trainer_spec、sequence_ctx、object_ctx、init_ctx、train_ctx、training_plan。
  • SolverTrainer 持有 solver、state、progress、paths、sequence，并负责 save result；这让数据流不如 preprocess 入口清楚。
  • fixed 和 free 被拆成 _run_fixed_global_optimization(...) / _run_free_global_optimization(...) 两条大分支。
  • track_3d_refine、joint_search、joint_smooth 被藏在 fixed branch 里，主流程看不到完整 pipeline。
  • build_training_plan(...) 本质只是把 config 重新包装一遍，增加了“配置搬运工”代码。

第一轮不要重写 solver/losses.py 的数学细节，也不要大改 tracks/refine3d.py。先把顶层流程压平。

Target layout

建议目标结构：

pipeline/reconstruction/optimization/
  entry.py          # argparse + run_optimization_from_config()
  run.py            # load/init/postprocess/save glue in one readable flow
  optimize.py       # build sequence state + unified optimize_sequence()
  data_loader/
    assets.py       # object asset / renderer / contact vertices
    sequence.py     # tracks, frames, masks, contact, human/world sequence inputs
  init/
    initialize.py   # initialize_object_state flow
    search.py       # first-frame scale / pose / joint search and candidate visualization
    support.py      # support-plane scale sweep, support patch, yaw/z correction
  solver/
    fit.py
    losses.py
    constraints.py
    joint_search.py
    state.py
  tracks/
    refine3d.py
    binding.py
    correspondences.py

删除目标：

pipeline/reconstruction/optimization/engine/
SolverTrainer
build_training_plan()
run_training_plan()
optimize_fixed_sequence()
optimize_free_sequence()
run_*_if_enabled()

run_*_if_enabled() 不需要存在。是否启用阶段应该在 run.py 里显式 if cfg...。

Main flow

run.py 里保留一个线性入口：

def run_optimization_from_config(args, cfg) -> None:
    paths = build_paths(args)
    sequence = load_sequence_inputs(paths, cfg)
    object_ctx = load_object_inputs(paths, cfg)

    init = initialize_object_state(paths, cfg, sequence, object_ctx)
    solver = build_solver(paths, cfg, sequence, object_ctx, init)
    progress = make_progress(...)
    run = {
        "paths": paths,
        "cfg": cfg,
        "sequence": sequence,
        "object": object_ctx,
        "init": init,
        "solver": solver,
        "progress": progress,
    }

    state, losses = optimize_sequence(
        paths=paths,
        cfg=cfg,
        sequence=sequence,
        object_ctx=object_ctx,
        solver=solver,
        init_state=init["state"],
        progress=progress,
    )

    if cfg.optimization.track_3d_refine.enabled:
        state = run_track_3d_refine(run, state)

    if cfg.optimization.joint_search.enabled:
        state = run_joint_search(run, state)

    if cfg.optimization.joint_search.smooth.iters > 0:
        state = smooth_joint_sequence(run, state)

    write_result_pt(paths["output_dir"] / "result.pt", run, state, losses=losses)

要求：

  • run.py 可以有少量 # 1. Load inputs. 这类流程注释。
  • 不新增 stage object、training_plan、runner class。
  • 固定路径在使用处清楚写出，例如 paths["output_dir"] / "result.pt"。
  • optional stage 的 config 判断放在调用点，不塞进 callee。

Unified state design

fixed/free 不拆 pipeline，只改变 root pose 的 batch shape。

mode root rotation / translation joint scale
fixed 共享 root，shape (3,) 或 yaw-only (1,) 每帧 (T, J) 共享 scalar
free 每帧 root，shape (T, 3) 或 yaw-only (T, 1) 每帧 (T, J) 共享 scalar

这里 fixed 的实际语义更接近 shared_root：root pose 可以优化，但所有帧共享同一个 root 参数。free 是 per_frame_root。

核心函数：

def build_sequence_state(init_state: ObjectState, *, num_frames: int, pose_mode: str, device) -> ObjectState:
    if pose_mode == "fixed":
        rotation = init_state.rotation.reshape(-1)
        translation = init_state.translation.reshape(3)
    elif pose_mode == "free":
        rotation = init_state.rotation.reshape(1, -1).expand(num_frames, -1).clone()
        translation = init_state.translation.reshape(1, 3).expand(num_frames, 3).clone()
    else:
        raise RuntimeError(f"Unsupported object.rigid_pose_mode={pose_mode!r}")

    hinge = expand_joint_to_frames(init_state.hinge, num_frames)
    slider = expand_joint_to_frames(init_state.slider, num_frames)
    scale = init_state.scale.reshape(())
    return make_object_state({
        "rotation": rotation,
        "translation": translation,
        "hinge": hinge,
        "slider": slider,
        "scale": scale,
        "device": device,
    })

ObjectState.at(frame_idx) 和 ObjectState.select(frame_indices) 已经能表达 shared vs per-frame tensor。不要再为 fixed/free 写两套 optimizer。

Optimizable variables

requires_grad 表示“这个变量是否优化”，不表示 fixed/free。

def make_optimizable_sequence_state(solver, state, *, optimize_root: bool, optimize_joints: bool, optimize_scale: bool):
    internal = solver._to_internal_state(state)

    rotation = clone_tensor(internal.rotation, requires_grad=optimize_root)
    translation = clone_tensor(internal.translation, requires_grad=optimize_root)
    hinge = clone_tensor(internal.hinge, requires_grad=optimize_joints and internal.hinge.numel() > 0)
    slider = clone_tensor(internal.slider, requires_grad=optimize_joints and internal.slider.numel() > 0)
    scale = clone_tensor(internal.scale, requires_grad=optimize_scale)

    params = []
    for value in (rotation, translation, hinge, slider, scale):
        if value.requires_grad:
            params.append(value)

    return ObjectState(rotation, translation, hinge, slider, scale), params

示例：

  • fixed + optimize root：rotation.shape == (3,)，所有 frame loss 更新同一个 root。
  • free + optimize root：rotation.shape == (T, 3)，每帧 root 独立更新。
  • floor yaw only：继续由 solver 的 internal rotation 处理，fixed 变 (1,)，free 变 (T, 1)。

Unified optimize_sequence

统一优化入口是 optimize_sequence(...)；ArticulatedSolver 只保留 FK、坐标变换、render 和 loss helper。

def optimize_sequence(paths, cfg, sequence, object_ctx, solver, init_state, progress):
    frame_data = sequence["frames"]
    num_frames = sequence["meta"]["num_frames"]
    state0 = build_sequence_state(
        init_state,
        num_frames=num_frames,
        pose_mode=cfg.object.rigid_pose_mode,
        device=paths["device"],
    )

    state, params = make_optimizable_sequence_state(
        solver,
        state0,
        optimize_root=True,
        optimize_joints=True,
        optimize_scale=cfg.optimization.solver.optimize_scale,
    )

    valid_frames = [
        frame_idx
        for frame_idx, frame_inputs in enumerate(frame_data)
        if has_frame_supervision(solver, frame_inputs)
    ]
    if not valid_frames:
        raise RuntimeError("Optimization found no valid frames with supervision")

    optimizer = torch.optim.Adam(params, lr=cfg.optimization.solver.lr)
    for iteration in range(cfg.optimization.solver.num_opt_iters):
        optimizer.zero_grad()
        batch_frames = sample_optimization_frames(valid_frames, cfg.optimization.solver.frame_batch_size)

        total = 0.0
        for chunk_frames in chunk_frame_indices(batch_frames, cfg.optimization.solver.render_chunk_size):
            chunk_state = state.select(chunk_frames)
            verts_world = solver._verts_world_batch_from_state(chunk_state)
            verts_cam = solver._transform_world_to_camera_points(verts_world)

            for local_idx, frame_idx in enumerate(chunk_frames):
                prev_state = state.at(frame_idx - 1) if frame_idx > 0 else None
                prev_prev_state = state.at(frame_idx - 2) if frame_idx > 1 else None
                breakdown = compute_frame_loss_breakdown(
                    solver,
                    state=state.at(frame_idx),
                    prev_state=prev_state,
                    prev_prev_state=prev_prev_state,
                    frame_inputs=frame_data[frame_idx],
                    verts_world=verts_world[local_idx],
                    verts_cam=verts_cam[local_idx],
                )
                total = total + breakdown.total

        (total / max(len(batch_frames), 1)).backward()
        torch.nn.utils.clip_grad_norm_(params, max_norm=1.0)
        optimizer.step()

    losses = compute_sequence_losses(solver, state, frame_data)
    return solver._export_state(state).detach(), losses

实现时不要机械照抄上面的伪代码；应该复用现有 mask render chunk、penetration batch、mask center warmup、progress report 逻辑。重点是：fixed/free 不再分叉。

Track correspondence

当前 free branch 里有 2D segment correspondence，fixed branch 禁用。重构后改成优化前给 frame_data 塞监督字段：

if solver.weights.track_2d > 0.0:
    attach_track_correspondences(sequence, solver, init_state, cfg)

attach_track_correspondences(...) 只负责填：

frame_inputs["vertex_indices"]
frame_inputs["target_2d"]
frame_inputs["visibility"]

后续 compute_frame_loss_breakdown(...) 不关心这些监督来自 fixed 还是 free。

第一轮可以保守处理：

  • 先让 track_2d == 0 的固定根主线跑通。
  • 再把原 free branch 的 segment correspondence 迁到 attach_track_correspondences(...)。
  • 如果 fixed mode 下 correspondence 暂时不可靠，显式报错说明原因，不要静默跳过。

Initialization

初始化编排放在 init/initialize.py 的 initialize_object_state(...)，不写 solver loop。

目标主流程：

def initialize_object_state(paths, cfg, sequence, object_ctx):
    init_mask0 = load_mask(sequence["object"]["mask_dir"], 0)
    scale0 = recover_object_scale(...)
    pose0, hinge0, slider0 = search_init_state(...)
    scale1, hinge1, slider1, meta = refine_init_joint_and_scale_from_depth(...)
    support = fit_support_plane_initial_state(...)

    state = make_object_state({
        "rotation": support["rotation"],
        "translation": support["translation"],
        "hinge": hinge1,
        "slider": slider1,
        "scale": scale1,
        "device": paths["device"],
    })
    return {
        "state": state,
        "tensor_asset": support["tensor_asset"],
        "support_vertex_indices": support["support_vertex_indices"],
    }

init/search.py 保留第一帧 scale、pose、joint candidate、candidate scoring、candidate visualization。 init/search.py 不保留未调用的 local sweep/debug wrapper；评分接口用 candidate + score_spec，避免长参数函数。 init/support.py 保留 floor/support-plane 相关初始化约束。

Postprocess

后处理函数名只描述动作，不带 if_enabled：

run_track_3d_refine(...)
run_joint_search(...)
smooth_joint_sequence(...)

调用点负责：

if cfg.optimization.track_3d_refine.enabled:
    state = run_track_3d_refine(...)

这样 run.py 一眼能看出实际 pipeline，不需要读 postprocess 函数才知道某个阶段是否运行。

Output

run.py 里的 write_result_pt(...) 只负责把优化后的 ObjectState 和 sequence context 写成 result.pt：

  • object.pose: (T, 4, 4) world-frame object pose。
  • object.joint_values: (T, J) hinge + slider。
  • object.scale: scalar。
  • world.T_c2w / world.T_w2c: 明确保留方向说明。
  • human.*: 保留当前 result schema。

不要让输出逻辑反过来依赖 optimizer/trainer class。输入应该是普通对象：

write_result_pt(paths["output_dir"] / "result.pt", run, state, losses=losses)

Migration steps

  1. 新增 run.py 和 optimize.py，先复制迁移逻辑，不改 loss 数学。
  2. 实现 build_sequence_state(...)，确认 fixed/free 只差 root tensor shape。
  3. 固定/自由 root pose 都走 optimize_sequence(...)，只通过 state shape 区分。
  4. 删除 solver 里的旧单帧/shared-pose 优化入口。
  5. 把 free branch 的 segment correspondence 迁成 attach_track_correspondences(...)。
  6. 把初始化编排放到 init/initialize.py，run.py 只调用 initialize_object_state(...)。
  7. 把 track_3d_refine、joint_search、joint_smooth 放到 run.py 显式 optional stage。
  8. 用一个 write_result_pt(...) 取代 build_result/save_result/save_optimization_result 三层输出。
  9. 删除 SolverTrainer、training_plan、engine/ 和临时拆出来的顶层 inputs.py / postprocess.py / output.py。
  10. 跑 smoke case，对比 result.pt key、shape、loss mean、可视化输出。

Acceptance criteria

  • pipeline/reconstruction/optimization/entry.py 只做 argparse 和调用核心函数。
  • run.py 主流程不超过一屏到两屏，能直接看到 load/init/opt/postprocess/save。
  • 没有 engine/ 目录。
  • 没有 SolverTrainer。
  • 没有 build_training_plan()。
  • 没有 optimize_fixed_sequence() / optimize_free_sequence() 两个大入口。
  • fixed/free 只在 build_sequence_state(...) 这类小函数里影响 root pose shape。
  • optional stage 的启用判断只在 run.py 调用点出现。
  • 中间 debug 文件不作为阶段 handoff。
  • 报错信息说明缺哪个文件、期望 shape、当前 shape。