diff --git a/python/griddly/util/environment_generator_generator.py b/python/griddly/util/environment_generator_generator.py
index 2191b7cd..08e5f1d8 100644
--- a/python/griddly/util/environment_generator_generator.py
+++ b/python/griddly/util/environment_generator_generator.py
@@ -1,5 +1,5 @@
 import os
-from typing import Optional, List, Any, Dict
+from typing import Optional, List, Any, Dict, Union
 import gymnasium as gym
 import numpy as np
 import yaml
diff --git a/python/griddly/util/rllib/environment/observer_episode_recorder.py b/python/griddly/util/rllib/environment/observer_episode_recorder.py
index 7735f6c7..b74b985d 100644
--- a/python/griddly/util/rllib/environment/observer_episode_recorder.py
+++ b/python/griddly/util/rllib/environment/observer_episode_recorder.py
@@ -5,8 +5,7 @@
 
 from griddly.gym import GymWrapper
 from griddly.util.render_tools import RenderToVideo
-from griddly.wrappers import RenderWrapper
-
+from griddly.wrappers.render_wrapper import RenderWrapper
 
 class RecordingState(Enum):
     NOT_RECORDING = 1