diff --git a/hw1/cs285/infrastructure/rl_trainer.py b/hw1/cs285/infrastructure/rl_trainer.py
index b78d8eb..1a36534 100644
--- a/hw1/cs285/infrastructure/rl_trainer.py
+++ b/hw1/cs285/infrastructure/rl_trainer.py
@@ -4,6 +4,7 @@
 import time
 
 import gym
+from gym import spaces
 import torch
 
 from cs285.infrastructure import pytorch_util as ptu
@@ -49,7 +50,7 @@ def __init__(self, params):
         MAX_VIDEO_LEN = self.params['ep_len']
 
         # Is this env continuous, or self.discrete?
-        discrete = isinstance(self.env.action_space, gym.spaces.Discrete)
+        discrete = isinstance(self.env.action_space, spaces.Discrete)  # type: ignore
         self.params['agent_params']['discrete'] = discrete
 
         # Observation and action sizes