381 lines
16 KiB
Python
381 lines
16 KiB
Python
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
|
|
DEFAULT_SCHEMES = {
|
|
"dense": list(range(50)),
|
|
"sparse_10": list(range(0, 50, 5)),
|
|
"sparse_5": list(range(0, 50, 10)),
|
|
"tail_heavy_10": [0, 5, 10, 20, 30, 38, 43, 46, 48, 49],
|
|
"tail_only_6": [40, 43, 46, 47, 48, 49],
|
|
"sparse_8": [0, 7, 14, 21, 28, 35, 42, 49],
|
|
"sparse_4": [0, 16, 32, 49],
|
|
"tail_only_4": [40, 43, 46, 49],
|
|
"tail_heavy_6": [0, 32, 40, 44, 47, 49],
|
|
"tail_heavy_8": [0, 10, 20, 30, 38, 43, 46, 49],
|
|
}
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Offline sparse head execution simulation from existing DDIM logs."
|
|
)
|
|
parser.add_argument(
|
|
"--root_dir",
|
|
type=str,
|
|
default=".",
|
|
help="Repository root or directory under which case outputs are stored.",
|
|
)
|
|
parser.add_argument(
|
|
"--output_dir",
|
|
type=str,
|
|
default=None,
|
|
help="Directory to write sparse-head simulation outputs. Defaults to <root_dir>/sparse_head_simulation.",
|
|
)
|
|
parser.add_argument(
|
|
"--target_threshold",
|
|
type=float,
|
|
default=0.95,
|
|
help="Target cosine threshold used by experiment 1.2.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def discover_stepwise_logs(root_dir: Path) -> list[Path]:
|
|
return sorted(root_dir.glob("unitree_*/case*/output/inference/stepwise_log.csv"))
|
|
|
|
|
|
def load_stepwise_table(stepwise_paths: list[Path]) -> pd.DataFrame:
|
|
frames = []
|
|
for path in stepwise_paths:
|
|
frame = pd.read_csv(path)
|
|
frame["dataset"] = path.parts[-5]
|
|
frame["case"] = path.parts[-4]
|
|
frames.append(frame)
|
|
if not frames:
|
|
raise FileNotFoundError("No stepwise_log.csv files found.")
|
|
stepwise_df = pd.concat(frames, ignore_index=True)
|
|
for column in [
|
|
"step",
|
|
"step_time_s",
|
|
"latent_delta",
|
|
"action_delta",
|
|
"state_delta",
|
|
"action_cosine_vs_full50",
|
|
"state_cosine_vs_full50",
|
|
"latent_l2_vs_full50",
|
|
]:
|
|
stepwise_df[column] = pd.to_numeric(stepwise_df[column], errors="coerce")
|
|
return stepwise_df
|
|
|
|
|
|
def simulate_schemes(stepwise_df: pd.DataFrame,
|
|
schemes: dict[str, list[int]]) -> pd.DataFrame:
|
|
rows = []
|
|
group_columns = [
|
|
"dataset",
|
|
"case",
|
|
"sample_id",
|
|
"scene",
|
|
"pass_type",
|
|
"round_id",
|
|
]
|
|
grouped = stepwise_df.groupby(group_columns)
|
|
for keys, group in grouped:
|
|
group = group.sort_values("step").reset_index(drop=True)
|
|
action_curve = dict(zip(group["step"] - 1, group["action_cosine_vs_full50"]))
|
|
state_curve = dict(zip(group["step"] - 1, group["state_cosine_vs_full50"]))
|
|
|
|
for scheme_name, checkpoints in schemes.items():
|
|
normalized_checkpoints = sorted(
|
|
checkpoint for checkpoint in checkpoints if 0 <= checkpoint <= 49)
|
|
if not normalized_checkpoints:
|
|
continue
|
|
last_checkpoint = normalized_checkpoints[-1]
|
|
rows.append({
|
|
"dataset": keys[0],
|
|
"case": keys[1],
|
|
"sample_id": keys[2],
|
|
"scene": keys[3],
|
|
"pass_type": keys[4],
|
|
"round_id": keys[5],
|
|
"scheme": scheme_name,
|
|
"head_exec_steps_zero_based": json.dumps(normalized_checkpoints),
|
|
"head_exec_count": len(normalized_checkpoints),
|
|
"head_compute_saving_ratio": 1.0 - len(normalized_checkpoints) / 50.0,
|
|
"final_checkpoint_zero_based": last_checkpoint,
|
|
"final_checkpoint_one_based": last_checkpoint + 1,
|
|
"final_action_cosine_vs_dense": action_curve[last_checkpoint],
|
|
"final_state_cosine_vs_dense": state_curve[last_checkpoint],
|
|
})
|
|
return pd.DataFrame(rows)
|
|
|
|
|
|
def summarize_scheme_results(simulation_df: pd.DataFrame) -> tuple[pd.DataFrame, pd.DataFrame]:
|
|
metric_columns = [
|
|
"final_action_cosine_vs_dense",
|
|
"final_state_cosine_vs_dense",
|
|
"head_exec_count",
|
|
"head_compute_saving_ratio",
|
|
]
|
|
overall = simulation_df.groupby(["scheme", "pass_type"])[metric_columns].agg(
|
|
["mean", "median", "std"]).reset_index()
|
|
per_case = simulation_df.groupby(
|
|
["dataset", "case", "scheme", "pass_type"])[metric_columns].mean().reset_index()
|
|
return overall, per_case
|
|
|
|
|
|
def compute_min_steps_needed(stepwise_df: pd.DataFrame,
|
|
threshold: float) -> pd.DataFrame:
|
|
rows = []
|
|
group_columns = [
|
|
"dataset",
|
|
"case",
|
|
"sample_id",
|
|
"scene",
|
|
"pass_type",
|
|
"round_id",
|
|
]
|
|
grouped = stepwise_df.groupby(group_columns)
|
|
for keys, group in grouped:
|
|
group = group.sort_values("step").reset_index(drop=True)
|
|
action_hits = group[group["action_cosine_vs_full50"] >= threshold]
|
|
state_hits = group[group["state_cosine_vs_full50"] >= threshold]
|
|
|
|
action_step = np.nan if action_hits.empty else float(action_hits.iloc[0]["step"] - 1)
|
|
state_step = np.nan if state_hits.empty else float(state_hits.iloc[0]["step"] - 1)
|
|
rows.append({
|
|
"dataset": keys[0],
|
|
"case": keys[1],
|
|
"sample_id": keys[2],
|
|
"scene": keys[3],
|
|
"pass_type": keys[4],
|
|
"round_id": keys[5],
|
|
"target_threshold": threshold,
|
|
"action_min_head_steps_needed": 1.0 if pd.notna(action_step) else np.nan,
|
|
"state_min_head_steps_needed": 1.0 if pd.notna(state_step) else np.nan,
|
|
"action_earliest_checkpoint_zero_based": action_step,
|
|
"state_earliest_checkpoint_zero_based": state_step,
|
|
"action_max_head_compute_saving_ratio": 0.98 if pd.notna(action_step) else np.nan,
|
|
"state_max_head_compute_saving_ratio": 0.98 if pd.notna(state_step) else np.nan,
|
|
})
|
|
return pd.DataFrame(rows)
|
|
|
|
|
|
def compare_tail_vs_uniform(simulation_df: pd.DataFrame) -> pd.DataFrame:
|
|
comparison_specs = [
|
|
("sparse_8", "tail_heavy_8", "budget_8_uniform_vs_tail"),
|
|
("sparse_4", "tail_only_4", "budget_4_uniform_vs_tail"),
|
|
("sparse_10", "tail_heavy_10", "budget_10_uniform_vs_tail"),
|
|
]
|
|
base_columns = [
|
|
"dataset",
|
|
"case",
|
|
"sample_id",
|
|
"scene",
|
|
"pass_type",
|
|
"round_id",
|
|
]
|
|
rows = []
|
|
for uniform_scheme, tail_scheme, comparison_name in comparison_specs:
|
|
uniform_df = simulation_df[simulation_df["scheme"] == uniform_scheme].copy()
|
|
tail_df = simulation_df[simulation_df["scheme"] == tail_scheme].copy()
|
|
merged = uniform_df.merge(
|
|
tail_df,
|
|
on=base_columns,
|
|
suffixes=("_uniform", "_tail"),
|
|
how="inner",
|
|
)
|
|
for _, row in merged.iterrows():
|
|
rows.append({
|
|
"comparison": comparison_name,
|
|
"dataset": row["dataset"],
|
|
"case": row["case"],
|
|
"sample_id": row["sample_id"],
|
|
"scene": row["scene"],
|
|
"pass_type": row["pass_type"],
|
|
"round_id": row["round_id"],
|
|
"uniform_scheme": uniform_scheme,
|
|
"tail_scheme": tail_scheme,
|
|
"uniform_head_exec_count": row["head_exec_count_uniform"],
|
|
"tail_head_exec_count": row["head_exec_count_tail"],
|
|
"uniform_action_cosine_vs_dense": row["final_action_cosine_vs_dense_uniform"],
|
|
"tail_action_cosine_vs_dense": row["final_action_cosine_vs_dense_tail"],
|
|
"uniform_state_cosine_vs_dense": row["final_state_cosine_vs_dense_uniform"],
|
|
"tail_state_cosine_vs_dense": row["final_state_cosine_vs_dense_tail"],
|
|
"tail_minus_uniform_action_cosine":
|
|
row["final_action_cosine_vs_dense_tail"] -
|
|
row["final_action_cosine_vs_dense_uniform"],
|
|
"tail_minus_uniform_state_cosine":
|
|
row["final_state_cosine_vs_dense_tail"] -
|
|
row["final_state_cosine_vs_dense_uniform"],
|
|
"tail_better_action": row["final_action_cosine_vs_dense_tail"] >
|
|
row["final_action_cosine_vs_dense_uniform"],
|
|
"tail_better_state": row["final_state_cosine_vs_dense_tail"] >
|
|
row["final_state_cosine_vs_dense_uniform"],
|
|
})
|
|
return pd.DataFrame(rows)
|
|
|
|
|
|
def build_summary_payload(simulation_df: pd.DataFrame,
|
|
min_steps_df: pd.DataFrame,
|
|
tail_compare_df: pd.DataFrame,
|
|
target_threshold: float) -> dict:
|
|
payload: dict[str, dict] = {
|
|
"config": {
|
|
"target_threshold": target_threshold,
|
|
"schemes": DEFAULT_SCHEMES,
|
|
},
|
|
"experiment_1_1": {},
|
|
"experiment_1_2": {},
|
|
"experiment_1_3": {},
|
|
}
|
|
|
|
for scheme, group in simulation_df.groupby("scheme"):
|
|
payload["experiment_1_1"][scheme] = {
|
|
"num_rows": int(len(group)),
|
|
"head_exec_count": float(group["head_exec_count"].iloc[0]),
|
|
"head_compute_saving_ratio": float(
|
|
group["head_compute_saving_ratio"].iloc[0]),
|
|
"final_action_cosine_vs_dense_mean": float(
|
|
group["final_action_cosine_vs_dense"].mean()),
|
|
"final_action_cosine_vs_dense_median": float(
|
|
group["final_action_cosine_vs_dense"].median()),
|
|
"final_state_cosine_vs_dense_mean": float(
|
|
group["final_state_cosine_vs_dense"].mean()),
|
|
"final_state_cosine_vs_dense_median": float(
|
|
group["final_state_cosine_vs_dense"].median()),
|
|
}
|
|
|
|
payload["experiment_1_2"] = {
|
|
"action_earliest_checkpoint_zero_based": {
|
|
"mean": float(
|
|
pd.to_numeric(
|
|
min_steps_df["action_earliest_checkpoint_zero_based"],
|
|
errors="coerce").dropna().mean()),
|
|
"median": float(
|
|
pd.to_numeric(
|
|
min_steps_df["action_earliest_checkpoint_zero_based"],
|
|
errors="coerce").dropna().median()),
|
|
},
|
|
"state_earliest_checkpoint_zero_based": {
|
|
"mean": float(
|
|
pd.to_numeric(
|
|
min_steps_df["state_earliest_checkpoint_zero_based"],
|
|
errors="coerce").dropna().mean()),
|
|
"median": float(
|
|
pd.to_numeric(
|
|
min_steps_df["state_earliest_checkpoint_zero_based"],
|
|
errors="coerce").dropna().median()),
|
|
},
|
|
"min_head_steps_needed_action_unique": sorted(
|
|
pd.to_numeric(min_steps_df["action_min_head_steps_needed"],
|
|
errors="coerce").dropna().unique().tolist()),
|
|
"min_head_steps_needed_state_unique": sorted(
|
|
pd.to_numeric(min_steps_df["state_min_head_steps_needed"],
|
|
errors="coerce").dropna().unique().tolist()),
|
|
}
|
|
|
|
for comparison, group in tail_compare_df.groupby("comparison"):
|
|
payload["experiment_1_3"][comparison] = {
|
|
"num_rows": int(len(group)),
|
|
"tail_better_action_share": float(group["tail_better_action"].mean()),
|
|
"tail_better_state_share": float(group["tail_better_state"].mean()),
|
|
"tail_minus_uniform_action_cosine_mean": float(
|
|
group["tail_minus_uniform_action_cosine"].mean()),
|
|
"tail_minus_uniform_state_cosine_mean": float(
|
|
group["tail_minus_uniform_state_cosine"].mean()),
|
|
}
|
|
return payload
|
|
|
|
|
|
def write_markdown_report(path: Path, payload: dict) -> None:
|
|
lines = [
|
|
"# Sparse Head Execution Simulation",
|
|
"",
|
|
"This report uses zero-order hold over logged stepwise action/state outputs.",
|
|
"For a sparse scheme, the final output at step 49 is approximated by the most recent checkpoint output.",
|
|
"",
|
|
"## Experiment 1.1",
|
|
"",
|
|
]
|
|
|
|
for scheme, stats in payload["experiment_1_1"].items():
|
|
lines.extend([
|
|
f"### {scheme}",
|
|
"",
|
|
f"- Head exec count: {stats['head_exec_count']:.0f}",
|
|
f"- Head compute saving ratio: {stats['head_compute_saving_ratio']:.4f}",
|
|
f"- Final action cosine vs dense: mean={stats['final_action_cosine_vs_dense_mean']:.4f}, median={stats['final_action_cosine_vs_dense_median']:.4f}",
|
|
f"- Final state cosine vs dense: mean={stats['final_state_cosine_vs_dense_mean']:.4f}, median={stats['final_state_cosine_vs_dense_median']:.4f}",
|
|
"",
|
|
])
|
|
|
|
lines.extend([
|
|
"## Experiment 1.2",
|
|
"",
|
|
f"- Target threshold: {payload['config']['target_threshold']}",
|
|
f"- Action earliest checkpoint mean: {payload['experiment_1_2']['action_earliest_checkpoint_zero_based']['mean']:.4f}",
|
|
f"- State earliest checkpoint mean: {payload['experiment_1_2']['state_earliest_checkpoint_zero_based']['mean']:.4f}",
|
|
f"- Unique min head steps needed for action: {payload['experiment_1_2']['min_head_steps_needed_action_unique']}",
|
|
f"- Unique min head steps needed for state: {payload['experiment_1_2']['min_head_steps_needed_state_unique']}",
|
|
"",
|
|
"## Experiment 1.3",
|
|
"",
|
|
])
|
|
|
|
for comparison, stats in payload["experiment_1_3"].items():
|
|
lines.extend([
|
|
f"### {comparison}",
|
|
"",
|
|
f"- Tail better action share: {stats['tail_better_action_share']:.4f}",
|
|
f"- Tail better state share: {stats['tail_better_state_share']:.4f}",
|
|
f"- Mean tail-minus-uniform action cosine: {stats['tail_minus_uniform_action_cosine_mean']:.4f}",
|
|
f"- Mean tail-minus-uniform state cosine: {stats['tail_minus_uniform_state_cosine_mean']:.4f}",
|
|
"",
|
|
])
|
|
|
|
with open(path, "w", encoding="utf-8") as file:
|
|
file.write("\n".join(lines) + "\n")
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
root_dir = Path(args.root_dir).resolve()
|
|
output_dir = Path(args.output_dir).resolve(
|
|
) if args.output_dir else root_dir / "sparse_head_simulation"
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
stepwise_paths = discover_stepwise_logs(root_dir)
|
|
stepwise_df = load_stepwise_table(stepwise_paths)
|
|
|
|
simulation_df = simulate_schemes(stepwise_df, DEFAULT_SCHEMES)
|
|
scheme_overall_df, scheme_per_case_df = summarize_scheme_results(simulation_df)
|
|
min_steps_df = compute_min_steps_needed(stepwise_df, args.target_threshold)
|
|
tail_compare_df = compare_tail_vs_uniform(simulation_df)
|
|
summary_payload = build_summary_payload(simulation_df, min_steps_df,
|
|
tail_compare_df,
|
|
args.target_threshold)
|
|
|
|
simulation_df.to_csv(output_dir / "scheme_simulation_per_round.csv",
|
|
index=False)
|
|
scheme_overall_df.to_csv(output_dir / "scheme_simulation_overall.csv",
|
|
index=False)
|
|
scheme_per_case_df.to_csv(output_dir / "scheme_simulation_per_case.csv",
|
|
index=False)
|
|
min_steps_df.to_csv(output_dir / "min_head_steps_needed.csv", index=False)
|
|
tail_compare_df.to_csv(output_dir / "tail_vs_uniform.csv", index=False)
|
|
with open(output_dir / "summary.json", "w", encoding="utf-8") as file:
|
|
json.dump(summary_payload, file, indent=2, ensure_ascii=False)
|
|
write_markdown_report(output_dir / "report.md", summary_payload)
|
|
|
|
print(f"Sparse head simulation written to: {output_dir}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|