Refactor configs to have env in seperate yaml + Fix training

tests/test_policies.py

@@ -6,12 +6,12 @@ from .utils import init_config
 
 
 @pytest.mark.parametrize(
-    "config_name",
+    "env_name",
     [
-        "default",
+        "simxarm",
         "pusht",
     ],
 )
-def test_factory(config_name):
-    cfg = init_config(config_name)
+def test_factory(env_name):
+    cfg = init_config(overrides=[f"env={env_name}"])
     policy = make_policy(cfg)