From 272264b5166407e9ba4db37f4e438c0b58f98e34 Mon Sep 17 00:00:00 2001 From: vincentchen Date: Mon, 8 Sep 2025 18:33:52 +0800 Subject: [PATCH] fix open loop script --- scripts/draw_openloop_plot.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/scripts/draw_openloop_plot.py b/scripts/draw_openloop_plot.py index a649392..2be5694 100644 --- a/scripts/draw_openloop_plot.py +++ b/scripts/draw_openloop_plot.py @@ -38,8 +38,7 @@ gt_traj = torch.zeros((total_frames, action_dim)) pred_traj = torch.zeros((total_frames, action_dim)) for idx, batch in enumerate(dataloader): - gt_traj[idx] = batch['action_chunk'][0, 0,:action_dim] - if idx % 32 ==0 and idx + 32 < total_frames: + if idx % pred_horizon ==0 and idx + pred_horizon < total_frames: batch = batch.to("cuda") with torch.no_grad(): outputs = model( @@ -50,7 +49,13 @@ for idx, batch in enumerate(dataloader): predict_mode="fast" ) 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() pred_traj_np = pred_traj.numpy()