feat: refactor dtensor v2 policy __init__ and introduce core types

[hemildesai]

What does this PR do ?

Refactors init of Dtensor policy v2. Depends on #1470 and #1556

Issues

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[github-actions]

⚠️ File Consistency Check

Check based on commit: 124860d (PR #1588 from hemil/dtensor-v2-types-init)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

• Please review if the changes in nemo_rl/models/policy/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/dtensor_policy_worker.py
• Update nemo_rl/models/policy/dtensor_policy_worker.py if necessary to maintain consistency
• If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

• Modified: nemo_rl/models/policy/dtensor_policy_worker_v2.py
• Not modified: nemo_rl/models/policy/dtensor_policy_worker.py

---

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

[github-actions]

❌ Submodule Fast-Forward Check Failed

Check based on commit: 124860d (PR #1588 from hemil/dtensor-v2-types-init)

❌ Submodules that need attention:

Automodel: ❌ Commits have DIVERGED from a common ancestor
TARGET (main branch): https://github.com/NVIDIA-NeMo/Automodel/commits/a2db048383cd54b3fafc928df4c30bf7bbf7c430/
CURRENT (PR #1588 from hemil/dtensor-v2-types-init): https://github.com/NVIDIA-NeMo/Automodel/commits/2d20e33a19d5e53a271b1403b507475e68ad14dc/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

[github-actions]

⚠️ File Consistency Check

Check based on commit: ffefeee (PR #1588 from hemil/dtensor-v2-types-init)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

• Please review if the changes in nemo_rl/models/policy/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/dtensor_policy_worker.py
• Update nemo_rl/models/policy/dtensor_policy_worker.py if necessary to maintain consistency
• If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

• Modified: nemo_rl/models/policy/dtensor_policy_worker_v2.py
• Not modified: nemo_rl/models/policy/dtensor_policy_worker.py

---

_{This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.}

[github-actions]

❌ Submodule Fast-Forward Check Failed

Check based on commit: ffefeee (PR #1588 from hemil/dtensor-v2-types-init)

❌ Submodules that need attention:

Automodel: ❌ Commits have DIVERGED from a common ancestor
TARGET (main branch): https://github.com/NVIDIA-NeMo/Automodel/commits/a2db048383cd54b3fafc928df4c30bf7bbf7c430/
CURRENT (PR #1588 from hemil/dtensor-v2-types-init): https://github.com/NVIDIA-NeMo/Automodel/commits/2d20e33a19d5e53a271b1403b507475e68ad14dc/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

[hemildesai]

Closing in favor of #1709