Repository: rlworkgroup/garage

Issue:
CEM and CMA don't support parallel sampling from TF policy
Currently, both ES are under the original branch. I managed to run them with some minor modifications (added the session initialization) with a single sampler but running with multiple samplers causes problems (that completely undermines the advantages of ES). 
Are there any immediate plans to move both algorithms into the main TF branch?


Discussion:
They are not in garage.tf because (to my knowledge) they do not depend on TensorFlow.

Can you post the error messages you got? These are both tested by the CI but perhaps some of the support code (e.g. sampler and/or plotter) have TF dependencies?
WHen one has to extract parameters from a policy one has to run a session. The absence of a default session is what it complained about initially (that is what I fixed by simply creating a session in the training routine). But it only works with a single sampler. When I am trying to increase n_parallel it keeps complaining about the lack of default session. I haven't dug into that yet hoping that maybe there was a plan to fix it. Here is the full post of the error I got now:
```sh
(garage) artem@artempc:~/prj/quad_metalearn/quad_dynalearn/quad_dynalearn$ ./train_garage_quad.py config/cem_quad.conf _results_temp/cem_quad_test --n_parallel 2
Reading parameter file config/cem_quad.conf ...
###############################################################
### PARAMETERS LOADED FROM CONFIG FILES (Later updated by arguments provided)
{'seed': 1, 'variant': {'env': 'QuadrotorEnv', 'alg_class': 'CEM', 'alg_param': {'max_path_length': 100, 'n_itr': 10, 'n_samples': 100, 'best_frac': 0.05, 'init_std': 0.1, 'plot': False, 'play_every_itr': 1, 'play_rollouts_num': 1}, 'policy_class': 'GaussianMLPPolicy', 'policy_param': {'hidden_sizes': [32, 32]}}}
+++++++++++++++++++++++++++++++++++++++++++++++++++
PARAMETERS TUPLE:  None ()  SEED:  1
python /home/artem/prj/drl/garage/scripts/run_experiment.py  --seed '1'  --n_parallel '2'  --snapshot_mode 'last'  --plot 'False'  --exp_name 'experiment_2019_01_09_14_39_27_0001'  --log_dir '_results_temp/cem_quad_test/seed_001/'  --use_cloudpickle 'True'  --args_data 'gASV3wUAAAAAAACMF2Nsb3VkcGlja2xlLmNsb3VkcGlja2xllIwOX2ZpbGxfZnVuY3Rpb26Uk5QoaACMD19tYWtlX3NrZWxfZnVuY5STlGgAjA1fYnVpbHRpbl90eXBllJOUjAhDb2RlVHlwZZSFlFKUKEsBSwBLC0sGS0NCbAEAAGQBZAJsAG0BfQEBAGQBZANsAm0DfQIBAGQBZARsBG0FfQMBAGQBZAVsBm0HfQQBAGQBZAZsCG0JfQUBAGQBZAdsCm0LfQYBAHwAZAgZAGQJawJyZnwCdAxkCmQLZAyNAoMBfQduFnwCdA10DmoPfABkCBkAgwGDAYMBfQd8AGQIPQB0EIMAfABkDRkAGQBmAGQOfAdqEWkBfABkDxkAlwKOAX0IfABkDT0AfABkDz0AfABkEBkAZBFrA5ABchJ0EIMAfABkEhkAGQBmAGQOfAdqEWkBfABkExkAlwKOAX0JfABkEj0AfABkEz0AdBCDAHwAZBAZABkAZgB8B3wIfAlkFJwDfABkFRkAlwKOAX0KbiJ0EIMAfABkEBkAGQBmAHwHfAhkFpwCfABkFRkAlwKOAX0KfABkED0AfABkFT0AfABkFz0AfABpAGsCkAFzYHQSZBh0E3wAgwEWAIMBggF8CmoUgwABAGQZUwCUKIxWCiAgICBXcmFwIFBQTyB0cmFpbmluZyB0YXNrIGluIHRoZSBydW5fdGFzayBmdW5jdGlvbi4KCiAgICA6cGFyYW0gXzoKICAgIDpyZXR1cm46CiAgICCUSwCME0dhdXNzaWFuTUxQQmFzZWxpbmWUhZSMBVRmRW52lIWUjBFHYXVzc2lhbk1MUFBvbGljeZSFlIwDQ0VNlIWUjANQUE+UhZSMBFRSUE+UhZSMA2VudpSMDFF1YWRyb3RvckVudpSIiYwLcmF3X2NvbnRyb2yUjAp0Zl9jb250cm9slIaUjAxwb2xpY3lfY2xhc3OUjAhlbnZfc3BlY5SMDHBvbGljeV9wYXJhbZSMCWFsZ19jbGFzc5RoEowOYmFzZWxpbmVfY2xhc3OUjA5iYXNlbGluZV9wYXJhbZRoGIwGcG9saWN5lIwIYmFzZWxpbmWUh5SMCWFsZ19wYXJhbZRoGGgjhpSMCGV4cF9uYW1llIwxRVJST1I6IFNvbWUgb2YgcGFyYW1ldGVyIHZhbHVlcyB3ZXJlIG5vdCB1c2VkOiAlc5ROdJQojBNnYXJhZ2UudGYuYmFzZWxpbmVzlGgMjA5nYXJhZ2UudGYuZW52c5RoDowSZ2FyYWdlLnRmLnBvbGljaWVzlGgQjBhxdWFkX2R5bmFsZWFybi5hbGdvcy5jZW2UaBKME2dhcmFnZS50Zi5hbGdvcy5wcG+UaBSMFGdhcmFnZS50Zi5hbGdvcy50cnBvlGgWaBmMCW5vcm1hbGl6ZZSMA2d5bZSMBG1ha2WUjAZsb2NhbHOUjARzcGVjlIwOQXNzZXJ0aW9uRXJyb3KUjANzdHKUjAV0cmFpbpR0lCiMCnRhc2tfcGFyYW2UaAxoDmgQaBJoFGgWaBhoI2gkjARhbGdvlHSUjBYuL3RyYWluX2dhcmFnZV9xdWFkLnB5lIwIcnVuX3Rhc2uUS0tDPgAHDAEMAQwCDAEMAQwCDAESAhYBBgIiAQYBBgIOASIBBgEGAg4BAgECAQYBDgIOAQIBBgEMAgYBBgUGARoClCkpdJRSlEr/////jAhfX21haW5fX5SHlFKUfZQojAdnbG9iYWxzlH2UKGgZjCNneW1fYXJ0LnF1YWRyb3Rvci5xdWFkcm90b3JfbW9kdWxhcpRoGZOUaDGMHGd5bV9hcnQucXVhZHJvdG9yLnF1YWRfdXRpbHOUaDGTlGgyaACMCXN1YmltcG9ydJSTlGgyhZRSlHWMCGRlZmF1bHRzlE6MBGRpY3SUfZSMDmNsb3N1cmVfdmFsdWVzlE6MBm1vZHVsZZRoQowEbmFtZZRoPowDZG9jlGgLjAhxdWFsbmFtZZRoPnV0Ui4='  --variant_data 'gAN9cQAoWAMAAABlbnZxAVgMAAAAUXVhZHJvdG9yRW52cQJYCQAAAGFsZ19jbGFzc3EDWAMAAABDRU1xBFgJAAAAYWxnX3BhcmFtcQV9cQYoWA8AAABtYXhfcGF0aF9sZW5ndGhxB0tkWAUAAABuX2l0cnEISwpYCQAAAG5fc2FtcGxlc3EJS2RYCQAAAGJlc3RfZnJhY3EKRz+pmZmZmZmaWAgAAABpbml0X3N0ZHELRz+5mZmZmZmaWAQAAABwbG90cQyJWA4AAABwbGF5X2V2ZXJ5X2l0cnENSwFYEQAAAHBsYXlfcm9sbG91dHNfbnVtcQ5LAXVYDAAAAHBvbGljeV9jbGFzc3EPWBEAAABHYXVzc2lhbk1MUFBvbGljeXEQWAwAAABwb2xpY3lfcGFyYW1xEX1xElgMAAAAaGlkZGVuX3NpemVzcRNdcRQoSyBLIGVzWAgAAABleHBfbmFtZXEVWCMAAABleHBlcmltZW50XzIwMTlfMDFfMDlfMTRfMzlfMjdfMDAwMXEWdS4='
2019-01-09 14:39:31 | Setting seed to 1
2019-01-09 14:39:31 | Setting seed to 2
2019-01-09 14:39:31 | tensorboard data will be logged into:_results_temp/cem_quad_test/seed_001/
WARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.
WARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.
2019-01-09 14:39:36.348996: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2019-01-09 14:39:36 | [experiment_2019_01_09_14_39_27_0001] Populating workers...
WARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.
WARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.
WARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.
WARN: gym.spaces.Box autodetected dtype as <class 'numpy.float32'>. Please provide explicit dtype.
*** Error in `python': double free or corruption (fasttop)multiprocessing.pool.RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/home/artem/prj/drl/garage/garage/sampler/stateful_pool.py", line 186, in _worker_run_each
    return runner(singleton_pool.G, *args)
  File "/home/artem/prj/drl/garage/garage/sampler/parallel_sampler.py", line 40, in _worker_populate_task
    g.policy = pickle.loads(policy)
  File "/home/artem/prj/drl/garage/garage/tf/core/parameterized.py", line 107, in __setstate__
    tf.get_default_session().run(
AttributeError: 'NoneType' object has no attribute 'run'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/opt/anaconda2/envs/garage/lib/python3.6/multiprocessing/pool.py", line 119, in worker
    result = (True, func(*args, **kwds))
  File "/opt/anaconda2/envs/garage/lib/python3.6/multiprocessing/pool.py", line 44, in mapstar
    return list(map(*args))
  File "/home/artem/prj/drl/garage/garage/sampler/stateful_pool.py", line 188, in _worker_run_each
    raise Exception("".join(traceback.format_exception(*sys.exc_info())))
Exception: Traceback (most recent call last):
  File "/home/artem/prj/drl/garage/garage/sampler/stateful_pool.py", line 186, in _worker_run_each
    return runner(singleton_pool.G, *args)
  File "/home/artem/prj/drl/garage/garage/sampler/parallel_sampler.py", line 40, in _worker_populate_task
    g.policy = pickle.loads(policy)
  File "/home/artem/prj/drl/garage/garage/tf/core/parameterized.py", line 107, in __setstate__
    tf.get_default_session().run(
AttributeError: 'NoneType' object has no attribute 'run'

"""

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/artem/prj/drl/garage/scripts/run_experiment.py", line 242, in <module>
    run_experiment(sys.argv)
  File "/home/artem/prj/drl/garage/scripts/run_experiment.py", line 185, in run_experiment
    method_call(variant_data)
  File "./train_garage_quad.py", line 125, in run_task
    algo.train()
  File "/home/artem/prj/quad_metalearn/quad_dynalearn/quad_dynalearn/algos/cem.py", line 137, in train
    parallel_sampler.populate_task(self.env, self.policy)
  File "/home/artem/prj/drl/garage/garage/sampler/parallel_sampler.py", line 58, in populate_task
    ] * singleton_pool.n_parallel)
  File "/home/artem/prj/drl/garage/garage/sampler/stateful_pool.py", line 72, in run_each
    return results.get()
  File "/opt/anaconda2/envs/garage/lib/python3.6/multiprocessing/pool.py", line 644, in get
    raise self._value
Exception: Traceback (most recent call last):
  File "/home/artem/prj/drl/garage/garage/sampler/stateful_pool.py", line 186, in _worker_run_each
    return runner(singleton_pool.G, *args)
  File "/home/artem/prj/drl/garage/garage/sampler/parallel_sampler.py", line 40, in _worker_populate_task
    g.policy = pickle.loads(policy)
  File "/home/artem/prj/drl/garage/garage/tf/core/parameterized.py", line 107, in __setstate__
    tf.get_default_session().run(
AttributeError: 'NoneType' object has no attribute 'run'
```
@naeioi thoughts?
@amolchanov86 Thanks for reporting. CEM itself depends on neither tf or theano. I notice that you were using `GaussianMLPPolicy` which is a tensorflow policy, that's why you need to initialize a tensorflow session manually for the policy to use. @ryanjulian CI does not break because it only tests CEM under theano policy. 

@amolchanov86 Can you also post your `train_garage_quad.py`? The magic `run_experiment` setups several processes and instruct them to run the callable `method_call` passed to it. I suspect that your tf session was not initialized within `method_call`. 
I see, so what you are trying to say is that making a session in the training routine itself is too late and it should be called earlier?


Here my full train_garage_quad.py
```
#!/usr/bin/env python
"""
This is a parametrized script to run TRPO/PPO 
with a custom env
"""
import argparse
import sys
import os
import datetime, time
import itertools
import os.path as osp
import uuid
import copy

import numpy as np

import dateutil.tz
import yaml

import gym

from garage.envs import normalize
from garage.experiment import run_experiment

# Custom stuff
import quad_dynalearn.config.config_loader as conf
import quad_dynalearn.misc.variants_utils as vu


########################################################################
## ARGUMENTS
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("config_file", help='yaml file with default settings of parameters')
parser.add_argument("log_dir", default='_results_temp/trpo_ppo_last', help='Directory to log into')
parser.add_argument("--seed", '-s', default=None, help='list of seeds to use separated by comma (or a single seed w/o comma). If None seeds from config_file will be used')
parser.add_argument("--n_parallel", '-n', type=int, default=1, help='Number of parallel workers to run a single task')
parser.add_argument("--snapshot_mode", '-snm', default='last', help='Snapshot mode. Opt: last')
parser.add_argument("--plot", '-plt', action="store_true", help='Plotting')
parser.add_argument("--param_name", '-p', help='task hyperparameter names separated by comma')
parser.add_argument("--param_val", '-pv', help='task hyperparam values.'+ 
                    ' For a single par separated by comma.' +
                    ' For adjacent params separated by double comma.' +
                    '   Ex: \"-p par1,par2 -pv pv11,pv12,,pv21,pv22\"' + 
                    '   where pv11,pv12 - par values for par1 , pv21,pv22 - par values for par2')
args = parser.parse_args()

########################################################################
## PARAMETERS (non grid)
# Loading parameters not specified in the arguments
print('Reading parameter file %s ...' % args.config_file)
params = conf.trpo_ppo_default_params()
yaml_stream = open(args.config_file, 'r')
params_new = yaml.load(yaml_stream)
params.update(params_new)
print('###############################################################')
print('### PARAMETERS LOADED FROM CONFIG FILES (Later updated by arguments provided)')
print(params)

## Get a grid of task variations and put it into list as parameter dictionaries
## WARN: when you add more parameters to add_arguments you will have to modify grid_of_variants()
variants_list = vu.grid_of_variants(args, params)

## Saving command line executing the script
cmd = " ".join(sys.argv)
if not os.path.isdir(args.log_dir):
    os.makedirs(args.log_dir)
with open(args.log_dir + os.sep + "cmd.sh", "w") as cmdfile:
    cmdfile.write("#!/usr/bin/bash\n")
    cmdfile.write(cmd)


def run_task(task_param):
    """
    Wrap PPO training task in the run_task function.

    :param _:
    :return:
    """
    from garage.tf.baselines import GaussianMLPBaseline
    from garage.tf.envs import TfEnv
    from garage.tf.policies import GaussianMLPPolicy, DeterministicMLPPolicy
    from garage.tf.algos.trpo import TRPO
    
    from quad_dynalearn.algos.cem import CEM
    from quad_dynalearn.algos.cma_es import CMAES
    from quad_dynalearn.algos.ppo import PPO

    if task_param["env"] == "QuadrotorEnv":
        # from gym_art.quadrotor.quadrotor_control import *
        from gym_art.quadrotor.quadrotor_modular import QuadrotorEnv
        env = TfEnv(QuadrotorEnv(**task_param["env_param"]))
        del task_param["env_param"]
    else:
        env = TfEnv(normalize(gym.make(task_param["env"])))
    del task_param["env"]
    
    policy = locals()[task_param["policy_class"]](env_spec=env.spec, **task_param["policy_param"])
    del task_param["policy_class"]
    del task_param["policy_param"]

    if task_param["alg_class"] != "CEM" and task_param["alg_class"] != "CMAES":
        baseline = locals()[task_param["baseline_class"]](env_spec=env.spec, **task_param["baseline_param"])
        del task_param["baseline_class"]
        del task_param["baseline_param"]

        algo = locals()[task_param["alg_class"]](
            env=env,
            policy=policy,
            baseline=baseline,
            **task_param["alg_param"])
    else:
        algo = locals()[task_param["alg_class"]](
            env=env,
            policy=policy,
            **task_param["alg_param"])

    del task_param["alg_class"]
    del task_param["alg_param"]

    # Check that we used all parameters:
    # It helps revealing situations where you thought you set certain parameter
    # But in fact made spelling mistake and it failed
    del task_param["exp_name"] #This is probably generated by garage
    assert task_param == {}, "ERROR: Some of parameter values were not used: %s" % str(task_param)

    algo.train()

start_time = time.time()
for var in variants_list:
    ## Dumping config
    with open(var["log_dir"] + os.sep + "config.yml", 'w') as yaml_file:
        yaml_file.write(yaml.dump(var, default_flow_style=False))

    ## Running
    run_experiment(
        run_task,
        **var
    )

end_time = time.time()
print("##################################################")
print("Total Runtime: ", end_time - start_time)
```
Yes. I assume the training routine you are referring to is the implementation of CEM, where there's nothing related to tf session. So you have to do it outside, earlier before algo.train() is called. 

I made some change to your script. There are some deps I don't have so I cannot test but it should work. Could you give it a try? https://gist.github.com/naeioi/28bb467abf6a598e814bdb136f7c2230
@amolchanov86 Would like to know if you have solved your problem?
Hi,
sorry, I have switched to other things and just missed your previous post. I will try today and report. Thanks a lot for the help!
I tried your example and I had the same result (i.e. works with one worker, fails if n_parallel > 1).
I even substituted the cem version I had with the original one. Hence I think if the bug is eliminated on the original version then my code should also work. Or maybe I am doing something wrong then, please, share a script that successfully runs CEM/CMAES with n_parallel > 1 and TF policies.
Thanks a lot for the support!
@amolchanov86 After some debugging, I found that both CEM and CMAES work only under theano policy when doing parallel sampling.

The root cause is that tf session is [not fork-safe](https://github.com/tensorflow/tensorflow/issues/2448), meaning that one tf session cannot be shared by multiple processes. Algos under `garage.tf` uses tf-specific samplers to manager isolated session in each worker, and since CEM and CMAES don't use those samplers, they cannot work under tf policy.

@ryanjulian @CatherineSue Do we need to add tf support for CEM and CMAES in this case?
@amolchanov86 @naeioi Thanks for pointing out the problem. I think there are two solutions: 
- Add tf support for CEM and CMAES; We always like more algorithms support. It'd be great that we have tf support for CEM and CMAES. 
- Add policies into garage.policies. Since CEM and CMAES don't depend on tf or theano, they should use primitives not depend on tf or theano, too. For instance, numpy. If we can add some policies to garage.policies, it could also solve this problem. 


uhh, these two solutions both require some effort. I think the first step we should take is to move those two algorithms under `garage.theano`. 
We are going to remove theano module soon. We also need to make sure the tf algo parity with theano. So move them to theano is equivalent to option one. And the algorithms don't actually depend on theano. The root reason is that we don't have a policy ready for the algirithm.

I prefer to leave them here then add tf support or no dep policy later.
I agree, let's just keep the issue open.

Resolving patch:
diff --git a/examples/cem_cartpole.py b/examples/cem_cartpole.py
new file mode 100644
--- /dev/null
+++ b/examples/cem_cartpole.py
@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+"""
+This is an example to train a task with Cross Entropy Method.
+
+Here it runs CartPole-v1 environment with 100 epoches.
+
+Results:
+    AverageReturn: 100
+    RiseTime: epoch 8
+"""
+from garage.experiment import LocalRunner, run_experiment
+from garage.np.algos import CEM
+from garage.np.baselines import LinearFeatureBaseline
+from garage.tf.envs import TfEnv
+from garage.tf.policies import CategoricalMLPPolicy
+from garage.tf.samplers import OnPolicyVectorizedSampler
+
+
+def run_task(*_):
+    """Train CEM with Cartpole-v1 environment."""
+    with LocalRunner() as runner:
+        env = TfEnv(env_name="CartPole-v1")
+
+        policy = CategoricalMLPPolicy(
+            name="policy", env_spec=env.spec, hidden_sizes=(32, 32))
+
+        baseline = LinearFeatureBaseline(env_spec=env.spec)
+
+        runner.initialize_tf_vars()
+
+        n_samples = 20
+
+        algo = CEM(
+            env_spec=env.spec,
+            policy=policy,
+            baseline=baseline,
+            best_frac=0.05,
+            max_path_length=100,
+            n_samples=n_samples)
+
+        runner.setup(algo, env, sampler_cls=OnPolicyVectorizedSampler)
+        # NOTE: make sure that n_epoch_cycles == n_samples !
+        runner.train(n_epochs=100, batch_size=1000, n_epoch_cycles=n_samples)
+
+
+run_experiment(
+    run_task,
+    snapshot_mode="last",
+    seed=1,
+)
diff --git a/garage/np/algos/cem.py b/garage/np/algos/cem.py
--- a/garage/np/algos/cem.py
+++ b/garage/np/algos/cem.py
@@ -1,185 +1,121 @@
-from itertools import chain, zip_longest
-
 import numpy as np
 
 from garage.core import Serializable
-from garage.logger import logger, snapshotter, tabular
-from garage.misc.special import discount_cumsum
+from garage.logger import logger, tabular
 from garage.np.algos.base import RLAlgorithm
 from garage.plotter import Plotter
-from garage.sampler import parallel_sampler, stateful_pool
-from garage.sampler.utils import rollout
-
-
-def _get_stderr_lb(x):
-    mu = np.mean(x, 0)
-    stderr = np.std(x, axis=0, ddof=1 if len(x) > 1 else 0) / np.sqrt(len(x))
-    return mu - stderr
-
-
-def _get_stderr_lb_varyinglens(x):
-    mus, stds, ns = [], [], []
-    for temp_list in zip_longest(*x, fillvalue=np.nan):
-        mus.append(np.nanmean(temp_list))
-        n = len(temp_list) - np.sum(np.isnan(temp_list))
-        stds.append(np.nanstd(temp_list, ddof=1 if n > 1 else 0))
-        ns.append(n)
-    return np.array(mus) - np.array(stds) / np.sqrt(ns)
-
-
-def _worker_rollout_policy(g, args):
-    sample_std = args['sample_std'].flatten()
-    cur_mean = args['cur_mean'].flatten()
-    n_evals = args['n_evals']
-    k = len(cur_mean)
-    params = np.random.standard_normal(k) * sample_std + cur_mean
-    g.policy.set_param_values(params)
-    paths, returns, undiscounted_returns = [], [], []
-    for _ in range(n_evals):
-        path = rollout(g.env, g.policy, args['max_path_length'])
-        path['returns'] = discount_cumsum(path['rewards'], args['discount'])
-        path['undiscounted_return'] = sum(path['rewards'])
-        paths.append(path)
-        returns.append(path['returns'])
-        undiscounted_returns.append(path['undiscounted_return'])
-
-    result_path = {'full_paths': paths}
-    result_path['undiscounted_return'] = _get_stderr_lb(undiscounted_returns)
-    result_path['returns'] = _get_stderr_lb_varyinglens(returns)
-
-    # not letting n_evals count towards below cases since n_evals is multiple
-    # eval for single paramset
-    if args['criterion'] == 'samples':
-        inc = len(path['rewards'])
-    elif args['criterion'] == 'paths':
-        inc = 1
-    else:
-        raise NotImplementedError
-    return (params, result_path), inc
 
 
 class CEM(RLAlgorithm, Serializable):
+    r"""Cross Entropy Method.
+
+    CEM works by iteratively optimizing a gaussian distribution of policy.
+
+    In each epoch, CEM does the following:
+    1. Sample n_samples policies from a gaussian distribution of
+       mean cur_mean and std cur_std.
+    2. Do rollouts for each policy.
+    3. Update cur_mean and cur_std by doing Maximum Likelihood Estimation
+       over the n_best top policies in terms of return.
+
+    Note:
+        When training CEM with LocalRunner, make sure that n_epoch_cycles for
+        runner equals to n_samples for CEM.
+
+        This implementation leverage n_epoch_cycles to do rollouts for a single
+        policy in an epoch cycle.
+
+    Attributes:
+        env_spec(EnvSpec): Environment specification.
+        policy(Policy): Action policy.
+        baseline(): Baseline for GAE (Generalized Advantage Estimation).
+        n_samples(int): Number of policies sampled in one epoch.
+        max_path_length(int):  Maximum length of a single rollout.
+        discount(float): Environment reward discount.
+        init_std(float): Initial std for policy param distribution.
+        extra_std(float): Decaying std added to param distribution.
+        extra_decay_time(float): Epochs that it takes to decay extra std.
+    """
+
     def __init__(self,
-                 env,
+                 env_spec,
                  policy,
-                 n_itr=500,
+                 baseline,
+                 n_samples,
+                 gae_lambda=1,
                  max_path_length=500,
                  discount=0.99,
-                 init_std=1.,
-                 n_samples=100,
-                 batch_size=None,
+                 init_std=1,
                  best_frac=0.05,
                  extra_std=1.,
                  extra_decay_time=100,
-                 plot=False,
-                 n_evals=1,
                  **kwargs):
-        """
-        :param n_itr: Number of iterations.
-        :param max_path_length: Maximum length of a single rollout.
-        :param batch_size: # of samples from trajs from param distribution,
-         when this is set, n_samples is ignored
-        :param discount: Discount.
-        :param plot: Plot evaluation run after each iteration.
-        :param init_std: Initial std for param distribution
-        :param extra_std: Decaying std added to param distribution at each
-         iteration
-        :param extra_decay_time: Iterations that it takes to decay extra std
-        :param n_samples: #of samples from param distribution
-        :param best_frac: Best fraction of the sampled params
-        :param n_evals: # of evals per sample from the param distr. returned
-         score is mean - stderr of evals
-        :return:
-        """
-        Serializable.quick_init(self, locals())
-        self.env = env
+        self.env_spec = env_spec
         self.policy = policy
-        self.batch_size = batch_size
-        self.plot = plot
+        self.baseline = baseline
+        self.n_samples = n_samples
         self.extra_decay_time = extra_decay_time
         self.extra_std = extra_std
         self.best_frac = best_frac
-        self.n_samples = n_samples
         self.init_std = init_std
         self.discount = discount
+        self.gae_lambda = gae_lambda
         self.max_path_length = max_path_length
-        self.n_itr = n_itr
-        self.n_evals = n_evals
         self.plotter = Plotter()
 
-    def train(self):
-        parallel_sampler.populate_task(self.env, self.policy)
-        if self.plot:
-            self.plotter.init_plot(self.env, self.policy)
-
-        cur_std = self.init_std
-        cur_mean = self.policy.get_param_values()
-        # K = cur_mean.size
-        n_best = max(1, int(self.n_samples * self.best_frac))
-
-        for itr in range(self.n_itr):
-            # sample around the current distribution
-            extra_var_mult = max(1.0 - itr / self.extra_decay_time, 0)
-            sample_std = np.sqrt(
-                np.square(cur_std) +
-                np.square(self.extra_std) * extra_var_mult)
-            if self.batch_size is None:
-                criterion = 'paths'
-                threshold = self.n_samples
-            else:
-                criterion = 'samples'
-                threshold = self.batch_size
-            infos = stateful_pool.singleton_pool.run_collect(
-                _worker_rollout_policy,
-                threshold=threshold,
-                args=(dict(
-                    cur_mean=cur_mean,
-                    sample_std=sample_std,
-                    max_path_length=self.max_path_length,
-                    discount=self.discount,
-                    criterion=criterion,
-                    n_evals=self.n_evals), ))
-            xs = np.asarray([info[0] for info in infos])
-            paths = [info[1] for info in infos]
-
-            fs = np.array([path['returns'][0] for path in paths])
-            print((xs.shape, fs.shape))
-            best_inds = (-fs).argsort()[:n_best]
-            best_xs = xs[best_inds]
-            cur_mean = best_xs.mean(axis=0)
-            cur_std = best_xs.std(axis=0)
-            best_x = best_xs[0]
-            logger.push_prefix('itr #{} | '.format(itr))
-            tabular.record('Iteration', itr)
-            tabular.record('CurStdMean', np.mean(cur_std))
-            undiscounted_returns = np.array(
-                [path['undiscounted_return'] for path in paths])
-            tabular.record('AverageReturn', np.mean(undiscounted_returns))
-            tabular.record('StdReturn', np.std(undiscounted_returns))
-            tabular.record('MaxReturn', np.max(undiscounted_returns))
-            tabular.record('MinReturn', np.min(undiscounted_returns))
-            tabular.record('AverageDiscountedReturn', np.mean(fs))
-            tabular.record('NumTrajs', len(paths))
-            paths = list(chain(
-                *[d['full_paths']
-                  for d in paths]))  # flatten paths for the case n_evals > 1
-            tabular.record('AvgTrajLen',
-                           np.mean([len(path['returns']) for path in paths]))
-
-            self.policy.set_param_values(best_x)
-            self.policy.log_diagnostics(paths)
-            snapshotter.save_itr_params(
-                itr,
-                dict(
-                    itr=itr,
-                    policy=self.policy,
-                    env=self.env,
-                    cur_mean=cur_mean,
-                    cur_std=cur_std,
-                ))
-            logger.log(tabular)
-            logger.pop_prefix()
-            if self.plot:
-                self.plotter.update_plot(self.policy, self.max_path_length)
-        parallel_sampler.terminate_task()
-        self.plotter.close()
+        # epoch-wise
+        self.cur_std = self.init_std
+        self.cur_mean = self.policy.get_param_values()
+        # epoch-cycle-wise
+        self.cur_params = self.cur_mean
+        self.all_returns = []
+        self.all_params = [self.cur_mean.copy()]
+        # fixed
+        self.n_best = int(n_samples * best_frac)
+        assert self.n_best >= 1, (
+            f"n_samples is too low. Make sure that n_samples * best_frac >= 1")
+        self.n_params = len(self.cur_mean)
+
+    def sample_params(self, epoch):
+        extra_var_mult = max(1.0 - epoch / self.extra_decay_time, 0)
+        sample_std = np.sqrt(
+            np.square(self.cur_std) +
+            np.square(self.extra_std) * extra_var_mult)
+        return np.random.standard_normal(
+            self.n_params) * sample_std + self.cur_mean
+
+    def train_once(self, itr, paths):
+        epoch = itr // self.n_samples
+        i_sample = itr - epoch * self.n_samples
+        tabular.record('Epoch', epoch)
+        tabular.record('# Sample', i_sample)
+        # -- Stage: Process path
+        rtn = paths['average_return']
+        self.all_returns.append(paths['average_return'])
+
+        # -- Stage: Update policy distribution.
+        if (itr + 1) % self.n_samples == 0:
+            avg_rtns = np.array(self.all_returns)
+            best_inds = np.argsort(-avg_rtns)[:self.n_best]
+            best_params = np.array(self.all_params)[best_inds]
+
+            # MLE of normal distribution
+            self.cur_mean = best_params.mean(axis=0)
+            self.cur_std = best_params.std(axis=0)
+            self.policy.set_param_values(self.cur_mean)
+
+            # Clear for next epoch
+            rtn = max(self.all_returns)
+            self.all_returns.clear()
+            self.all_params.clear()
+
+        # -- Stage: Generate a new policy for next path sampling
+        self.cur_params = self.sample_params(itr)
+        self.all_params.append(self.cur_params.copy())
+        self.policy.set_param_values(self.cur_params)
+
+        logger.log(tabular)
+        return rtn
+
+    def get_itr_snapshot(self, itr):
+        return dict(itr=itr, policy=self.policy, baseline=self.baseline)
diff --git a/garage/tf/algos/batch_polopt.py b/garage/tf/algos/batch_polopt.py
--- a/garage/tf/algos/batch_polopt.py
+++ b/garage/tf/algos/batch_polopt.py
@@ -1,5 +1,3 @@
-import tensorflow as tf
-
 from garage.logger import logger
 from garage.np.algos import RLAlgorithm
 
@@ -54,9 +52,6 @@ def __init__(self,
         self.fixed_horizon = fixed_horizon
         self.init_opt()
 
-    def initialize(self, sess=None):
-        self.sess = sess or tf.Session()
-
     def train_once(self, itr, paths):
         self.log_diagnostics(paths)
         logger.log('Optimizing policy...')


Test patch:
diff --git a/tests/garage/np/algos/test_algos.py b/tests/garage/np/algos/test_algos.py
--- a/tests/garage/np/algos/test_algos.py
+++ b/tests/garage/np/algos/test_algos.py
@@ -7,7 +7,6 @@
 from garage.envs import GridWorldEnv
 from garage.envs import PointEnv
 from garage.logger import logger
-from garage.np.algos import CEM
 from garage.np.algos import CMAES
 from garage.np.algos.nop import NOP
 from garage.np.baselines import LinearFeatureBaseline
@@ -22,11 +21,6 @@
 )
 
 algo_args = {
-    CEM: dict(
-        n_itr=1,
-        max_path_length=100,
-        n_samples=5,
-    ),
     CMAES: dict(
         n_itr=1,
         max_path_length=100,
@@ -36,7 +30,7 @@
 }
 
 polopt_cases = []
-for algo in [CEM, CMAES, NOP]:
+for algo in [CMAES, NOP]:
     polopt_cases.extend([
         (algo, GridWorldEnv, DummyPolicy, ZeroBaseline),
         (algo, PointEnv, DummyPolicy, ZeroBaseline),
diff --git a/tests/garage/tf/algos/test_cem.py b/tests/garage/tf/algos/test_cem.py
new file mode 100644
--- /dev/null
+++ b/tests/garage/tf/algos/test_cem.py
@@ -0,0 +1,38 @@
+from garage.experiment import LocalRunner
+from garage.np.algos import CEM
+from garage.np.baselines import LinearFeatureBaseline
+from garage.tf.envs import TfEnv
+from garage.tf.policies import CategoricalMLPPolicy
+from garage.tf.samplers import OnPolicyVectorizedSampler
+from tests.fixtures import TfGraphTestCase
+
+
+class TestCEM(TfGraphTestCase):
+    def test_cem_cartpole(self):
+        """Test CEM with Cartpole-v1 environment."""
+        with LocalRunner() as runner:
+            env = TfEnv(env_name="CartPole-v1")
+
+            policy = CategoricalMLPPolicy(
+                name="policy", env_spec=env.spec, hidden_sizes=(32, 32))
+
+            baseline = LinearFeatureBaseline(env_spec=env.spec)
+
+            runner.initialize_tf_vars()
+
+            n_samples = 10
+
+            algo = CEM(
+                env_spec=env.spec,
+                policy=policy,
+                baseline=baseline,
+                best_frac=0.1,
+                max_path_length=100,
+                n_samples=n_samples)
+
+            runner.setup(algo, env, sampler_cls=OnPolicyVectorizedSampler)
+            rtn = runner.train(
+                n_epochs=5, batch_size=2000, n_epoch_cycles=n_samples)
+            assert rtn > 40
+
+            env.close()