Fix: new observation shape dimension should be converted to int in FlattenObservationsWrapper.
Fix: new observation shape dimension should be converted to int in FlattenObservationsWrapper.

PiperOrigin-RevId: 307853888
Change-Id: Ib45367a53ddcbdecb1668303b7c79221aea940c7

tensorflow/agents,tensorflow/agents