File size: 5,856 Bytes
fbd9366 | 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 | #!/usr/bin/env python3
"""Run one independent T-Rex Track-Force 16-step chunk from an NPZ prefix.
Required NPZ arrays:
head_left, left_wrist, right_wrist: uint8 [9,H,W,3]
state_eef62: float [62]
track_past_xy: float [16,250,2] in [0,1]
track_past_visibility: float/bool [16,250]
tactile_force_history: float [16,10,6] (raw sensor units)
The last force-history sample is used as current force. Output contains the
normalized model action, physical delta-base EEF62 action, absolute EEF62
targets, and predicted future tracks.
"""
from __future__ import annotations
import argparse
from pathlib import Path
import numpy as np
import torch
from transformers import AutoTokenizer
from groot.vla.model.dreamzero.transform.dreamzero_cotrain import (
basic_clean,
whitespace_clean,
)
from groot.vla.model.trex_track_force.runtime import (
TrexRuntimeStatistics,
delta_base_to_absolute,
)
from groot.vla.model.trex_track_force.vla import TrexTrackForceVLA
def _grid_three_views(archive: np.lib.npyio.NpzFile) -> np.ndarray:
views = [
np.asarray(archive[name], dtype=np.uint8)
for name in ("head_left", "left_wrist", "right_wrist")
]
if any(view.ndim != 4 or view.shape[0] != 9 or view.shape[-1] != 3 for view in views):
raise ValueError("each RGB view must be uint8 [9,H,W,3]")
if len({view.shape for view in views}) != 1:
raise ValueError("all three RGB histories must have the same shape")
_, height, width, channels = views[0].shape
grid = np.zeros((9, 2 * height, 2 * width, channels), dtype=np.uint8)
grid[:, :height, :width] = views[0]
grid[:, height:, :width] = views[1]
grid[:, :height, width:] = views[2]
return grid
def _pad64(values: np.ndarray) -> np.ndarray:
return np.pad(values, ((0, 0), (0, 2)), mode="constant")
def run(args: argparse.Namespace) -> None:
device = torch.device(args.device)
dtype = torch.bfloat16 if args.bf16 else torch.float32
stats = TrexRuntimeStatistics.from_dataset(args.dataset_root)
with np.load(args.input_npz, allow_pickle=False) as archive:
history_images = _grid_three_views(archive)
reference_state = np.asarray(archive["state_eef62"], dtype=np.float32)
track_xy = np.asarray(archive["track_past_xy"], dtype=np.float32)
track_visibility = np.asarray(
archive["track_past_visibility"], dtype=np.float32
)
force_history_raw = np.asarray(
archive["tactile_force_history"], dtype=np.float32
)
if reference_state.shape != (62,):
raise ValueError("state_eef62 must be [62]")
if track_xy.shape != (16, 250, 2) or track_visibility.shape != (16, 250):
raise ValueError("past tracks must be [16,250,2] and [16,250]")
if force_history_raw.shape != (16, 10, 6):
raise ValueError("tactile_force_history must be [16,10,6]")
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
instruction = whitespace_clean(basic_clean(args.instruction))
text = tokenizer(
instruction,
max_length=512,
padding="max_length",
truncation=True,
return_tensors="pt",
)
model = TrexTrackForceVLA.load_lora(str(args.checkpoint))
model.eval().requires_grad_(False)
model.to(device=device, dtype=dtype)
normalized_state = _pad64(stats.normalize_state(reference_state)[None])
normalized_force_history = stats.normalize_force(force_history_raw)
inputs = {
"history_images": torch.from_numpy(history_images[None]).to(device),
"state": torch.from_numpy(normalized_state[:, None]).to(device, dtype),
"track_past_xy": torch.from_numpy(track_xy[None, None]).to(device, dtype),
"track_past_visibility": torch.from_numpy(
track_visibility[None, None]
).to(device, dtype),
"current_force": torch.from_numpy(
normalized_force_history[-1:][None]
).to(device, dtype),
"tactile_force_history": torch.from_numpy(
normalized_force_history[None, None]
).to(device, dtype),
"text": text.input_ids.to(device),
"text_attention_mask": text.attention_mask.to(device),
}
with torch.inference_mode():
prediction = model.get_action(inputs)
normalized_action = prediction["action_pred"].float().cpu().numpy()[0]
delta_base_action = stats.denormalize_action(normalized_action)
absolute_action = delta_base_to_absolute(reference_state, delta_base_action)
output = {
"normalized_action64": normalized_action,
"delta_base_action62": delta_base_action,
"absolute_action62": absolute_action,
"track_pred": prediction["track_pred"].float().cpu().numpy()[0],
}
if "video_latents_pred" in prediction:
output["video_latents_pred"] = (
prediction["video_latents_pred"].float().cpu().numpy()[0]
)
args.output_npz.parent.mkdir(parents=True, exist_ok=True)
np.savez_compressed(args.output_npz, **output)
print(f"wrote {args.output_npz} with 16 actions at 20 Hz")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--checkpoint", type=Path, required=True)
parser.add_argument("--input-npz", type=Path, required=True)
parser.add_argument("--output-npz", type=Path, required=True)
parser.add_argument("--dataset-root", type=Path, required=True)
parser.add_argument("--tokenizer-path", type=Path, required=True)
parser.add_argument(
"--instruction",
default="Perform the requested bimanual manipulation.",
)
parser.add_argument("--device", default="cuda:0")
parser.add_argument("--bf16", action=argparse.BooleanOptionalAction, default=True)
return parser.parse_args()
if __name__ == "__main__":
run(parse_args())
|