Collect data.
(
env: envs.Env,
env_state: envs.State,
policy: Policy,
encoder: Policy,
key: PRNGKey,
extra_fields: Sequence[str] = ()
)
| 32 | |
| 33 | |
| 34 | def actor_step( |
| 35 | env: envs.Env, |
| 36 | env_state: envs.State, |
| 37 | policy: Policy, |
| 38 | encoder: Policy, |
| 39 | key: PRNGKey, |
| 40 | extra_fields: Sequence[str] = () |
| 41 | ): |
| 42 | """Collect data.""" |
| 43 | key, key_encoder = jax.random.split(key) |
| 44 | latent, kl_div = encoder(env_state.obs, key_encoder) |
| 45 | obs_latent = jnp.concatenate([jnp.split(env_state.obs, 2, axis=-1)[0], latent], axis=-1) |
| 46 | actions, policy_extras = policy(obs_latent, key) |
| 47 | nstate = env.step(env_state, actions) |
| 48 | nstate.metrics.update(kl_div=jax.lax.stop_gradient(kl_div)) |
| 49 | # nstate = nstate.replace(reward=nstate.reward - kl_div) |
| 50 | state_extras = {x: nstate.info[x] for x in extra_fields} |
| 51 | return nstate, (env_state.qp, latent) |
| 52 | |
| 53 | |
| 54 | def generate_unroll( |