| 110 | |
| 111 | @jax.jit |
| 112 | def update_once(carry, unused): |
| 113 | t, rng, mu_0t, sigma = carry |
| 114 | |
| 115 | # sample from q_i |
| 116 | rng, Y0s_rng = jax.random.split(rng) |
| 117 | eps_u = jax.random.normal(Y0s_rng, (args.Nsample, args.Hsample, Nu)) * sigma |
| 118 | Y0s = eps_u + mu_0t |
| 119 | Y0s = jnp.clip(Y0s, -1.0, 1.0) |
| 120 | |
| 121 | # esitimate mu_0tm1 |
| 122 | rews = jax.vmap(eval_us, in_axes=(None, 0))(state_init, Y0s).mean(axis=-1) |
| 123 | logp0 = (rews - rews.mean()) / rews.std() / args.temp_sample |
| 124 | weights = jax.nn.softmax(logp0) |
| 125 | mu_0tm1, sigma = update_fn(weights, Y0s, sigma, mu_0t) |
| 126 | |
| 127 | return (t - 1, rng, mu_0tm1, sigma), rews.mean() |
| 128 | |
| 129 | # run reverse |
| 130 | def update(mu_0T, rng): |