Bug Fix for IM Loss for Action Spaces > 1

[Zach-Attach]

Bug fix for adjusting the size of the mask to match the size of the im loss, which is not 1D when action space is greater than 1.

Description

In rewards that use im_loss (ICM, E3B, RIDE, and pseudo counts), I repeated the same mask value by the action dim. Since the im_loss and fm_loss share the same mask anyway, I chose to keep the same mask when adding values.

Motivation and Context

closes #52

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

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

Checklist

• I've read the CONTRIBUTION guide (required)
• I have updated the changelog accordingly (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.
• I have opened an associated PR on the rllte-hub repository (if necessary)
• I have reformatted the code using make format (required)
• I have checked the codestyle using make check-codestyle and make lint (required)
• I have ensured make pytest and make type both pass. (required)
• I have checked that the documentation builds using make doc (required)

Note: You can run most of the checks using make commit-checks.

Note: we are using a maximum length of 127 characters per line