def test_fit(self, obs_dim):
box_env = TfEnv(DummyBoxEnv(obs_dim=obs_dim))
with mock.patch(('garage.tf.baselines.'
'continuous_mlp_baseline_with_model.'
'ContinuousMLPRegressorWithModel'),
new=SimpleMLPRegressor):
cmb = ContinuousMLPBaselineWithModel(env_spec=box_env.spec)
paths = [{
'observations': [np.full(obs_dim, 1)],
'returns': [1]
}, {
'observations': [np.full(obs_dim, 2)],
'returns': [2]
}]
cmb.fit(paths)
obs = {'observations': [np.full(obs_dim, 1), np.full(obs_dim, 2)]}
prediction = cmb.predict(obs)
assert np.array_equal(prediction, [1, 2])
def test_is_pickleable(self):
box_env = TfEnv(DummyBoxEnv(obs_dim=(1, )))
with mock.patch(('garage.tf.baselines.'
'continuous_mlp_baseline_with_model.'
'ContinuousMLPRegressorWithModel'),
new=SimpleMLPRegressor):
cmb = ContinuousMLPBaselineWithModel(env_spec=box_env.spec)
obs = {'observations': [np.full(1, 1), np.full(1, 1)]}
with tf.variable_scope('ContinuousMLPBaselineWithModel', reuse=True):
return_var = tf.get_variable('SimpleMLPModel/return_var')
return_var.load(1.0)
prediction = cmb.predict(obs)
h = pickle.dumps(cmb)
with tf.Session(graph=tf.Graph()):
cmb_pickled = pickle.loads(h)
prediction2 = cmb_pickled.predict(obs)
assert np.array_equal(prediction, prediction2)