import errors

Hi all, thanks for the code release! I followed to instructions for installing dependencies via `requirements_static.txt`. I then tried `from language_table.environments.oracles import push_oracle_rrt_slowdown` and got the following error:

```
Traceback (most recent call last):
File "tmp.py", line 1, in
from language_table.environments.oracles import push_oracle_rrt_slowdown
File "/local/vondrick/sy/repos/language-table/language_table/environments/oracles/push_oracle_rrt_slowdown.py", line 23, in
from language_table.environments.oracles import oriented_push_oracle
File "/local/vondrick/sy/repos/language-table/language_table/environments/oracles/oriented_push_oracle.py", line 22, in
from tf_agents.policies import py_policy
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tf_agents/__init__.py", line 70, in
from tf_agents import agents
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tf_agents/agents/__init__.py", line 17, in
from tf_agents.agents import behavioral_cloning
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tf_agents/agents/behavioral_cloning/__init__.py", line 17, in
from tf_agents.agents.behavioral_cloning import behavioral_cloning_agent
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tf_agents/agents/behavioral_cloning/behavioral_cloning_agent.py", line 37, in
import tensorflow_probability as tfp
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tensorflow_probability/__init__.py", line 20, in
from tensorflow_probability import substrates
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tensorflow_probability/substrates/__init__.py", line 17, in
from tensorflow_probability.python.internal import all_util
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tensorflow_probability/python/__init__.py", line 138, in
dir(globals()[pkg_name]) # Forces loading the package from its lazy loader.
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tensorflow_probability/python/internal/lazy_loader.py", line 57, in __dir__
module = self._load()
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tensorflow_probability/python/internal/lazy_loader.py", line 40, in _load
module = importlib.import_module(self.__name__)
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/importlib/__init__.py", line 127, in import_module
return _bootstrap._gcd_import(name[level:], package, level)
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tensorflow_probability/python/experimental/__init__.py", line 31, in
from tensorflow_probability.python.experimental import bijectors
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tensorflow_probability/python/experimental/bijectors/__init__.py", line 17, in
from tensorflow_probability.python.bijectors.ldj_ratio import forward_log_det_jacobian_ratio
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tensorflow_probability/python/bijectors/__init__.py", line 19, in
from tensorflow_probability.python.bijectors.absolute_value import AbsoluteValue
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tensorflow_probability/python/bijectors/absolute_value.py", line 19, in
from tensorflow_probability.python.bijectors import bijector
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tensorflow_probability/python/bijectors/bijector.py", line 26, in
from tensorflow_probability.python.internal import batch_shape_lib
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tensorflow_probability/python/internal/batch_shape_lib.py", line 23, in
from tensorflow_probability.python.internal import prefer_static as ps
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tensorflow_probability/python/internal/prefer_static.py", line 421, in
size = _copy_docstring(tf.size, _size)
File "/local/vondrick/sy/miniconda3/envs/lt/lib/python3.8/site-packages/tensorflow_probability/python/internal/prefer_static.py", line 91, in _copy_docstring
raise ValueError(
ValueError: Arg specs do not match: original=FullArgSpec(args=['input', 'out_type', 'name'], varargs=None, varkw=None, defaults=(None, None), kwonlyargs=[], kwonlydefaults=None, annotations={}), new=FullArgSpec(args=['input', 'out_type', 'name'], varargs=None, varkw=None, defaults=(tf.int32, None), kwonlyargs=[], kwonlydefaults=None, annotations={}), fn=
```

Not super familiar with the `tf_agents` library, any idea what could be up?