From 19f206273cae6390e4603bd102c1ab39dbd1d2e5 Mon Sep 17 00:00:00 2001 From: Reggie McLean Date: Fri, 30 Aug 2024 12:40:14 -0400 Subject: [PATCH] type hinting --- metaworld/wrappers.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/metaworld/wrappers.py b/metaworld/wrappers.py index 448ba566..3d9a29d0 100644 --- a/metaworld/wrappers.py +++ b/metaworld/wrappers.py @@ -59,13 +59,15 @@ def __init__( def toggle_sample_tasks_on_reset(self, on: bool): self.sample_tasks_on_reset = on - def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None): + def reset( + self, *, seed: Optional[int] = None, options: Optional[dict[str, Any]] = None + ): if self.sample_tasks_on_reset: self._set_random_task() return self.env.reset(seed=seed, options=options) def sample_tasks( - self, *, seed: int | None = None, options: dict[str, Any] | None = None + self, *, seed: Optional[int] = None, options: Optional[dict[str, Any]] = None ): self._set_random_task() return self.env.reset(seed=seed, options=options) @@ -98,7 +100,7 @@ def __init__( env: Env, tasks: list[Task], sample_tasks_on_reset: bool = False, - seed: int | None = None, + seed: Optional[int] = None, ): super().__init__(env) self.sample_tasks_on_reset = sample_tasks_on_reset @@ -107,13 +109,15 @@ def __init__( if seed: np.random.seed(seed) - def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None): + def reset( + self, *, seed: Optional[int] = None, options: Optional[dict[str, Any]] = None + ): if self.sample_tasks_on_reset: self._set_pseudo_random_task() return self.env.reset(seed=seed, options=options) def sample_tasks( - self, *, seed: int | None = None, options: dict[str, Any] | None = None + self, *, seed: Optional[int] = None, options: Optional[dict[str, Any]] = None ): self._set_pseudo_random_task() return self.env.reset(seed=seed, options=options)