Run one timestep of the environment's dynamics.
(self, state: State, action: jax.Array)
| 76 | |
| 77 | @partial(jax.jit, static_argnums=(0,)) |
| 78 | def step(self, state: State, action: jax.Array) -> State: |
| 79 | """Run one timestep of the environment's dynamics.""" |
| 80 | action = jnp.clip(action, -1.0, 1.0) |
| 81 | q = state.pipeline_state |
| 82 | q_new = rk4(car_dynamics, state.pipeline_state, action, self.dt) |
| 83 | collide = check_collision(q_new, self.obs_center, self.obs_radius) |
| 84 | q = jnp.where(collide, q, q_new) |
| 85 | reward = self.get_reward(q) |
| 86 | return state.replace(pipeline_state=q, obs=q, reward=reward, done=0.0) |
| 87 | |
| 88 | @partial(jax.jit, static_argnums=(0,)) |
| 89 | def get_reward(self, q): |
nothing calls this directly
no test coverage detected