Skip to content

Commit

Permalink
type hinting
Browse files Browse the repository at this point in the history
  • Loading branch information
reginald-mclean committed Aug 30, 2024
1 parent 3a1d74f commit 19f2062
Showing 1 changed file with 9 additions and 5 deletions.
14 changes: 9 additions & 5 deletions metaworld/wrappers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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)
Expand Down

0 comments on commit 19f2062

Please sign in to comment.