title: Soft actor critic with discrete action space
body: <p>I'm trying to implement the soft actor critic algorithm for discrete action space and I have trouble with the loss function.</p>
  <p>Here is the link from SAC with continuous action space:
  <a href="https://spinningup.openai.com/en/latest/algorithms/sac.html" rel="nofollow noreferrer">https://spinningup.openai.com/en/latest/algorithms/sac.html</a></p>
  <p>I do not know what I'm doing wrong.</p>
  <p>The problem is the network do not learn anything on the cartpole environment.</p>
  <p>The full code on github: <a href="https://github.com/tk2232/sac_discrete/blob/master/sac_discrete.py" rel="nofollow noreferrer">https://github.com/tk2232/sac_discrete/blob/master/sac_discrete.py</a></p>
  <p>Here is my idea how to calculate the loss for discrete actions.</p>
  <h2>Value Network</h2>
  <pre class="lang-py prettyprint-override"><code>class ValueNet:
      def __init__(self, sess, state_size, hidden_dim, name):
          self.sess = sess
  
          with tf.variable_scope(name):
              self.states = tf.placeholder(dtype=tf.float32, shape=[None, state_size], name='value_states')
              self.targets = tf.placeholder(dtype=tf.float32, shape=[None, 1], name='value_targets')
              x = Dense(units=hidden_dim, activation='relu')(self.states)
              x = Dense(units=hidden_dim, activation='relu')(x)
              self.values = Dense(units=1, activation=None)(x)
  
              optimizer = tf.train.AdamOptimizer(0.001)
  
              loss = 0.5 * tf.reduce_mean((self.values - tf.stop_gradient(self.targets)) ** 2)
              self.train_op = optimizer.minimize(loss, var_list=_params(name))
  
      def get_value(self, s):
          return self.sess.run(self.values, feed_dict={self.states: s})
  
      def update(self, s, targets):
          self.sess.run(self.train_op, feed_dict={self.states: s, self.targets: targets})
  </code></pre>
  <p>In the Q_Network I'm gather the values with the collected actions</p>
  <h2>Example</h2>
  <pre><code>q_out = [[0.5533, 0.4444], [0.2222, 0.6666]]
  collected_actions = [0, 1]
  gather = [[0.5533], [0.6666]]
  </code></pre>
  <h2>gather function</h2>
  <pre class="lang-py prettyprint-override"><code>def gather_tensor(params, idx):
      idx = tf.stack([tf.range(tf.shape(idx)[0]), idx[:, 0]], axis=-1)
      params = tf.gather_nd(params, idx)
      return params
  </code></pre>
  <h2>Q Network</h2>
  <pre class="lang-py prettyprint-override"><code>class SoftQNetwork:
      def __init__(self, sess, state_size, action_size, hidden_dim, name):
          self.sess = sess
  
          with tf.variable_scope(name):
              self.states = tf.placeholder(dtype=tf.float32, shape=[None, state_size], name='q_states')
              self.targets = tf.placeholder(dtype=tf.float32, shape=[None, 1], name='q_targets')
              self.actions = tf.placeholder(dtype=tf.int32, shape=[None, 1], name='q_actions')
  
              x = Dense(units=hidden_dim, activation='relu')(self.states)
              x = Dense(units=hidden_dim, activation='relu')(x)
              x = Dense(units=action_size, activation=None)(x)
              self.q = tf.reshape(gather_tensor(x, self.actions), shape=(-1, 1))
  
              optimizer = tf.train.AdamOptimizer(0.001)
  
              loss = 0.5 * tf.reduce_mean((self.q - tf.stop_gradient(self.targets)) ** 2)
              self.train_op = optimizer.minimize(loss, var_list=_params(name))
  
      def update(self, s, a, target):
          self.sess.run(self.train_op, feed_dict={self.states: s, self.actions: a, self.targets: target})
  
      def get_q(self, s, a):
          return self.sess.run(self.q, feed_dict={self.states: s, self.actions: a})
  </code></pre>
  <h2>Policy Net</h2>
  <pre class="lang-py prettyprint-override"><code>class PolicyNet:
      def __init__(self, sess, state_size, action_size, hidden_dim):
          self.sess = sess
  
          with tf.variable_scope('policy_net'):
              self.states = tf.placeholder(dtype=tf.float32, shape=[None, state_size], name='policy_states')
              self.targets = tf.placeholder(dtype=tf.float32, shape=[None, 1], name='policy_targets')
              self.actions = tf.placeholder(dtype=tf.int32, shape=[None, 1], name='policy_actions')
  
              x = Dense(units=hidden_dim, activation='relu')(self.states)
              x = Dense(units=hidden_dim, activation='relu')(x)
              self.logits = Dense(units=action_size, activation=None)(x)
              dist = Categorical(logits=self.logits)
  
              optimizer = tf.train.AdamOptimizer(0.001)
  
              # Get action
              self.new_action = dist.sample()
              self.new_log_prob = dist.log_prob(self.new_action)
  
              # Calc loss
              log_prob = dist.log_prob(tf.squeeze(self.actions))
              loss = tf.reduce_mean(tf.squeeze(self.targets) - 0.2 * log_prob)
              self.train_op = optimizer.minimize(loss, var_list=_params('policy_net'))
  
      def get_action(self, s):
          action = self.sess.run(self.new_action, feed_dict={self.states: s[np.newaxis, :]})
          return action[0]
  
      def get_next_action(self, s):
          next_action, next_log_prob = self.sess.run([self.new_action, self.new_log_prob], feed_dict={self.states: s})
          return next_action.reshape((-1, 1)), next_log_prob.reshape((-1, 1))
  
      def update(self, s, a, target):
          self.sess.run(self.train_op, feed_dict={self.states: s, self.actions: a, self.targets: target})
  </code></pre>
  <h2>Update function</h2>
  <pre class="lang-py prettyprint-override"><code>def soft_q_update(batch_size, frame_idx):
      gamma = 0.99
      alpha = 0.2
  
      state, action, reward, next_state, done = replay_buffer.sample(batch_size)
      action = action.reshape((-1, 1))
      reward = reward.reshape((-1, 1))
      done = done.reshape((-1, 1))
  </code></pre>
  <h3>Q_target</h3>
  <p><img src="https://latex.codecogs.com/gif.latex?r&space;&plus;&space;%5Cgamma&space;(1-d)V_%7B%5Cpsi&space;targ%7D(s%7B%7D%27)" title="r + \gamma (1-d)V_{\psi targ}(s{}')" /></p>
  <pre class="lang-py prettyprint-override"><code>v_ = value_net_target.get_value(next_state)
  q_target = reward + (1 - done) * gamma * v_
  </code></pre>
  <h3>V_target</h3>
  <p><img src="https://latex.codecogs.com/gif.latex?%5Cmin_%7Bi=1,2%7D&space;Q_%7B%5Cphi_i%7D(s,%5Ctilde%7Ba%7D)&space;-&space;%5Calpha&space;%5Clog&space;%5Cpi_%7B%5Ctheta%7D(%5Ctilde%7Ba%7D|s)" title="\min_{i=1,2} Q_{\phi_i}(s,\tilde{a}) - \alpha \log \pi_{\theta}(\tilde{a}|s)" /></p>
  <pre class="lang-py prettyprint-override"><code>next_action, next_log_prob = policy_net.get_next_action(state)
  q1 = soft_q_net_1.get_q(state, next_action)
  q2 = soft_q_net_2.get_q(state, next_action)
  q = np.minimum(q1, q2)
  v_target = q - alpha * next_log_prob
  
  </code></pre>
  <h3>Policy_target</h3>
  <p><img src="https://latex.codecogs.com/gif.latex?Q%5E%7B%5Cpi&space;%7D(s,a))" title="Q^{\pi }(s,a))" /></p>
  <pre class="lang-py prettyprint-override"><code>q1 = soft_q_net_1.get_q(state, action)
  q2 = soft_q_net_2.get_q(state, action)
  policy_target = np.minimum(q1, q2)
  </code></pre>