How to properly train an actor-critic - reinforcement-learning

Question

I'm trying to figure out how the actor-critic method works. And implement it in tensorflow.js
I have designed a small environment where you need a board to catch a falling ball. If caught +1 reward point, if not -1 reward point.
My implementation https://codepen.io/redrickShuhart/pen/popKmRj
The model is implemented with one body and two heads.
I am trying to train the model, but it seems to me that I have implemented the training process incorrectly.
Please help. Where am I wrong.
I would be glad to get a link to a tfjs example to help me figure it out.
train() {
this.memory.reverse()
let V_next = 0
for (let s of this.memory) {
let V_current = s.reward + V_next * this.gamma
s['V'] = V_current
s['V_next'] = V_next
V_next = V_current
}
this.memory.reverse()
for (let mem of this.memory) {
let f = () => tf.tidy(() => {
let {state, action, reward, next_state, V, V_next} = mem
state = tf.tensor(state, [1, this.env.state_size])
let [probs, value] = this.model.predict(state)
let target_V = tf.tensor(V, [1, 1])
let value_loss = tf.losses.meanSquaredError(target_V, value).asScalar()
let advantages = new Array(this.action_size).fill(0)
advantages[action] = V_next - V
advantages = tf.tensor(advantages, [1, this.action_size])
let policy_loss = tf.log(probs).mul(advantages).sum().mul(tf.scalar(-1)).asScalar()
return policy_loss.add(value_loss)
})
let grads = tf.variableGrads(f).grads
this.ALL_GRADIENTS.push(grads)
}
tf.tidy(() => {
/* grads names */
let keys = Object.keys(this.ALL_GRADIENTS[0])
const new_grad = {}
for (let key of keys) {
new_grad[key] = null
}
for (let key in new_grad) {
let grad_by_key = this.ALL_GRADIENTS.map(grad => grad[key])
let mean = grad_by_key.reduce((prev_grad, cur_grad) => prev_grad.add(cur_grad)).div(tf.scalar(grad_by_key.length))
new_grad[key] = mean
tf.dispose(grad_by_key)
}
this.model.optimizer.applyGradients(new_grad)
})
// for (let grad of this.ALL_GRADIENTS) {
// this.model.optimizer.applyGradients(grad)
// }
this.memory = []
tf.dispose(this.ALL_GRADIENTS)
this.ALL_GRADIENTS = []
// console.log(tf.memory())
}