Training now use embodiment viz and made the viz more modular #275

This check will pass when downstack PRs merge

#274 needs to be merged into main before you can merge this PR.
- This PR is stacked on top of #274. Merging this PR on GitHub will cause it to merge incorrectly into #274 instead of main.
- Learn more about stacked PRs here
To ensure that these PRs merge correctly, view and merge this PR with Graphite
To disable this check, edit your organization's settings on Graphite.