Resets the environment to an initial state.
(self, rng: jax.Array)
| 46 | super().__init__(sys=sys, backend="positional", n_frames=5) |
| 47 | |
| 48 | def reset(self, rng: jax.Array) -> State: |
| 49 | """Resets the environment to an initial state.""" |
| 50 | |
| 51 | qpos = self.sys.init_q |
| 52 | qvel = jnp.zeros(self.sys.qd_size()) |
| 53 | |
| 54 | pipeline_state = self.pipeline_init(qpos, qvel) |
| 55 | obs = self._get_obs(pipeline_state) |
| 56 | reward, done, zero = jnp.zeros(3) |
| 57 | metrics = { |
| 58 | "reward_linup": zero, |
| 59 | "reward_quadctrl": zero, |
| 60 | } |
| 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.""" |