minor refactoring to avoid pylance error

hyeok9855 · Sep 8, 2023 · eccd142 · eccd142
1 parent dc5b2a4
commit eccd142
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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