Feat/appo g1 motion tracking

conf/appo/config.yaml

@@ -58,6 +58,10 @@ training:
   no_play: false
   load_run: "-1"
   play_env_num: 16
+  play_steps: 150
+  cam_distance: 6.0
+  cam_elevation: -20.0
+  cam_azimuth: 90.0
   replay_queue_size: null
 
 hydra:

conf/appo/reward/g1_motion_tracking.yaml

@@ -0,0 +1,21 @@
+# @package _global_
+reward:
+  scales:
+    motion_global_root_pos: 0.5
+    motion_global_root_ori: 0.5
+    motion_body_pos: 1.0
+    motion_body_ori: 1.0
+    motion_body_lin_vel: 1.0
+    motion_body_ang_vel: 1.0
+    motion_joint_pos: 0.0
+    motion_joint_vel: 0.0
+    action_rate_l2: -0.1
+    joint_limit: -10.0
+  std_root_pos: 0.3
+  std_root_ori: 0.4
+  std_body_pos: 0.3
+  std_body_ori: 0.4
+  std_body_lin_vel: 1.0
+  std_body_ang_vel: 3.14
+  std_joint_pos: 0.2
+  std_joint_vel: 1.0

conf/appo/task/g1_motion_tracking.yaml

@@ -0,0 +1,12 @@
+# @package _global_
+defaults:
+  - override /reward: g1_motion_tracking
+
+training:
+  task_name: G1MotionTracking
+  play_steps: 1000
+
+algo:
+  num_envs: 1024
+  max_iterations: 5000
+  save_interval: 500

scripts/train_appo.py

@@ -18,7 +18,11 @@
 
 
 def ensure_registries():
-    for pkg_name in ("unilab.envs.locomotion", "unilab.envs.manipulation"):
+    for pkg_name in (
+        "unilab.envs.locomotion",
+        "unilab.envs.manipulation",
+        "unilab.envs.motion_tracking",
+    ):
         try:
             package = importlib.import_module(pkg_name)
             if hasattr(package, "__path__"):
@@ -147,13 +151,11 @@ def play_appo(cfg: DictConfig, rl_cfg: dict):
     if env.state is None:
         env.init_state()
     env_indices = np.arange(cfg.training.play_env_num, dtype=np.int32)
-    from unilab.utils.obs_utils import flatten_obs_dict
-
     obs_out, _ = env.reset(env_indices)
     obs_np = np.asarray(obs_out["obs"], dtype=np.float32)
 
     state_list = []
-    num_steps = 150
+    num_steps = int(getattr(cfg.training, "play_steps", 1000))
 
     print("Collecting physics states...")
     with torch.inference_mode():
@@ -168,7 +170,14 @@ def play_appo(cfg: DictConfig, rl_cfg: dict):
 
     print("Rendering frames...")
     frames = render_many.render_states_get_frames(
-        state_list, env.cfg.model_file, width=1280, height=720, camera_id=-1
+        state_list,
+        env.cfg.model_file,
+        width=1280,
+        height=720,
+        camera_id=-1,
+        cam_distance=cfg.training.cam_distance,
+        cam_elevation=cfg.training.cam_elevation,
+        cam_azimuth=cfg.training.cam_azimuth,
     )
 
     print(f"Saving video to {output_video} with mediapy...")

src/unilab/utils/algo_utils.py

@@ -6,7 +6,11 @@
 
 def ensure_registries():
     """Import all env modules so they are registered."""
-    for pkg_name in ("unilab.envs.locomotion", "unilab.envs.manipulation"):
+    for pkg_name in (
+        "unilab.envs.locomotion",
+        "unilab.envs.manipulation",
+        "unilab.envs.motion_tracking",
+    ):
         try:
             package = importlib.import_module(pkg_name)
             if hasattr(package, "__path__"):

tests/scripts/test_train_script_configs.py

@@ -16,7 +16,7 @@
 @pytest.mark.veryslow
 @pytest.mark.parametrize(
     "task",
-    ["go1_joystick", "go2_joystick", "g1_joystick"],
+    ["go1_joystick", "go2_joystick", "g1_joystick", "g1_motion_tracking"],
 )
 def test_appo_task_configs_load(task):
     """APPO can start training with all task configs."""