Runs one timestep of the environment's dynamics.
(self, state: State, action: jax.Array)
| 61 | return State(pipeline_state, obs, reward, done, metrics) |
| 62 | |
| 63 | def step(self, state: State, action: jax.Array) -> State: |
| 64 | """Runs one timestep of the environment's dynamics.""" |
| 65 | pipeline_state = self.pipeline_step(state.pipeline_state, action) |
| 66 | # set reference state for visualization |
| 67 | for i, idx in enumerate(self.ref_body_idx): |
| 68 | pipeline_state = pipeline_state.replace( |
| 69 | x=pipeline_state.x.replace( |
| 70 | pos=pipeline_state.x.pos.at[idx].set( |
| 71 | self.xref[i, jnp.int32(state.done)] |
| 72 | ), |
| 73 | ) |
| 74 | ) |
| 75 | # quad_impact_cost is not computed here |
| 76 | |
| 77 | obs = self._get_obs(pipeline_state) |
| 78 | reward = self._get_reward(state) |
| 79 | |
| 80 | return state.replace( |
| 81 | pipeline_state=pipeline_state, obs=obs, reward=reward, done=state.done + 1 |
| 82 | ) |
| 83 | |
| 84 | def _get_obs(self, pipeline_state: base.State) -> jax.Array: |
| 85 | return jnp.concatenate([pipeline_state.q, pipeline_state.qd], axis=-1) |
nothing calls this directly
no test coverage detected