adding LSTM support to pretrain

stable_baselines/a2c/a2c.py

@@ -87,9 +87,20 @@ def __init__(self, policy, env, gamma=0.99, n_steps=5, vf_coef=0.25, ent_coef=0.
 
     def _get_pretrain_placeholders(self):
         policy = self.train_model
+
+        if self.initial_state is None:
+            states_ph = None
+            snew_ph = None
+            dones_ph = None
+        else:
+            states_ph = policy.states_ph
+            snew_ph = policy.snew
+            dones_ph = policy.dones_ph
+
         if isinstance(self.action_space, gym.spaces.Discrete):
-            return policy.obs_ph, self.actions_ph, policy.policy
-        return policy.obs_ph, self.actions_ph, policy.deterministic_action
+            return policy.obs_ph, self.actions_ph, states_ph, snew_ph, dones_ph, policy.policy
+        return policy.obs_ph, self.actions_ph, states_ph, snew_ph, dones_ph,\
+               policy.deterministic_action
 
     def setup_model(self):
         with SetVerbosity(self.verbose):

stable_baselines/acer/acer_simple.py

@@ -152,8 +152,18 @@ def __init__(self, policy, env, gamma=0.99, n_steps=20, num_procs=1, q_coef=0.5,
     def _get_pretrain_placeholders(self):
         policy = self.step_model
         action_ph = policy.pdtype.sample_placeholder([None])
+
+        if self.initial_state is None:
+            states_ph = None
+            snew_ph = None
+            dones_ph = None
+        else:
+            states_ph = policy.states_ph
+            snew_ph = policy.snew
+            dones_ph = policy.dones_ph
+
         if isinstance(self.action_space, Discrete):
-            return policy.obs_ph, action_ph, policy.policy
+            return policy.obs_ph, action_ph, states_ph, snew_ph, dones_ph, policy.policy
         raise NotImplementedError('Only discrete actions are supported for ACER for now')
 
     def set_env(self, env):

stable_baselines/acktr/acktr_disc.py

@@ -104,8 +104,18 @@ def __init__(self, policy, env, gamma=0.99, nprocs=1, n_steps=20, ent_coef=0.01,
 
     def _get_pretrain_placeholders(self):
         policy = self.train_model
+
+        if self.initial_state is None:
+            states_ph = None
+            snew_ph = None
+            dones_ph = None
+        else:
+            states_ph = policy.states_ph
+            snew_ph = policy.snew
+            dones_ph = policy.dones_ph
+
         if isinstance(self.action_space, Discrete):
-            return policy.obs_ph, self.action_ph, policy.policy
+            return policy.obs_ph, self.action_ph, states_ph, snew_ph, dones_ph, policy.policy
         raise NotImplementedError("WIP: ACKTR does not support Continuous actions yet.")
 
     def setup_model(self):

stable_baselines/common/base_class.py

@@ -50,6 +50,10 @@ def __init__(self, policy, env, verbose=0, *, requires_vec_env, policy_base, pol
         self.sess = None
         self.params = None
         self._param_load_ops = None
+        self.initial_state = None
+        self.n_batch = None
+        self.nminibatches = None
+        self.n_steps = None
 
         if env is not None:
             if isinstance(env, str):
@@ -246,13 +250,24 @@ def pretrain(self, dataset, n_epochs=10, learning_rate=1e-4,
             else:
                 val_interval = int(n_epochs / 10)
 
+        use_lstm = self.initial_state is not None
+
+        if use_lstm:
+            if self.nminibatches is None:
+                envs_per_batch = self.n_envs * self.n_steps
+            else:
+                batch_size = self.n_batch // self.nminibatches
+                envs_per_batch = batch_size // self.n_steps
+
         with self.graph.as_default():
             with tf.variable_scope('pretrain'):
                 if continuous_actions:
-                    obs_ph, actions_ph, deterministic_actions_ph = self._get_pretrain_placeholders()
+                    obs_ph, actions_ph, states_ph, snew_ph, dones_ph, \
+                    deterministic_actions_ph = self._get_pretrain_placeholders()
                     loss = tf.reduce_mean(tf.square(actions_ph - deterministic_actions_ph))
                 else:
-                    obs_ph, actions_ph, actions_logits_ph = self._get_pretrain_placeholders()
+                    obs_ph, actions_ph, states_ph, snew_ph, dones_ph, \
+                    actions_logits_ph = self._get_pretrain_placeholders()
                     # actions_ph has a shape if (n_batch,), we reshape it to (n_batch, 1)
                     # so no additional changes is needed in the dataloader
                     actions_ph = tf.expand_dims(actions_ph, axis=1)
@@ -272,13 +287,23 @@ def pretrain(self, dataset, n_epochs=10, learning_rate=1e-4,
 
         for epoch_idx in range(int(n_epochs)):
             train_loss = 0.0
+            if use_lstm:
+                state = self.initial_state[:envs_per_batch]
+
             # Full pass on the training set
             for _ in range(len(dataset.train_loader)):
-                expert_obs, expert_actions = dataset.get_next_batch('train')
+                expert_obs, expert_actions, expert_mask = dataset.get_next_batch('train')
                 feed_dict = {
                     obs_ph: expert_obs,
                     actions_ph: expert_actions,
                 }
+
+                if use_lstm:
+                    feed_dict.update({states_ph: state, dones_ph: expert_mask})
+                    state, train_loss_, _ = self.sess.run([snew_ph, loss, optim_op], feed_dict)
+                else:
+                    train_loss_, _ = self.sess.run([loss, optim_op], feed_dict)
+
                 train_loss_, _ = self.sess.run([loss, optim_op], feed_dict)
                 train_loss += train_loss_
 
@@ -288,9 +313,19 @@ def pretrain(self, dataset, n_epochs=10, learning_rate=1e-4,
                 val_loss = 0.0
                 # Full pass on the validation set
                 for _ in range(len(dataset.val_loader)):
-                    expert_obs, expert_actions = dataset.get_next_batch('val')
-                    val_loss_, = self.sess.run([loss], {obs_ph: expert_obs,
-                                                        actions_ph: expert_actions})
+                    expert_obs, expert_actions, expert_mask = dataset.get_next_batch('val')
+
+                    feed_dict = {
+                        obs_ph: expert_obs,
+                        actions_ph: expert_actions,
+                    }
+
+                    if use_lstm:
+                        feed_dict.update({states_ph: state, dones_ph: expert_mask})
+                        val_loss_, = self.sess.run([loss], feed_dict)
+                    else:
+                        val_loss_, = self.sess.run([loss], feed_dict)
+
                     val_loss += val_loss_
 
                 val_loss /= len(dataset.val_loader)

stable_baselines/ddpg/ddpg.py

@@ -308,7 +308,7 @@ def _get_pretrain_placeholders(self):
         policy = self.policy_tf
         # Rescale
         deterministic_action = self.actor_tf * np.abs(self.action_space.low)
-        return policy.obs_ph, self.actions, deterministic_action
+        return policy.obs_ph, self.actions, None, None, None, deterministic_action
 
     def setup_model(self):
         with SetVerbosity(self.verbose):

stable_baselines/deepq/dqn.py

@@ -99,7 +99,7 @@ def __init__(self, policy, env, gamma=0.99, learning_rate=5e-4, buffer_size=5000
 
     def _get_pretrain_placeholders(self):
         policy = self.step_model
-        return policy.obs_ph, tf.placeholder(tf.int32, [None]), policy.q_values
+        return policy.obs_ph, tf.placeholder(tf.int32, [None]), None, None, None, policy.q_values
 
     def setup_model(self):
 

stable_baselines/gail/__init__.py

@@ -1,3 +1,3 @@
 from stable_baselines.gail.model import GAIL
-from stable_baselines.gail.dataset.dataset import ExpertDataset, DataLoader
+from stable_baselines.gail.dataset.dataset import ExpertDataset, ExpertDatasetLSTM, DataLoader
 from stable_baselines.gail.dataset.record_expert import generate_expert_traj