Fixed the bug in flow matching (#68)

* fix a bug in flow matching

* fixed the bugs in flow matching

* fixed format

---------

Co-authored-by: Geoffrey19 <sympathischmann35@gmail.com>

wall_x/model/action_head.py

@@ -307,7 +307,7 @@ class ActionProcessor(nn.Module):
             torch.Tensor: Sampled timesteps of shape [batch_size]
         """
         sample = self.beta_dist.sample([batch_size]).to(device=device, dtype=dtype)
-        time = (1 - sample) / self.s
+        time = (1 - sample) * self.s
         return time
 
     def proprioception_proj(

wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py

@@ -2046,7 +2046,7 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
             times = torch.linspace(
                 0,
                 1,
-                num_inference_timesteps,
+                num_inference_timesteps + 1,
                 device=inputs_embeds.device,
                 dtype=inputs_embeds.dtype,
             )