Automatically resets Brax envs that are done.
| 21 | |
| 22 | |
| 23 | class AutoResetWrapper(wrappers.AutoResetWrapper): |
| 24 | """Automatically resets Brax envs that are done.""" |
| 25 | |
| 26 | def reset_ref(self, rng: jp.ndarray, ref_traj: jp.ndarray, mask: jp.ndarray) -> brax_env.State: |
| 27 | state = self.env.reset_ref(rng, ref_traj, mask) |
| 28 | state.info['first_qp'] = state.qp |
| 29 | state.info['first_obs'] = state.obs |
| 30 | return state |
| 31 | |
| 32 | |
| 33 | class EpisodeWrapper(wrappers.EpisodeWrapper): |
nothing calls this directly
no outgoing calls
no test coverage detected