(step_env, state, us)
| 4 | |
| 5 | # evaluate the diffused uss |
| 6 | def eval_us(step_env, state, us): |
| 7 | def step(state, u): |
| 8 | state = step_env(state, u) |
| 9 | return state, state.reward |
| 10 | |
| 11 | _, rews = jax.lax.scan(step, state, us) |
| 12 | return rews |
| 13 | |
| 14 | def rollout_us(step_env, state, us): |
| 15 | def step(state, u): |