[BugFix] Model maker functions for A2C and PPO fail for discrete action space envs

[albertbou92]

Description

For make_ppo_model and make_a2c_model helper functions, using an environment with a discrete action space raised the following error:

TypeError: distribution keywords and tensordict keys indicated by ProbabilisticTensorDictModule.dist_in_keys must match.Got this error message:
    __init__() got an unexpected keyword argument 'loc'
with dist_in_keys={'loc': 'loc', 'scale': 'scale'}

In this PR I suggest a fix to the problem. Also extended the help_helpers.py script to test discrete environments for both PPO and A2C.

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)
• Example (update in the folder of examples)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[codecov]

Codecov Report

Merging #717 (4e15136) into main (bac229a) will increase coverage by 0.07%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main     #717      +/-   ##
==========================================
+ Coverage   88.66%   88.73%   +0.07%     
==========================================
  Files         120      120              
  Lines       20186    20220      +34     
==========================================
+ Hits        17898    17943      +45     
+ Misses       2288     2277      -11     

Flag	Coverage Δ
habitat-gpu	25.05% <ø> (ø)
linux-cpu	85.67% <100.00%> (+0.07%)	⬆️
linux-gpu	86.63% <100.00%> (+0.09%)	⬆️
linux-jumanji	30.27% <ø> (ø)
linux-outdeps-gpu	72.10% <8.33%> (-0.12%)	⬇️
linux-stable-cpu	85.52% <100.00%> (+0.07%)	⬆️
linux-stable-gpu	86.26% <100.00%> (+0.07%)	⬆️
linux_examples-gpu	43.07% <0.00%> (-0.03%)	⬇️
macos-cpu	85.34% <100.00%> (+0.07%)	⬆️
olddeps-gpu	76.10% <100.00%> (+0.09%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
test/test_helpers.py	92.06% <100.00%> (+0.46%)	⬆️
torchrl/trainers/helpers/models.py	93.38% <100.00%> (+1.88%)	⬆️
torchrl/modules/distributions/discrete.py	64.28% <0.00%> (+2.38%)	⬆️

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[vmoens]

LGTM thanks for this