[BugFix] Batch size for data generation

[aziabatz]

Description

This PR addresses an issue that appears when generating the data for an environment. The problem was that the whole dataset size was used for training, validation, and testing, instead of the specified batch sizes.

Motivation

Using total dataset sizes led memory allocation problems during environment setup. By switching to batch sizes, we assure that the data we handle has the right size.

Types of Changes

• [x] Fix: Change data generation to use batch sizes.

Checklist

• [x] No new linting errors introduced.
• [x] Verified all changes are compatible with existing codebase.

[fedebotu]

Please let us know if you find some better way of describing variable, right now the behavior is correct but naming might be confusing @aziabatz

[cbhua]

Thanks for reporting your concern @aziabatz. I think the variable name of batch_size in rl4co/envs/common/base.py is misleading. This batch_size actually is the problem size for the whole train/valid phase. We could revise this variable if necessary.