Repository: tensorflow/agents

Issue:
PyEnvironment Methods Incompatible with TF
The docstring for `tf_py_environment.__getattr__` indicates that certain PyEnvironment methods might be incompatible with TF.  

```python
def __getattr__(self, name: Text) -> Any:
    """Enables access attributes of the wrapped PyEnvironment.
    Use with caution since methods of the PyEnvironment can be incompatible
    with TF.
    Args:
      name: Name of the attribute.
    Returns:
      The attribute.
    """
    if name in self.__dict__:
      return getattr(self, name)
    return getattr(self._env, name)
```
What makes a Py Environment method incompatible with tensorflow?  

I ran across this issue when trying to call `.seed` on the `tf_py_environment`.  I implemented a `.seed` function for my subclass of `py_environment`, but calling `.seed` on the wrapper doesn't lead to the `.seed` function of the subclass being called.  Perhaps this is intentional?

Discussion:
@egonina: cc @kbanoop, can you PTAL at this question about envs? Thanks!@kbanoop: Hi, This docstring is just meant to indicate that the methods in PyEnvironment do not accept/return Tensors or do other things that are not TF compatible. 

But calling .seed on TFPyEnvironment should call .seed on self._env, whether that function is compatible or not. If that is not happening, it is likely a bug. What happens if you print(wrapper.seed)?@ergsfe: ```python
import tensorflow as tf
import numpy as np
import tf_agents
import tf_agents.environments.py_environment
import tf_agents.environments.tf_py_environment
import tf_agents.specs
class ExamplePyEnv(tf_agents.environments.py_environment.PyEnvironment):
    def __init__(self):
        super().__init__()
    def _step(self,action):
        return np.array([0.0])
    def _reset(self,action):
        return np.array([0.0])
    def seed(self,seed):
        return seed
    def action_spec(self):
        return tf_agents.specs.ArraySpec(shape=(),dtype=np.int32)
    def observation_spec(self):
        return tf_agents.specs.ArraySpec(shape=(1,),dtype=np.float32)
epe = ExamplePyEnv()
epe.seed(50)
tfpyenv = tf_agents.environments.tf_py_environment.TFPyEnvironment(epe)
try:
    tfpyenv.seed(51)
except Exception as e:
    print(repr(e))
print(tfpyenv.seed)
```
```
NotImplementedError('No seed support for this environment.')
<bound method PyEnvironment.seed of <tf_agents.environments.batched_py_environment.BatchedPyEnvironment object at 0x0000018E2E8A6BE0>>
```
@ebrevdo: A full stack trace would be more useful, but I think the issue comes from the fact that we wrap the env in a BatchedPyEnv, for which we haven't implemented batched seeding.  Can you provide the full stack trace to confirm?@ergsfe: I think that the .seed function is called on the `BatchedPyEnv`, which doesn't implement it.  

```python
NotImplementedError('No seed support for this environment.')
Traceback (most recent call last):
  File "<ipython-input-1-eee22a88ac7c>", line 26, in <module>
    tfpyenv.seed(51)
  File "d:\python\env\lib\site-packages\tf_agents\environments\py_environment.py", line 260, in seed
    raise NotImplementedError('No seed support for this environment.')
NotImplementedError: No seed support for this environment.

<bound method PyEnvironment.seed of <tf_agents.environments.batched_py_environment.BatchedPyEnvironment object at 0x00000237A2DBF340>>
```@ebrevdo: Would you be interested in submitting a PR adding `BatchedPyEnvironment` seed setters and getters that accept/return lists of seeds?@ergsfe: I can try doing that.  

Solution patch:
diff --git a/tf_agents/environments/batched_py_environment.py b/tf_agents/environments/batched_py_environment.py
index 99fbc2b3..b5041374 100644
--- a/tf_agents/environments/batched_py_environment.py
+++ b/tf_agents/environments/batched_py_environment.py
@@ -26,7 +26,7 @@ from __future__ import print_function
 from multiprocessing import dummy as mp_threads
 from multiprocessing import pool
 # pylint: enable=line-too-long
-from typing import Sequence, Optional
+from typing import Any, Optional, Sequence
 
 import gin
 import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import
@@ -182,6 +182,21 @@ class BatchedPyEnvironment(py_environment.PyEnvironment):
       )
       return nest_utils.stack_nested_arrays(time_steps)
 
+  def seed(self, seed: types.Seed) -> Any:
+    """Seeds the environment."""
+    return self._execute(lambda env: env.seed(seed), self._envs)
+
+  def get_state(self) -> Any:
+    """Returns the `state` of the environment."""
+    return self._execute(lambda env: env.get_state(), self._envs)
+
+  def set_state(self, state: Sequence[Any]) -> None:
+    """Restores the environment to a given `state`."""
+    self._execute(
+        lambda env_state: env_state[0].set_state(env_state[1]),
+        zip(self._envs, state)
+    )
+
   def render(self, mode="rgb_array") -> Optional[types.NestedArray]:
     if self._num_envs == 1:
       img = self._envs[0].render(mode)


Test patch:
diff --git a/tf_agents/environments/batched_py_environment_test.py b/tf_agents/environments/batched_py_environment_test.py
index 9cdf9637..3fc6e4a4 100644
--- a/tf_agents/environments/batched_py_environment_test.py
+++ b/tf_agents/environments/batched_py_environment_test.py
@@ -38,10 +38,21 @@ class GymWrapperEnvironmentMock(random_py_environment.RandomPyEnvironment):
   def __init__(self, *args, **kwargs):
     super(GymWrapperEnvironmentMock, self).__init__(*args, **kwargs)
     self._info = {}
+    self._state = {'seed': 0}
 
   def get_info(self):
     return self._info
 
+  def seed(self, seed):
+    self._state['seed'] = seed
+    return super(GymWrapperEnvironmentMock, self).seed(seed)
+
+  def get_state(self):
+    return self._state
+
+  def set_state(self, state):
+    self._state = state
+
   def _step(self, action):
     self._info['last_action'] = action
     return super(GymWrapperEnvironmentMock, self)._step(action)
@@ -116,6 +127,32 @@ class BatchedPyEnvironmentTest(tf.test.TestCase, parameterized.TestCase):
     self.assertAllEqual(info['last_action'], action)
     gym_env.close()
 
+  @parameterized.parameters(*COMMON_PARAMETERS)
+  def test_seed_gym_env(self, multithreading):
+    num_envs = 5
+    gym_env = self._make_batched_mock_gym_py_environment(
+        multithreading, num_envs=num_envs
+    )
+
+    gym_env.seed(42)
+
+    actual_seeds = [state['seed'] for state in gym_env.get_state()]
+    self.assertEqual(actual_seeds, [42] * num_envs)
+    gym_env.close()
+
+  @parameterized.parameters(*COMMON_PARAMETERS)
+  def test_state_gym_env(self, multithreading):
+    num_envs = 5
+    gym_env = self._make_batched_mock_gym_py_environment(
+        multithreading, num_envs=num_envs
+    )
+    state = [{'value': i * 10} for i in range(num_envs)]
+
+    gym_env.set_state(state)
+
+    self.assertEqual(gym_env.get_state(), state)
+    gym_env.close()
+
   @parameterized.parameters(*COMMON_PARAMETERS)
   def test_step(self, multithreading):
     num_envs = 5