Bug in state and action prediction

[ezhang7423]

Although decision transformer does not predict states or returns, if one were to add these losses the following prediction appears to be broken as the predict_return and predict_state projections do not take in the state, but only the action.

decision-transformer/gym/decision_transformer/models/decision_transformer.py

Line 97 in e2d82e6

return_preds = self.predict_return(x[:,2]) # predict next return given state and action