I am trying to learn a custom environment using the TFAgents package. I am following the Hands-on-ML book (Code in colab see cell 129). My aim is to use DQN agent on a custo