(self, rng: jp.ndarray, ref_traj: jp.ndarray, mask: jp.ndarray, text_embedding:jp.ndarray)
| 44 | """Vectorizes Brax env.""" |
| 45 | |
| 46 | def reset_ref(self, rng: jp.ndarray, ref_traj: jp.ndarray, mask: jp.ndarray, text_embedding:jp.ndarray) -> brax_env.State: |
| 47 | return jp.vmap(self.env.reset_ref)(rng, ref_traj, mask, text_embedding) |
nothing calls this directly
no outgoing calls
no test coverage detected