From 2cf509795e2fb92344e5d888dd9b95cf007a207f Mon Sep 17 00:00:00 2001 From: Geoffrey19 Date: Thu, 20 Nov 2025 17:17:25 +0800 Subject: [PATCH] fix bugs in flow --- src/lerobot/policies/wall_x/modeling_wall_x.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/lerobot/policies/wall_x/modeling_wall_x.py b/src/lerobot/policies/wall_x/modeling_wall_x.py index 99cc29f08..ebab36e54 100644 --- a/src/lerobot/policies/wall_x/modeling_wall_x.py +++ b/src/lerobot/policies/wall_x/modeling_wall_x.py @@ -119,7 +119,7 @@ class ActionHead(nn.Module): torch.tensor(self.beta_beta, dtype=dtype, device=device) ) sample = beta_dist.sample([batch_size]) - time = (1 - sample) / self.s + time = (1 - sample) * self.s return time def forward(self, action_chunk, dof_mask=None): @@ -420,7 +420,7 @@ class WallXPolicy(PreTrainedPolicy): num_steps = self.config.num_inference_timesteps dt = 1.0 / num_steps - for step_idx in range(num_steps): + for step_idx in range(num_steps + 1): t = torch.tensor(step_idx * dt, device=device, dtype=dtype) timestep = t.unsqueeze(0).repeat(batch_size)