fix last batch index error reinforce

[sidhantls]

What does this PR do?

This PR replaces getting batch size from self when calculating loss with using the length of the tensor, fixing an IndexError in Reinforce

The error was an IndexError (size mismatch) occurring on the last batch of the epoch because tensor size wasn't equal to self.batch_size. The number of samples in the dataset may not be divisible by self.batch_size, because trajectory collection here is done based on the number of episodes rather than the number of samples. Rather than fixing the error using drop_last=True in dataloader, this approach uses all trajectory data

Fixes #381

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you make sure your PR does only one thing, instead of bundling different changes together? Otherwise, we ask you to create a separate PR for every change.
• Did you make sure to update the documentation with your changes? No changes
• Did you write any new necessary tests?

Did you have fun?

Make sure you had fun coding 🙃

[codecov]

Codecov Report

Merging #389 (4e8ad0c) into master (fa5a944) will increase coverage by 0.06%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master     #389      +/-   ##
==========================================
+ Coverage   81.11%   81.18%   +0.06%     
==========================================
  Files         100      100              
  Lines        5693     5714      +21     
==========================================
+ Hits         4618     4639      +21     
  Misses       1075     1075              

Flag	Coverage Δ
cpu	24.27% <0.00%> (ø)
pytest	24.27% <0.00%> (ø)
unittests	80.46% <100.00%> (+0.07%)	⬆️

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

Impacted Files	Coverage Δ
pl_bolts/models/rl/reinforce_model.py	90.08% <100.00%> (ø)
pl_bolts/datasets/base_dataset.py	95.45% <0.00%> (-4.55%)	⬇️
pl_bolts/utils/arguments.py	96.15% <0.00%> (+0.04%)	⬆️
.../models/autoencoders/basic_vae/basic_vae_module.py	91.57% <0.00%> (+0.08%)	⬆️
...ts/models/autoencoders/basic_ae/basic_ae_module.py	87.83% <0.00%> (+0.16%)	⬆️
pl_bolts/losses/self_supervised_learning.py	71.33% <0.00%> (+0.18%)	⬆️
pl_bolts/models/rl/common/gym_wrappers.py	89.91% <0.00%> (+0.26%)	⬆️
...lts/models/self_supervised/simclr/simclr_module.py	71.49% <0.00%> (+0.27%)	⬆️
pl_bolts/datasets/imagenet_dataset.py	18.45% <0.00%> (+0.38%)	⬆️
pl_bolts/optimizers/lars_scheduling.py	95.74% <0.00%> (+0.39%)	⬆️
... and 3 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update fa5a944...4e8ad0c. Read the comment docs.

[akihironitta]

@sid-sundrani Thank you for reporting and investigating the bug!
Seems you've completed this PR... Make sure to mark this PR ready for review when it's ready :]

[akihironitta]

@sid-sundrani LGTM! Thank you for your contribution and swift action :]

[Borda]

may we add also a test for this issue you have described?

[sidhantls]

@Borda Hmm okay. Considering this issue, we'd want a test to check if there is an error when current batch size isn't equal to given batch size when calculating the loss? If we set num_batch_episodes=1 and batch_size=1000 (batch_size larger than the number of samples in the first episode), reinforce would only run for one episode and the size of its only training batch won't be equal to the given batch size. Thus it would check if the loss function is compatible with batch size !=self.batch_size.

If we set those parameters then would running something like this in TestReinforce work?

def test_loss_epoch_end(self):
    """Test the reinforce loss function with dataloader"""
    for batch_states, batch_actions, batch_qvals in  self.rl_dataloader:
        loss = self.model.loss(batch_states, batch_actions, batch_qvals)

[Borda]

@sid-sundrani mind send a PR with adding test so we can talk about it there...
cc: @tchaton

[sidhantls]

@Borda Sure. #402