Hotfix - Return the new updated key in function _train

[theovincent]

Description

This PR addresses the issue #45

Motivation and Context

closes #45

• 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)

Checklist:

• I've read the CONTRIBUTION guide (required)
• I have updated the changelog accordingly (required).
  I do not think they exist in SBX. If I am wrong, could you please tell me how to modify them?
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
  I do not think I need to add extra tests. I ran the test with make pytest and they are all passing.
• I have updated the documentation accordingly.
• 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)

>>> make commit-checks
# Sort imports
ruff check --select I sbx/ tests/ setup.py --fix
All checks passed!
# Reformat using black
black sbx/ tests/ setup.py
All done! ✨ 🍰 ✨
36 files left unchanged.
mypy sbx/ tests/ setup.py
Success: no issues found in 36 source files
# stop the build if there are Python syntax errors or undefined names
# see https://www.flake8rules.com/
ruff check sbx/ tests/ setup.py --select=E9,F63,F7,F82 --output-format=full
All checks passed!
# exit-zero treats all errors as warnings.
ruff check sbx/ tests/ setup.py --exit-zero
All checks passed!

[araffin]

I've launched several runs of CrossQ on https://wandb.ai/openrlbenchmark/sbx
will compare them to https://wandb.ai/openrlbenchmark/sbx/reports/CrossQ-SBX-Perf-Report--Vmlldzo3MzQxOTAw

[araffin]

LGTM, thanks =)

[araffin]

I didn't notice any significant difference in performance (that's also explain why we didn't realize earlier), there was probably enough randomness from the sampled transitions (this was still correct, same for collecting new transitions).