fix open loop script

This commit is contained in:
vincentchen
2025-09-08 18:37:11 +08:00
parent cfafa3923f
commit 272264b516
+7 -2
View File
@@ -38,8 +38,7 @@ gt_traj = torch.zeros((total_frames, action_dim))
pred_traj = torch.zeros((total_frames, action_dim)) pred_traj = torch.zeros((total_frames, action_dim))
for idx, batch in enumerate(dataloader): for idx, batch in enumerate(dataloader):
gt_traj[idx] = batch['action_chunk'][0, 0,:action_dim] if idx % pred_horizon ==0 and idx + pred_horizon < total_frames:
if idx % 32 ==0 and idx + 32 < total_frames:
batch = batch.to("cuda") batch = batch.to("cuda")
with torch.no_grad(): with torch.no_grad():
outputs = model( outputs = model(
@@ -51,6 +50,12 @@ for idx, batch in enumerate(dataloader):
) )
pred_traj[idx : idx + pred_horizon] = outputs['predict_action'].detach().cpu() pred_traj[idx : idx + pred_horizon] = outputs['predict_action'].detach().cpu()
# Denormalize ground truth actions
gt_action_chunk = batch['action_chunk'][:, :, :action_dim]
dof_mask = batch["dof_mask"].to(gt_action_chunk.dtype)
denormalized_gt = model.action_preprocessor.normalizer_action.unnormalize_data(gt_action_chunk, ["x2_normal"], dof_mask)
gt_traj[idx : idx + pred_horizon] = denormalized_gt.detach().cpu()
gt_traj_np = gt_traj.numpy() gt_traj_np = gt_traj.numpy()
pred_traj_np = pred_traj.numpy() pred_traj_np = pred_traj.numpy()