Method
reset_ref
(self, rng: jp.ndarray, ref_traj: jp.ndarray, mask: jp.ndarray, text_embedding: jp.ndarray)
Source from the content-addressed store, hash-verified
| 24 | """Automatically resets Brax envs that are done.""" |
| 25 | |
| 26 | def reset_ref(self, rng: jp.ndarray, ref_traj: jp.ndarray, mask: jp.ndarray, text_embedding: jp.ndarray) -> brax_env.State: |
| 27 | state = self.env.reset_ref(rng, ref_traj, mask, text_embedding) |
| 28 | state.info['first_qp'] = state.qp |
| 29 | state.info['first_obs'] = state.obs |
| 30 | return state |
| 31 | |
| 32 | |
| 33 | class EpisodeWrapper(wrappers.EpisodeWrapper): |
Callers
nothing calls this directly
Tested by
no test coverage detected