| 157 | |
| 158 | |
| 159 | def env_step(carry: Tuple[envs.State, PRNGKey], step_index: int, |
| 160 | policy: types.Policy, encoder: types.Policy): |
| 161 | env_state, key = carry |
| 162 | key, key_sample = jax.random.split(key) |
| 163 | key_encoder, key_sample = jax.random.split(key_sample) |
| 164 | latent, kl_div = encoder(env_state.obs, key_encoder) |
| 165 | obs_latent = jnp.concatenate([jnp.split(env_state.obs, 2, axis=-1)[0], latent], axis=-1) |
| 166 | actions = policy(obs_latent, key_sample)[0] |
| 167 | nstate = env.step(env_state, actions) |
| 168 | nstate.metrics.update(kl_div=jax.lax.stop_gradient(kl_div)) |
| 169 | nstate = nstate.replace(reward=nstate.reward - kl_div * beta) |
| 170 | if truncation_length is not None: |
| 171 | nstate = jax.lax.cond( |
| 172 | jnp.mod(step_index + 1, truncation_length) == 0., |
| 173 | jax.lax.stop_gradient, lambda x: x, nstate) |
| 174 | |
| 175 | return (nstate, key), (nstate.reward, env_state.obs, obs_latent, nstate.metrics) |
| 176 | |
| 177 | def loss(cvae_params, normalizer_params, key, ref_traj, mask): |
| 178 | encoder_params, policy_params = cvae_params[0], cvae_params[1] |