Internal change

PiperOrigin-RevId: 524095538 Change-Id: I01c75f74490a394dcdad0dbcbe807748bad62c21
google · Apr 13, 2023 · ff3ff64 · ff3ff64
1 parent d9049fa
commit ff3ff64
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/brax/training/acting.py b/brax/training/acting.py
@@ -15,7 +15,7 @@
 """Brax training acting functions."""
 
 import time
-from typing import Callable, Sequence, Tuple
+from typing import Callable, Sequence, Tuple, Union
 
 from brax import envs
 from brax.training.types import Metrics
@@ -27,8 +27,8 @@
 import jax
 import numpy as np
 
-State = envs.State | envs_v1.State
-Env = envs.Env | envs_v1.Env | envs_v1.Wrapper
+State = Union[envs.State, envs_v1.State]
+Env = Union[envs.Env, envs_v1.Env, envs_v1.Wrapper]
 
 
 def actor_step(