Learning rewards using graph neural networks

[akpas001]

I need to train a Gnn based on the data of pre existing next_states and rewards. next states are one hot encoded and rewards are random numbers(can be negative as well as float). I have constructed a gnn using GINConv.

According to Markov's decision process in Reinforcement learning, a state_action pair as an input to the network gives the next_state and reward. (state, action, next_states are one-hot encoded).

So, i am forming a state action pair inside the network and trying to calculate next state and reward loss. the loss graph for next states is good but the graph for rewards is very bad(rewards are not learning). I donot understand what mistake i am doing here. Can somebody help me?

Here I am attaching the snippets of the code, and graphs of next_state loss and reward loss respectively. Can somebody help me with that?