(self, state: State, action: jnp.ndarray)
| 37 | return State(pipeline_state, obs, reward, done, metrics) |
| 38 | |
| 39 | def step(self, state: State, action: jnp.ndarray) -> State: |
| 40 | pipeline_state = self.pipeline_step(state.pipeline_state, action) |
| 41 | obs = self._get_obs(pipeline_state) |
| 42 | reward = self._get_reward(pipeline_state) |
| 43 | done = self._get_done(pipeline_state) |
| 44 | return state.replace( |
| 45 | pipeline_state=pipeline_state, obs=obs, reward=reward, done=done |
| 46 | ) |
| 47 | |
| 48 | def _get_obs(self, pipeline_state: pipeline.State) -> jnp.ndarray: |
| 49 | return jnp.concat([pipeline_state.q, pipeline_state.qd], axis=-1) |
nothing calls this directly
no test coverage detected