created a two-hot encoding function

[DavideTr8]

Summary

Describe the purpose of the pull request, including:

A function for converting tensors representing floating point to a two-hot representation, as done in MuZero and Dreamer V3 papers.

Type of Change

Please select the one relevant option below:

• New feature (non-breaking change that adds functionality)

Checklist

Please confirm that the following tasks have been completed:

• I have tested my changes locally and they work as expected. (Please describe the tests you performed.)
• I have added unit tests for my changes, or updated existing tests if necessary.
• I have updated the documentation, if applicable.
• I have installed pre-commit and run locally for my code changes.

[belerico]

I think that the implementation in the Dreamer-V3 paper is more general than this one (https://arxiv.org/abs/2301.04104). Could we adjust it?