Function update

mbd/planners/path_integral.py:130–141

(mu_0T, rng)

    # run reverse
    def update(mu_0T, rng):
        sigma = 1.0
        mu_0t = mu_0T
        mu_0ts = []
        with tqdm(range(args.Nrefine - 1, 0, -1), desc="Path Integrating") as pbar:
            for t in pbar:
                carry_once = (t, rng, mu_0t, sigma)
                (t, rng, mu_0t, sigma), rew = update_once(carry_once, None)
                mu_0ts.append(mu_0t)
                # Update the progress bar's suffix to show the current reward
                pbar.set_postfix({"rew": f"{rew:.2e}"})
        return mu_0ts

    rng_exp, rng = jax.random.split(rng)
    mu_0ts = update(mu_0T, rng_exp)