Adding back the missing fully connected layers to the state based WTL model.
Adding back the missing fully connected layers to the state based WTL model.

PiperOrigin-RevId: 254899813

google-research/tensor2robot