Flow-CPS hyper-params v.s. Flow-GRPO

[chenkang455]

Hi @happynear, thanks for your great work!

I have a question about replacing flow-cps with flow-sde. When I made the switch, I noticed my model's rollout performance improved, but the training process became very unstable.

Did you just swap flow-cps for flow-sde directly, or did you have to change other hyperparameters to get it to work? I tried increasing num_steps and the noise level, thinking it might help since the final step in flow-cps has no standard deviation, but that approach failed. Thanks a lot!

[happynear]

I do find that the optimization becomes unstable when switching to flow-cps. After some diagnosis, I find the reason is the /std term in logprob. Thus, I removed it as in 814c359 . After that, I find it becomes more stable. My modification is no more than this commit. The hyperparameter settings, eta and beta, are written in the paper.
If you have already removed the std, maybe you can try to turn down the gradnorm? I didn't tune this hyperparameter during paper experiments (only for our company's service model).
Maybe you can post the training curves, including reward, zero_std_ratio, clipfrac_gt_one, approx_kl and maybe also check the ratio between p_theta and p_ref? Let's find out the reason from the curves.

[happynear]

BTW, adding noise in the last step is meaningless. Flow-GRPO ignores the last step training by setting config.train.timestep_fraction = 0.99.

[chenkang455]

I do find that the optimization becomes unstable when switching to flow-cps. After some diagnosis, I find the reason is the /std term in logprob. Thus, I removed it as in 814c359 . After that, I find it becomes more stable. My modification is no more than this commit. The hyperparameter settings, eta and beta, are written in the paper. In my experiments, the original flow-grpo is also not stable as shown in Figure 4 & 5, both have sudden drop of rewards. Removing the std makes optimization faster and more stable for both Flow-CPS and Flow-GRPO. If you have already done so, maybe you can try to turn down the gradnorm? I didn't tune this hyperparameter during paper experiments (only for our company's service model). Maybe you can post the training curves, including reward, zero_std_ratio, clipfrac_gt_one, approx_kl and maybe also check the ratio between p_theta and p_ref? Let's find out the reason from the curves.

Hi, thanks for the response! I saw your modification to the std value. I'm a bit confused about why this affects the results. Since it's a constant term, shouldn't it cancel out when calculating the ratio (ratio = new_log_ratio - old_log_ratio)?

[chenkang455]

I do find that the optimization becomes unstable when switching to flow-cps. After some diagnosis, I find the reason is the /std term in logprob. Thus, I removed it as in 814c359 . After that, I find it becomes more stable. My modification is no more than this commit. The hyperparameter settings, eta and beta, are written in the paper. In my experiments, the original flow-grpo is also not stable as shown in Figure 4 & 5, both have sudden drop of rewards. Removing the std makes optimization faster and more stable for both Flow-CPS and Flow-GRPO. If you have already done so, maybe you can try to turn down the gradnorm? I didn't tune this hyperparameter during paper experiments (only for our company's service model). Maybe you can post the training curves, including reward, zero_std_ratio, clipfrac_gt_one, approx_kl and maybe also check the ratio between p_theta and p_ref? Let's find out the reason from the curves.

Hi, thanks for the response! I saw your modification to the std value. I'm a bit confused about why this affects the results. Since it's a constant term, shouldn't it cancel out when calculating the ratio (ratio = new_log_prob - old_log_prob)?

[happynear]

The ratio can only cancel out the add/minus terms, which is - torch.log(std_dev_t * torch.sqrt(-1*dt)) - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi))), not multiple/divide terms.

[chenkang455]

The ratio can only cancel out the add/minus terms, which is - torch.log(std_dev_t * torch.sqrt(-1*dt)) - torch.log(torch.sqrt(2 * torch.as_tensor(math.pi))), not multiple/divide terms.

Hi, @happynear thanks for that! I will have a try and post the results.

[chenkang455]

Hi @happynear,

Thank you for your previous advice. I'm working on a grasping task using a Flow Matching model and have encountered a significant performance difference between flow-sde and flow-cps.

My observations are:

• The flow-sde method works well, showing consistent improvement in both rollout and evaluation metrics.
• The flow-cps method shows good initial rollout results, but its performance on the final evaluation metrics stagnates and does not improve.

I am currently using num-steps=3 and noise-level=0.5.

Could you offer some insight into why there is such a large performance gap between these two methods, especially when the initial rollout for flow-cps seems promising? I have tried to enlarge the noise to make them comparable, but the results are similar.

Below are the relevant code snippets for my implementation.

Path Generation Logic:

if self.config.noise_method == "flow_sde":
    sigmas = noise_level * torch.sqrt(
        timesteps / (1 - torch.where(timesteps == 1, timesteps[1], timesteps))
    )[:-1]
    sigma_i = sigmas[idx][:,None,None].expand_as(x_t)
    x0_weight = torch.ones_like(t_input) - (t_input - delta)
    x1_weight = t_input - delta - sigma_i**2 * delta / (2 * t_input) 
    x_t_std = torch.sqrt(delta) * sigma_i
elif self.config.noise_method == "flow_cps":
    pi = torch.pi
    cos_term = torch.cos(pi * noise_level / 2).to(device)
    sin_term = torch.sin(pi * noise_level / 2).to(device)
    x0_weight = torch.ones_like(t_input) - (t_input - delta)
    x1_weight = (t_input - delta) * cos_term
    x_t_std = (t_input - delta) * sin_term

Get log-prob code:

def get_logprob_norm(self, sample, mu, sigma):
    # logprob = log p(x|mu,sigma) = -log(sigma) - 0.5 * log(2 * pi) - 0.5 * ((x - mu) / sigma) ** 2
    if self.config.safe_get_logprob:
        log_prob = -torch.pow((sample - mu), 2)
    else:
        mask = sigma == 0
        sigma_safe = torch.where(mask, torch.ones_like(sigma), sigma)
        constant_term = -torch.log(sigma_safe) - 0.5 * torch.log(2 * torch.pi * torch.ones_like(sample))
        exponent_term = -0.5 * torch.pow((sample - mu) / sigma_safe, 2)
        log_prob = constant_term + exponent_term
        log_prob = torch.where(mask, torch.zeros_like(log_prob), log_prob)
    return log_prob

Rollout results:

Eval results:

Train log:

[happynear]

Currently I have no idea about it. In practice, the reward of flow-cps increases very fast and quickly hacks the rewards. Then I have to turn on the KL loss to reduce the reward hacking. Maybe you can try it.

[happynear]

I see that the inital reward of flow-sde and flow-cps cannot match. In evaluation, both methods are ODE and should have the same reward at epoch 0. Maybe you can check the code?

[chenkang455]

I see that the inital reward of flow-sde and flow-cps cannot match. In evaluation, both methods are ODE and should have the same reward at epoch 0. Maybe you can check the code?

This minor difference might be related to the environment initialization, and thanks for your advice i will try it!