Set max_episode_steps

tensortrade-org · Aug 23, 2022 · 475c71e · 475c71e
1 parent 1091e33
commit 475c71e
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/tensortrade/env/generic/environment.py b/tensortrade/env/generic/environment.py
@@ -66,6 +66,7 @@ def __init__(self,
                  informer: Informer,
                  renderer: Renderer,
                  min_periods: int = None,
+                 max_episode_steps: int = None,
                  random_start_pct: float = 0.00,
                  **kwargs) -> None:
         super().__init__()
@@ -80,6 +81,13 @@ def __init__(self,
         self.min_periods = min_periods
         self.random_start_pct = random_start_pct
 
+        # Register the environment in Gym and fetch spec
+        gym.envs.register(
+            id='TensorTrade-v0',
+            max_episode_steps=max_episode_steps,
+        )
+        self.spec = gym.spec(id='TensorTrade-v0')
+
         for c in self.components.values():
             c.clock = self.clock