Change the way reward visualization works

[YouGuessedMyName]

I should account for the fact that rewards are set on states + actions, not just states.