Resets the environment to an initial state.
(self, rng: jax.Array)
| 71 | self.rew_xref = jax.vmap(self.get_reward)(self.xref).mean() |
| 72 | |
| 73 | def reset(self, rng: jax.Array): |
| 74 | """Resets the environment to an initial state.""" |
| 75 | return State(self.x0, self.x0, 0.0, 0.0) |
| 76 | |
| 77 | @partial(jax.jit, static_argnums=(0,)) |
| 78 | def step(self, state: State, action: jax.Array) -> State: |