pytorch version

baselines/nets.py

@@ -0,0 +1,95 @@
+import tensorflow as tf
+import baselines.baselines_common.tf_util as U
+from baselines.baselines_common.mpi_running_mean_std import RunningMeanStd
+from baselines.baselines_common.distributions import make_pdtype, DiagGaussianPdType, BernoulliPdType
+
+
+def mlp_block(x, name, num_hid_layers, hid_size, activation_fn=tf.nn.tanh):
+    with tf.variable_scope(name_or_scope=name):
+        for i in range(num_hid_layers):
+            x = U.dense(
+                x, hid_size,
+                name="fc%i" % (i + 1), weight_init=U.normc_initializer(1.0))
+            x = activation_fn(x)
+        return x
+
+
+def feature_net(x, name, num_hid_layers, hid_size, activation_fn=tf.nn.tanh):
+    with tf.variable_scope(name_or_scope=name):
+        x = mlp_block(
+            x, name="mlp",
+            hid_size=hid_size, num_hid_layers=num_hid_layers, activation_fn=activation_fn)
+        return x
+
+
+class Actor(object):
+    def __init__(self, name, *args, **kwargs):
+        with tf.variable_scope(name):
+            self._init(*args, **kwargs)
+            self.scope = tf.get_variable_scope().name
+
+    def _init(self, ob_space, ac_space, hid_size, num_hid_layers, gaussian_fixed_var=True, noise_type=None):
+        if noise_type == "gaussian":
+            self.pdtype = pdtype = DiagGaussianPdType(ac_space.shape[0])
+        else:
+            self.pdtype = pdtype = make_pdtype(ac_space)
+
+        ob = U.get_placeholder(
+            name="ob", dtype=tf.float32,
+            shape=[None] + list(ob_space.shape))
+
+        with tf.variable_scope("obfilter"):
+            self.ob_rms = RunningMeanStd(shape=ob_space.shape)
+        obz = (ob - self.ob_rms.mean) / self.ob_rms.std
+        obz = tf.clip_by_value(obz, -5.0, 5.0)
+
+        # critic net (value network)
+        last_out = feature_net(
+            obz, name="vf",
+            num_hid_layers=num_hid_layers, hid_size=hid_size,
+            activation_fn=tf.nn.tanh)
+        self.vpred = U.dense(
+            last_out, 1,
+            name="vf_final", weight_init=U.normc_initializer(1.0))[:, 0]
+
+        # actor net (policy network)
+        last_out = feature_net(
+            obz, name="pol",
+            num_hid_layers=num_hid_layers, hid_size=hid_size,
+            activation_fn=tf.nn.tanh)
+
+        if gaussian_fixed_var and isinstance(self.pdtype, DiagGaussianPdType):
+            mean = U.dense(
+                last_out, pdtype.param_shape()[0] // 2,
+                name="pol_final", weight_init=U.normc_initializer(0.01))
+            logstd = tf.get_variable(
+                name="logstd", shape=[1, pdtype.param_shape()[0] // 2],
+                initializer=tf.zeros_initializer())
+            pdparam = U.concatenate([mean, mean * 0.0 + logstd], axis=1)
+        else:
+            pdparam = U.dense(
+                last_out, pdtype.param_shape()[0],
+                name="pol_final", weight_init=U.normc_initializer(0.01))
+
+        # pd - probability distribution
+        self.pd = pdtype.pdfromflat(pdparam)
+
+        self.state_in = []
+        self.state_out = []
+
+        stochastic = tf.placeholder(dtype=tf.bool, shape=())
+        ac = U.switch(stochastic, self.pd.sample(), self.pd.mode())
+        self._act = U.function([stochastic, ob], [ac, self.vpred])
+
+    def act(self, stochastic, ob):
+        ac1, vpred1 = self._act(stochastic, ob[None])
+        return ac1[0], vpred1[0]
+
+    def get_variables(self):
+        return tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, self.scope)
+
+    def get_trainable_variables(self):
+        return tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.scope)
+
+    def get_initial_state(self):
+        return []