Runs one timestep of the environment's dynamics.
(self, state: State, action: jax.Array)
| 35 | return State(pipeline_state, obs, reward, done, {}) |
| 36 | |
| 37 | def step(self, state: State, action: jax.Array) -> State: |
| 38 | """Runs one timestep of the environment's dynamics.""" |
| 39 | pipeline_state0 = state.pipeline_state |
| 40 | assert pipeline_state0 is not None |
| 41 | pipeline_state = self.pipeline_step(pipeline_state0, action) |
| 42 | |
| 43 | obs = self._get_obs(pipeline_state) |
| 44 | reward = self._get_reward(pipeline_state) |
| 45 | |
| 46 | return state.replace( |
| 47 | pipeline_state=pipeline_state, obs=obs, reward=reward, done=0.0 |
| 48 | ) |
| 49 | |
| 50 | def _get_obs(self, pipeline_state: base.State) -> jax.Array: |
| 51 | """Returns the environment observations.""" |
nothing calls this directly
no test coverage detected