Update state checking logic in PPO policy to work with autograph for stateful networks.
Update state checking logic in PPO policy to work with autograph for stateful networks.

PiperOrigin-RevId: 411136020
Change-Id: I8aa838f30e22f137835ee4e6e9c5291e806b8485

tensorflow/agents,tensorflow/agents