TODO:

• Create single pyproject.toml
• Refactor configs
• Parallelise 2d dataset
• Simplify unet
• Fix value unet
• Make classifier policy
• Add data viz script
• Simplify command manager
• Add IK solver
• Simplify classifier policy
• Make classifier mlp
• Batch classifier training
• Simplify classifier runner
• Fix jit tracing
• simplify normalizer
• Remove RL code
• Refactor env cfgs
• Add IK goals and init
• Tune performance
• Make vis class