(carry, unused_t)
| 64 | |
| 65 | @jax.jit |
| 66 | def f(carry, unused_t): |
| 67 | state, current_key = carry |
| 68 | current_key, next_key = jax.random.split(current_key) |
| 69 | nstate, (qp, latent) = actor_step( |
| 70 | env, state, policy, encoder, current_key, extra_fields=extra_fields) |
| 71 | return (nstate, next_key), (qp, latent) |
| 72 | |
| 73 | (final_state, _), (qp_list, latent_list) = jax.lax.scan( |
| 74 | f, (env_state, key), (), length=unroll_length) |
nothing calls this directly
no test coverage detected