DDPG training in Tensorflow.js

Question

I'm trying to figure out how to implement DDPG in Tensorflow.js using Python examples such as this one from keras website. I got stuck on the training code:

        with tf.GradientTape() as tape:
            target_actions = target_actor(next_state_batch, training=True)
            y = reward_batch + gamma * target_critic(
                [next_state_batch, target_actions], training=True
            )
            critic_value = critic_model([state_batch, action_batch], training=True)
            critic_loss = tf.math.reduce_mean(tf.math.square(y - critic_value))

        critic_grad = tape.gradient(critic_loss, critic_model.trainable_variables)
        critic_optimizer.apply_gradients(
            zip(critic_grad, critic_model.trainable_variables)
        )

        with tf.GradientTape() as tape:
            actions = actor_model(state_batch, training=True)
            critic_value = critic_model([state_batch, actions], training=True)
            # Used `-value` as we want to maximize the value given
            # by the critic for our actions
            actor_loss = -tf.math.reduce_mean(critic_value)

        actor_grad = tape.gradient(actor_loss, actor_model.trainable_variables)
        actor_optimizer.apply_gradients(
            zip(actor_grad, actor_model.trainable_variables)
        )

So far my typescript version looks like this:

    const batch = this.memory.getMinibatch(this.config.replayBatchSize);
    const states = this.actorService.getStateTensor(this.actor, ...batch.map(s => s.state));
    const nextStates = this.actorService.getStateTensor(this.actor, ...batch.map(s => s.nextState));
    const rewards = tf.tensor2d(batch.map(s => s.reward), [batch.length, 1], 'float32');
    const actions = this.actorService.getActionTensor(...batch.map(s => s.action));

    const criticLossFunction = () => tf.tidy(() => {
      let targetQs: tf.Tensor;
      if (this.config.discountRate === 0) {
        targetQs = rewards;
      } else {
        const targetActions = this.targetActorModel.predict(nextStates) as tf.Tensor;
        const targetCriticQs = this.targetCriticModel.predict(tf.concat([nextStates, targetActions], 1)) as tf.Tensor;
        targetQs = rewards.add(targetCriticQs.mul(this.config.discountRate));
      }
      const criticQs = this.criticModel.predict(tf.concat([states, actions], 1)) as tf.Tensor;
      const criticLoss = tf.losses.meanSquaredError(targetQs, criticQs);
      return criticLoss.asScalar();
    });
    const criticTrainableVars = this.criticModel.getWeights(true) as tf.Variable<tf.Rank>[];
    const criticGradient = tf.variableGrads(criticLossFunction, criticTrainableVars);
    // HOWTO:  zip(critic_grad, critic_model.trainable_variables)
    this.criticModel.optimizer.applyGradients(criticGradient.grads);
    tf.dispose(criticGradient);

    const actorLossFunction = () => tf.tidy(() => {
      const policyActions = this.actorModel.predict(states) as tf.Tensor;
      const criticQs = this.criticModel.predict(tf.concat([states, policyActions], 1)) as tf.Tensor;
      const actorLoss = tf.mean(criticQs.mul(-1));
      return actorLoss.asScalar()
    });
    const actorTrainableVars = this.actorModel.getWeights(true) as tf.Variable<tf.Rank>[];
    const actorGradient = tf.variableGrads(actorLossFunction, actorTrainableVars);
    // HOWTO: zip(actor_grad, actor_model.trainable_variables)
    this.actorModel.optimizer.applyGradients(actorGradient.grads);
    const actorLoss = actorGradient.value.dataSync()[0];
    tf.dispose(actorGradient);

but my code does not work correctly (the loss is too high on a very simple task) due to it's missing one major step (?): zipping critic grads with critic trainable vars before passing it to applyGradients.

Answer 1

False alarm, the code works correctly, no need to zip or anything. High loss was due to I used actorLoss as the value instead of criticLoss (criticGradient.value.dataSync()[0])