Move vectorized sampler from tf to garage.sampler (#840)

- refactor tests accordingly
- fix tf_vars initialize bug in examples/np

examples/np/cem_cartpole.py

@@ -11,10 +11,10 @@
 from garage.experiment import run_experiment
 from garage.np.algos import CEM
 from garage.np.baselines import LinearFeatureBaseline
+from garage.sampler import OnPolicyVectorizedSampler
 from garage.tf.envs import TfEnv
 from garage.tf.experiment import LocalTFRunner
 from garage.tf.policies import CategoricalMLPPolicy
-from garage.tf.samplers import OnPolicyVectorizedSampler
 
 
 def run_task(snapshot_config, *_):
@@ -27,6 +27,8 @@ def run_task(snapshot_config, *_):
 
         baseline = LinearFeatureBaseline(env_spec=env.spec)
 
+        runner.initialize_tf_vars()
+
         n_samples = 20
 
         algo = CEM(

examples/np/cma_es_cartpole.py

@@ -11,10 +11,10 @@
 from garage.experiment import run_experiment
 from garage.np.algos import CMAES
 from garage.np.baselines import LinearFeatureBaseline
+from garage.sampler import OnPolicyVectorizedSampler
 from garage.tf.envs import TfEnv
 from garage.tf.experiment import LocalTFRunner
 from garage.tf.policies import CategoricalMLPPolicy
-from garage.tf.samplers import OnPolicyVectorizedSampler
 
 
 def run_task(snapshot_config, *_):
@@ -27,6 +27,8 @@ def run_task(snapshot_config, *_):
 
         baseline = LinearFeatureBaseline(env_spec=env.spec)
 
+        runner.initialize_tf_vars()
+
         n_samples = 20
 
         algo = CMAES(

src/garage/np/algos/batch_polopt.py

@@ -5,7 +5,8 @@
 
 from garage.misc import special, tensor_utils
 from garage.np.algos.base import RLAlgorithm
-from garage.tf.samplers import BatchSampler, OnPolicyVectorizedSampler
+from garage.sampler import OnPolicyVectorizedSampler
+from garage.tf.samplers import BatchSampler
 
 
 class BatchPolopt(RLAlgorithm):

src/garage/np/algos/off_policy_rl_algorithm.py

@@ -1,5 +1,5 @@
 from garage.np.algos import RLAlgorithm
-from garage.tf.samplers import OffPolicyVectorizedSampler
+from garage.sampler import OffPolicyVectorizedSampler
 
 
 class OffPolicyRLAlgorithm(RLAlgorithm):

src/garage/sampler/__init__.py

@@ -2,12 +2,17 @@
 from garage.sampler.base import Sampler
 from garage.sampler.batch_sampler import BatchSampler
 from garage.sampler.is_sampler import ISSampler
+from garage.sampler.off_policy_vectorized_sampler import (
+    OffPolicyVectorizedSampler)
+from garage.sampler.on_policy_vectorized_sampler import (
+    OnPolicyVectorizedSampler)
 from garage.sampler.parallel_vec_env_executor import ParallelVecEnvExecutor
 from garage.sampler.ray_sampler import RaySampler, SamplerWorker
 from garage.sampler.stateful_pool import singleton_pool
 from garage.sampler.vec_env_executor import VecEnvExecutor
 
 __all__ = [
     'BaseSampler', 'BatchSampler', 'Sampler', 'ISSampler', 'singleton_pool',
-    'RaySampler', 'SamplerWorker', 'ParallelVecEnvExecutor', 'VecEnvExecutor'
+    'RaySampler', 'SamplerWorker', 'ParallelVecEnvExecutor', 'VecEnvExecutor',
+    'OffPolicyVectorizedSampler', 'OnPolicyVectorizedSampler'
 ]

src/garage/sampler/batch_sampler.py

@@ -13,7 +13,7 @@ class BatchSampler(BaseSampler):
     """
 
     def __init__(self, algo, env):
-        super(BatchSampler, self).__init__(algo, env)
+        super().__init__(algo, env)
 
     def start_worker(self):
         """Start worker function."""

src/garage/tf/samplers/off_policy_vectorized_sampler.py → src/garage/sampler/off_policy_vectorized_sampler.py

@@ -15,8 +15,8 @@
 
 from garage.misc import tensor_utils
 from garage.misc.overrides import overrides
-from garage.sampler import VecEnvExecutor
-from garage.tf.samplers.batch_sampler import BatchSampler
+from garage.sampler.batch_sampler import BatchSampler
+from garage.sampler.vec_env_executor import VecEnvExecutor
 
 
 class OffPolicyVectorizedSampler(BatchSampler):
@@ -33,7 +33,7 @@ class OffPolicyVectorizedSampler(BatchSampler):
     def __init__(self, algo, env, n_envs=None, no_reset=True):
         if n_envs is None:
             n_envs = int(algo.rollout_batch_size)
-        super(OffPolicyVectorizedSampler, self).__init__(algo, env, n_envs)
+        super().__init__(algo, env)
         self.n_envs = n_envs
         self.no_reset = no_reset
 

src/garage/tf/samplers/on_policy_vectorized_sampler.py → src/garage/sampler/on_policy_vectorized_sampler.py

@@ -7,14 +7,14 @@
 from garage.misc import tensor_utils
 from garage.misc.overrides import overrides
 from garage.misc.prog_bar_counter import ProgBarCounter
-from garage.sampler import VecEnvExecutor
+from garage.sampler.batch_sampler import BatchSampler
 from garage.sampler.utils import truncate_paths
-from garage.tf.samplers.batch_sampler import BatchSampler
+from garage.sampler.vec_env_executor import VecEnvExecutor
 
 
 class OnPolicyVectorizedSampler(BatchSampler):
     def __init__(self, algo, env, n_envs=1):
-        super(OnPolicyVectorizedSampler, self).__init__(algo, env, n_envs)
+        super().__init__(algo, env)
         self.n_envs = n_envs
 
     @overrides

src/garage/tf/algos/batch_polopt.py

@@ -5,8 +5,9 @@
 
 from garage.misc import special
 from garage.np.algos import RLAlgorithm
+from garage.sampler import OnPolicyVectorizedSampler
 from garage.tf.misc import tensor_utils
-from garage.tf.samplers import BatchSampler, OnPolicyVectorizedSampler
+from garage.tf.samplers import BatchSampler
 
 
 class BatchPolopt(RLAlgorithm):

src/garage/tf/samplers/__init__.py

@@ -1,11 +1,4 @@
 from garage.tf.samplers.batch_sampler import BatchSampler
-from garage.tf.samplers.off_policy_vectorized_sampler import (
-    OffPolicyVectorizedSampler)
-from garage.tf.samplers.on_policy_vectorized_sampler import (
-    OnPolicyVectorizedSampler)
 from garage.tf.samplers.ray_sampler import (RaySamplerTF, SamplerWorkerTF)
 
-__all__ = [
-    'BatchSampler', 'OffPolicyVectorizedSampler', 'OnPolicyVectorizedSampler',
-    'RaySamplerTF', 'SamplerWorkerTF'
-]
+__all__ = ['BatchSampler', 'RaySamplerTF', 'SamplerWorkerTF']

src/garage/tf/samplers/batch_sampler.py

@@ -1,7 +1,8 @@
 import tensorflow as tf
 
-from garage.sampler import parallel_sampler, singleton_pool
+from garage.sampler import parallel_sampler
 from garage.sampler.base import BaseSampler
+from garage.sampler.stateful_pool import singleton_pool
 from garage.sampler.utils import truncate_paths
 
 

tests/garage/np/algos/test_batch_polopt.py

@@ -2,8 +2,9 @@
 
 from garage.envs import normalize
 from garage.np.baselines import LinearFeatureBaseline
+from garage.sampler import OnPolicyVectorizedSampler
 from garage.tf.envs import TfEnv
-from garage.tf.samplers import BatchSampler, OnPolicyVectorizedSampler
+from garage.tf.samplers import BatchSampler
 from tests.fixtures.algos import DummyAlgo
 from tests.fixtures.policies import DummyPolicy, DummyPolicyWithoutVectorized
 

tests/garage/np/algos/test_cem.py

@@ -2,10 +2,10 @@
 
 from garage.np.algos import CEM
 from garage.np.baselines import LinearFeatureBaseline
+from garage.sampler import OnPolicyVectorizedSampler
 from garage.tf.envs import TfEnv
 from garage.tf.experiment import LocalTFRunner
 from garage.tf.policies import CategoricalMLPPolicy
-from garage.tf.samplers import OnPolicyVectorizedSampler
 from tests.fixtures import TfGraphTestCase
 
 

tests/garage/np/algos/test_cma_es.py

@@ -1,9 +1,9 @@
 from garage.np.algos import CMAES
 from garage.np.baselines import LinearFeatureBaseline
+from garage.sampler import OnPolicyVectorizedSampler
 from garage.tf.envs import TfEnv
 from garage.tf.experiment import LocalTFRunner
 from garage.tf.policies import CategoricalMLPPolicy
-from garage.tf.samplers import OnPolicyVectorizedSampler
 from tests.fixtures import TfGraphTestCase
 
 

tests/garage/tf/samplers/test_off_policy_vectorized_sampler.py → tests/garage/sampler/test_off_policy_vectorized_sampler_integration.py

@@ -5,12 +5,12 @@
 from garage.envs import normalize
 from garage.np.exploration_strategies import OUStrategy
 from garage.replay_buffer import SimpleReplayBuffer
+from garage.sampler import OffPolicyVectorizedSampler
 from garage.tf.algos import DDPG
 from garage.tf.envs import TfEnv
 from garage.tf.experiment import LocalTFRunner
 from garage.tf.policies import ContinuousMLPPolicyWithModel
 from garage.tf.q_functions import ContinuousMLPQFunction
-from garage.tf.samplers import OffPolicyVectorizedSampler
 from tests.fixtures import TfGraphTestCase
 from tests.fixtures.envs.dummy import DummyDictEnv
 from tests.fixtures.policies import DummyPolicy

tests/garage/sampler/test_ray_batched_sampler.py

@@ -5,10 +5,8 @@
 
 from garage.envs.grid_world_env import GridWorldEnv
 from garage.np.policies import ScriptedPolicy
-from garage.sampler import RaySampler, SamplerWorker
+from garage.sampler import OnPolicyVectorizedSampler, RaySampler, SamplerWorker
 from garage.tf.envs import TfEnv
-from garage.tf.samplers.on_policy_vectorized_sampler \
-    import OnPolicyVectorizedSampler
 
 
 class TestSampler:

tests/garage/tf/algos/test_batch_polopt.py

@@ -2,8 +2,9 @@
 
 from garage.envs import normalize
 from garage.np.baselines import LinearFeatureBaseline
+from garage.sampler import OnPolicyVectorizedSampler
 from garage.tf.envs import TfEnv
-from garage.tf.samplers import BatchSampler, OnPolicyVectorizedSampler
+from garage.tf.samplers import BatchSampler
 from tests.fixtures.algos import DummyTFAlgo
 from tests.fixtures.policies import DummyPolicy, DummyPolicyWithoutVectorized