Update README.md

README.md

@@ -14,20 +14,27 @@ import torch
 import numpy as np
 
 from alpamayo_r1.models.alpamayo_r1 import AlpamayoR1
+
 from alpamayo_r1.load_physical_aiavdataset import load_physical_aiavdataset
+
 from alpamayo_r1 import helper
 
  
 model_path = "Alpamayo-R1-10B-4bit"
+
 model = AlpamayoR1.from_pretrained(model_path, dtype=torch.bfloat16).to("cuda")
+
 processor = helper.get_processor(model.tokenizer)
 
  
 clip_id = "030c760c-ae38-49aa-9ad8-f5650a545d26"
+
 print(f"Loading dataset for clip_id: {clip_id}...")
+
 data = load_physical_aiavdataset(clip_id, t0_us=15_100_000,num_frames=1)
-print(f"{data}")
+
 print("Dataset loaded.")
+
 messages = helper.create_message(data["image_frames"].flatten(0, 1))
 
 inputs = processor.apply_chat_template(
@@ -48,7 +55,9 @@ model_inputs = {
 model_inputs = helper.to_device(model_inputs, "cuda")
 
 torch.cuda.manual_seed_all(42)
+
 with torch.autocast("cuda", dtype=torch.bfloat16):
+
     pred_xyz, pred_rot, extra = model.sample_trajectories_from_data_with_vlm_rollout(
         data=model_inputs,
         top_p=0.98,
@@ -62,10 +71,15 @@ with torch.autocast("cuda", dtype=torch.bfloat16):
 print("Chain-of-Causation (per trajectory):\n", extra["cot"][0])
 
 gt_xy = data["ego_future_xyz"].cpu()[0, 0, :, :2].T.numpy()
+
 pred_xy = pred_xyz.cpu().numpy()[0, 0, :, :, :2].transpose(0, 2, 1)
+
 diff = np.linalg.norm(pred_xy - gt_xy[None, ...], axis=1).mean(-1)
+
 min_ade = diff.min()
+
 print("minADE:", min_ade, "meters")
+
 print(
     "Note: VLA-reasoning models produce nondeterministic outputs due to trajectory sampling, "
     "hardware differences, etc. With num_traj_samples=1 (set for GPU memory compatibility), "