[RLlib] Hot fix for PPOTorchRLModule._compute_values with non-shared stateful encoder and batch slicing with non-empty infos.

[simonsays1980]

Why are these changes needed?

Running PPO with use_lstm=True and vf_share_layers=False results in an error in the PPOTorchRLModule._compute_values method as the specs checker expects a different spec for the state_in:

raise SpecCheckingError(
ray.rllib.core.models.specs.checker.SpecCheckingError: input spec validation failed on TorchLSTMEncoder.forward, The data dict does not match the model specs. Keys ('state_in', 'h') are in the spec dict but not on the data dict. Data keys are {('state_in', 'critic', 'c'), ('rewards',), ('action_logp',), ('state_in', 'actor', 'h'), ('state_in', 'critic', 'h'), ('terminateds',), ('obs',), ('truncateds',), ('vf_preds',), ('loss_mask',), ('seq_lens',), ('action_dist_inputs',), ('infos',), ('actions',), ('state_in', 'actor', 'c')}.

Exctracting the state_in for the critic solves this problem.

Another problem is solved related to non-empty infos in batch slicing (mainly occuring in MinibatchIterators). The reason is that slicing via tree.map_structure tries to slice also the entries of the infos which are usually singular values:

data = tree.map_structure(lambda value: value[start:stop], self)
IndexError: slice() cannot be applied to a 0-dim tensor.

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[sven1977]

Makes sense!

[sven1977]

You mean a SampleBatch with B=0, correct?

[simonsays1980]

B>0. But in this case the infos are a list of dicts. When they are empty, tree.map_structure(infos) works, but when they are filled, tree.map_structure will fail as it tries to apply the slicing on singular values slicing fails.

[sven1977]

Looks good, thanks for the fix @simonsays1980!

[sven1977]

Awesome! Ran into this issue yesterday as well (and continued testing then with a shared value function :) ).