fix support for preconstructed gym environments (#170)

* fix support for preconstructed gym environments

* add documentation to GymEnvironment

* whitespace

all/environments/gym.py

@@ -5,10 +5,27 @@
 gym.logger.set_level(40)
 
 class GymEnvironment(Environment):
+    '''
+    A wrapper for OpenAI Gym environments (see: https://gym.openai.com).
+
+    This wrapper converts the output of the gym environment to PyTorch tensors,
+    and wraps them in a State object that can be passed to an Agent.
+    This constructor supports either a string, which will be passed to the
+    gym.make(name) function, or a preconstructed gym environment. Note that
+    in the latter case, the name property is set to be the whatever the name
+    of the outermost wrapper on the environment is.
+
+    Args:
+        env: Either a string or an OpenAI gym environment
+        device (optional): the device on which tensors will be stored
+    '''
     def __init__(self, env, device=torch.device('cpu')):
-        self._name = env
         if isinstance(env, str):
+            self._name = env
             env = gym.make(env)
+        else:
+            self._name = env.__class__.__name__
+
         self._env = env
         self._state = None
         self._action = None

all/environments/gym_test.py

@@ -1,8 +1,17 @@
 import unittest
+import gym
 from all.environments import GymEnvironment
 
 
 class GymEnvironmentTest(unittest.TestCase):
+    def test_env_name(self):
+        env = GymEnvironment('CartPole-v0')
+        self.assertEqual(env.name, 'CartPole-v0')
+
+    def test_preconstructed_env_name(self):
+        env = GymEnvironment(gym.make('Blackjack-v0'))
+        self.assertEqual(env.name, 'BlackjackEnv')
+
     def test_reset(self):
         env = GymEnvironment('CartPole-v0')
         state = env.reset()
@@ -11,6 +20,14 @@ def test_reset(self):
         self.assertFalse(state.done)
         self.assertEqual(state.mask, 1)
 
+    def test_reset_preconstructed_env(self):
+        env = GymEnvironment(gym.make('CartPole-v0'))
+        state = env.reset()
+        self.assertEqual(state.observation.shape, (4,))
+        self.assertEqual(state.reward, 0)
+        self.assertFalse(state.done)
+        self.assertEqual(state.mask, 1)
+
     def test_step(self):
         env = GymEnvironment('CartPole-v0')
         env.reset()