File size: 8,226 Bytes
eafbe80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 | #!/usr/bin/env python3
"""
Build per-chunk action JSONs for a composite action sequence, for static-consistency revisit tests.
First version supports the requested pattern:
rotate_left_45 -> translate_forward -> rotate_right_45 -> translate_backward
Notes:
- Each JSON is a dict: frame_index(str) -> RT list length 12: [tx,ty,tz,R11..R33] (row-major 3x3).
- Actions are *relative to each chunk's first frame* (matching training / existing eval conventions).
"""
from __future__ import annotations
import argparse
import json
import math
import os
import random
from typing import Dict, List, Tuple
def _load_json(path: str) -> Dict[str, List[float]]:
with open(path, "r", encoding="utf-8") as f:
return json.load(f)
def build_rotation_yaw_chunk(yaw_total_deg: float, clockwise: bool, chunk_frames: int) -> Dict[str, List[float]]:
"""Linear yaw 0→±yaw_total_deg in chunk (Z-only RT), same convention as run_replay_loop_two_chunk.build_action_chunk."""
denom = max(1, chunk_frames - 1)
sign = -1.0 if clockwise else 1.0
out: Dict[str, List[float]] = {}
for i in range(chunk_frames):
yaw = sign * (i / denom) * float(yaw_total_deg)
rad = math.radians(yaw)
c, s = math.cos(rad), math.sin(rad)
r_flat = [c, -s, 0.0, s, c, 0.0, 0.0, 0.0, 1.0]
out[str(i)] = [0.0, 0.0, 0.0] + r_flat
return out
def build_translation_only(direction: str, translation_delta: float, chunk_frames: int) -> Dict[str, List[float]]:
"""Mirror run_replay_loop_two_chunk.build_action_translation_only without importing heavy deps."""
identity_rot = [1.0, 0.0, 0.0,
0.0, 1.0, 0.0,
0.0, 0.0, 1.0]
denom = max(1, chunk_frames - 1)
out: Dict[str, List[float]] = {}
for i in range(chunk_frames):
t = (i / denom) * translation_delta
if direction == "forward":
tx, ty, tz = 0.0, t, 0.0
elif direction == "backward":
tx, ty, tz = 0.0, -t, 0.0
elif direction == "left":
tx, ty, tz = -t, 0.0, 0.0
elif direction == "right":
tx, ty, tz = t, 0.0, 0.0
else:
tx, ty, tz = 0.0, 0.0, 0.0
out[str(i)] = [tx, ty, tz] + identity_rot
return out
def save_action_json(actions: Dict[str, List[float]], path: str) -> None:
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "w", encoding="utf-8") as f:
json.dump(actions, f, indent=2)
def build_combo(
exp_dir: str,
out_dir: str,
chunk_frames: int = 81,
translation_delta: float = 0.1,
) -> Tuple[str, str, str, str]:
"""
Create 4 chunk action jsons under out_dir:
chunk0_rotate_left_45.json
chunk1_translate_forward.json
chunk2_rotate_right_45.json
chunk3_translate_backward.json
"""
left_path = os.path.join(exp_dir, "action_rotation_left_45.json")
right_path = os.path.join(exp_dir, "action_rotation_right_45.json")
if not (os.path.isfile(left_path) and os.path.isfile(right_path)):
raise FileNotFoundError(f"Missing rotation jsons under exp_dir: {left_path} / {right_path}")
rot_left = _load_json(left_path)
rot_right = _load_json(right_path)
# Sanity: ensure expected frame keys exist; if not, allow but warn via truncation
def _trim(d: Dict[str, List[float]]) -> Dict[str, List[float]]:
return {str(i): d[str(i)] for i in range(chunk_frames) if str(i) in d}
rot_left = _trim(rot_left)
rot_right = _trim(rot_right)
trans_fwd = build_translation_only("forward", translation_delta, chunk_frames)
trans_bwd = build_translation_only("backward", translation_delta, chunk_frames)
p0 = os.path.join(out_dir, "chunk0_rotate_left_45.json")
p1 = os.path.join(out_dir, "chunk1_translate_forward.json")
p2 = os.path.join(out_dir, "chunk2_rotate_right_45.json")
p3 = os.path.join(out_dir, "chunk3_translate_backward.json")
save_action_json(rot_left, p0)
save_action_json(trans_fwd, p1)
save_action_json(rot_right, p2)
save_action_json(trans_bwd, p3)
return p0, p1, p2, p3
def build_random_symmetric_closed_loop(
out_dir: str,
chunk_frames: int,
rng: random.Random,
yaw_min: float = 20.0,
yaw_max: float = 55.0,
translation_min: float = 0.05,
translation_max: float = 0.18,
) -> Tuple[List[str], Dict]:
"""
Symmetric motion that composes to ~identity in the training RT convention:
chunk0: CCW yaw (left) 0→+Y
chunk1: forward +d along Y
chunk2: CW yaw (right) 0→-Y (cancels chunk0 in world yaw if chunk frames align)
chunk3: backward -d along Y (cancels chunk1 translation)
Same filenames as fixed 45° combo for drop-in use with run_combo_revisit_fixed_first.py.
"""
yaw = rng.uniform(float(yaw_min), float(yaw_max))
d = rng.uniform(float(translation_min), float(translation_max))
rot_left = build_rotation_yaw_chunk(yaw, clockwise=False, chunk_frames=chunk_frames)
rot_right = build_rotation_yaw_chunk(yaw, clockwise=True, chunk_frames=chunk_frames)
trans_fwd = build_translation_only("forward", d, chunk_frames)
trans_bwd = build_translation_only("backward", d, chunk_frames)
p0 = os.path.join(out_dir, "chunk0_rotate_left_45.json")
p1 = os.path.join(out_dir, "chunk1_translate_forward.json")
p2 = os.path.join(out_dir, "chunk2_rotate_right_45.json")
p3 = os.path.join(out_dir, "chunk3_translate_backward.json")
save_action_json(rot_left, p0)
save_action_json(trans_fwd, p1)
save_action_json(rot_right, p2)
save_action_json(trans_bwd, p3)
meta = {
"pattern": "symmetric_closed_loop_random",
"yaw_deg": yaw,
"translation_delta": d,
"chunk_frames": chunk_frames,
"chunks": [
{"file": os.path.basename(p0), "desc": "ccw_yaw_0_to_+yaw"},
{"file": os.path.basename(p1), "desc": "forward_d"},
{"file": os.path.basename(p2), "desc": "cw_yaw_0_to_-yaw"},
{"file": os.path.basename(p3), "desc": "backward_d"},
],
}
return [p0, p1, p2, p3], meta
def main() -> None:
p = argparse.ArgumentParser(description="Build composite action JSONs for revisit tests")
p.add_argument("--exp_dir", type=str, default="", help="exp dir with action_rotation_*.json (fixed 45° mode)")
p.add_argument("--out_dir", type=str, required=True, help="output directory to write per-chunk action jsons")
p.add_argument("--chunk_frames", type=int, default=81)
p.add_argument("--translation_delta", type=float, default=0.1)
p.add_argument(
"--random_symmetric",
action="store_true",
help="Random yaw/translation magnitudes with symmetric closed-loop (left→fwd→right→back); ignores exp_dir rotations",
)
p.add_argument("--combo_seed", type=int, default=42, help="RNG seed for --random_symmetric")
p.add_argument("--yaw_min", type=float, default=20.0)
p.add_argument("--yaw_max", type=float, default=55.0)
p.add_argument("--translation_min", type=float, default=0.05)
p.add_argument("--translation_max", type=float, default=0.18)
args = p.parse_args()
out_dir = os.path.abspath(args.out_dir)
os.makedirs(out_dir, exist_ok=True)
if args.random_symmetric:
rng = random.Random(int(args.combo_seed))
paths, meta = build_random_symmetric_closed_loop(
out_dir,
chunk_frames=args.chunk_frames,
rng=rng,
yaw_min=args.yaw_min,
yaw_max=args.yaw_max,
translation_min=args.translation_min,
translation_max=args.translation_max,
)
with open(os.path.join(out_dir, "combo_manifest.json"), "w", encoding="utf-8") as f:
json.dump(meta, f, indent=2)
print("\n".join(paths))
return
if not args.exp_dir:
raise SystemExit("build_action_combo: need --exp_dir unless --random_symmetric")
paths = build_combo(
exp_dir=os.path.abspath(args.exp_dir),
out_dir=out_dir,
chunk_frames=args.chunk_frames,
translation_delta=args.translation_delta,
)
print("\n".join(paths))
if __name__ == "__main__":
main()
|