[RLlib|New API|Inconsistency] LSTM Encoder lacks the output Linear, but stated in the docstring (#47625)

[brieyla1]

I don't know who to add for review

Why are these changes needed?

Fixes an inconsistency between the docstring and the implementation of the TorchLSTMEncoder class in the RLlib new API.

See Issue for more info

Related issue number

Closes #47625

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 (Not sure, ran all tests from python -m pytest -v -s python/ray/tests/test_basic.py like state in https://docs.ray.io/en/latest/ray-contribute/getting-involved.html#testing and didn't run any C++ Tests are none have been changed.)

[simonsays1980]

@brieyla1 Thanks a lot for your PR and the great catch! Because we are rethinking the design of encoders and their configs to make implementation of modules easier for users, we are reluctant to changes in the encoders themselves at this moment and opt for a change of all docstrings.

This PR is great as it is in its code implementation.

[brieyla1]

Interesting, I wonder what's the impact on the action processing of directly using an LSTM connected to a head, some of that knowledge may be making the training easier when the VF layers are not shared because of the extra linear processing of the data & their relationships before splitting it.

At the same time, I understand the reluctance to modify the Catalog & Encoder, as models built with it may have issues.. e.g. when reinitializing the RLModule and migrating the weights.

We should do some benchmarks to weigh the benefits of such an "improvement" since we already know the downsides.

[simonsays1980]

Interesting, I wonder what's the impact on the action processing of directly using an LSTM connected to a head, some of that knowledge may be making the training easier when the VF layers are not shared because of the extra linear processing of the data & their relationships before splitting it.

At the same time, I understand the reluctance to modify the Catalog & Encoder, as models built with it may have issues.. e.g. when reinitializing the RLModule and migrating the weights.

We should do some benchmarks to weigh the benefits of such an "improvement" since we already know the downsides.

Benchmarking can be done already without adding the linear layer to the encoder but to the head(s): using post_fcnet_hiddens adds layers to the head network(s). The more complex your environment is the more layers will be needed to succeed in the task(s). PPO learns Atari Pong with just a single fc-layer after the CNN encoder (should take around 8 minutes on 4 GPUs and 95 EnvRunners). In this case layers are shared between the policy and the critic (just make sure that you tune the vf_loss_coeff).