Predict output is NaN on RL Card Game DQN - reinforcement-learning

Question

I'm trying to make a reinforcement learning project for a popular portuguese card game. I have the environment working. The code can run on it's own, for n rounds, using random picks.
For a quick explanation, the game is similar to Hearts. It has 4 players in 2 teams. Each players has 10 cards, playing following suit, when possible, until no cards left. Each trick earns points, after 10 tricks the team with more than 60 points wins round.
First, i had some doubts on how to "encode" the cards/deck. I'm passing to the model the cards on table, cards on hand and already played cards. I one-hot encoded the deck, 4 suits, 10 numbers/figures. I also have to include the trump (taken from the begining of the round). So as an example 1 card looks like: [ 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0 ], where the first four numbers are suit, and the last is trump (boolean) and the rest is what card (2, 3, 4, 5, 6, j, q, k, 7, a). I'm passing the deck for each hand, table, played, so 40*3, 15 features each (1800 features when flattened). Possibly this is not the best approach, if someone could advise on this would be great.
With this data, i'm setting 1 player as the AI agent. When it's AI turn, state is cards in hand, cards on table, cards played. For next state i'm passing the trick end, with the same type of data. Reward from that trick is a cumulative value (max 120 points). When i run the code, the loss outputs a value initialy, like the first time, then displays NaN. So predictions are coming out as NaN too (i have 10 outputs, 1 for each card in hand. Also having some doubts on this one because they start with 10 cards, but as the game goes and the cards get played, this number goes to zero.)
Here's the code for the training:
async expReplay() {
console.debug('Training...')
const minibatch = await this.memory.concat().sort(() => .5 - Math.random()).slice(0, this.batchSize)
for (let i = 0; i < minibatch.length - 1; i++) {
let [state, action, reward, next_state, done] = minibatch[i]
state = await tf.concat(state).flatten().reshape([1, this.stateSize])
next_state = await tf.concat(next_state).flatten().reshape([1, this.stateSize])
let target = reward
if (!done) {
let predictNext = await this.model.predict(next_state)
target = reward + this.gamma * predictNext.argMax().dataSync()[0]
}
let target_f = await this.model.predict(state).dataSync()
target_f[action] = target
target_f = await tf.tensor2d(target_f, [1, this.actionSize])
await this.model.fit(state, target_f, {
epochs: 1,
verbose: 1,
callbacks: {
onEpochEnd: (epoch, logs) => {
process.stdout.write(`${logs.loss} ${logs.acc} \r`)
}
}
})
await state.dispose()
await next_state.dispose()
await target_f.dispose()
}
if (this.epsilon > this.epsilonMin) {
this.epsilon *= this.epsilonDecay
}
return 'Training... stop!'
}
This loop, i've used it before, on a DQN also for a Bitcoin trader i was experimenting and it worked fine. So i'm guessing my data is wrong somewhere. I've logged the state and next_state to check for NaN but didn't catch any...
If you need more info please ask!