Merge remote-tracking branch 'Rofunc-secret/main'

examples/learning_rl/IsaacGym_RofuncRL/example_HumanoidHOTU_RofuncRL.py

@@ -0,0 +1,116 @@
+"""
+HumanoidAMP (RofuncRL)
+===========================
+
+Humanoid backflip/walk/run/dance/hop, trained by RofuncRL
+"""
+
+import isaacgym
+import argparse
+
+from rofunc.config.utils import omegaconf_to_dict, get_config
+from rofunc.learning.RofuncRL.tasks import Tasks
+from rofunc.learning.RofuncRL.trainers import Trainers
+from rofunc.learning.pre_trained_models.download import model_zoo
+from rofunc.learning.utils.utils import set_seed
+
+
+def train(custom_args):
+    # Config task and trainer parameters for Isaac Gym environments
+    args_overrides = ["task={}".format(custom_args.task),
+                      "train=Humanoid{}RofuncRL".format(custom_args.agent.upper()),
+                      "device_id={}".format(custom_args.sim_device),
+                      "rl_device=cuda:{}".format(custom_args.rl_device),
+                      "headless={}".format(custom_args.headless),
+                      "num_envs={}".format(custom_args.num_envs)]
+    cfg = get_config('./learning/rl', 'config', args=args_overrides)
+    cfg.task.env.motion_file = custom_args.motion_file
+    cfg.task.env.object_motion_file = custom_args.object_motion_file
+    cfg.task.env.object_asset.assetFileName = custom_args.object_asset_files
+    cfg.task.env.object_asset.assetName = custom_args.object_asset_names
+    cfg.task.env.object_asset.assetSize = custom_args.object_asset_sizes
+    cfg_dict = omegaconf_to_dict(cfg.task)
+
+    set_seed(cfg.train.Trainer.seed)
+
+    # Instantiate the Isaac Gym environment
+    env = Tasks().task_map[custom_args.task](cfg=cfg_dict,
+                                             rl_device=cfg.rl_device,
+                                             sim_device=f'cuda:{cfg.device_id}',
+                                             graphics_device_id=cfg.device_id,
+                                             headless=cfg.headless,
+                                             virtual_screen_capture=cfg.capture_video,  # TODO: check
+                                             force_render=cfg.force_render)
+
+    # Instantiate the RL trainer
+    trainer = Trainers().trainer_map[custom_args.agent](cfg=cfg,
+                                                        env=env,
+                                                        device=cfg.rl_device,
+                                                        env_name=custom_args.task)
+    # Start training
+    trainer.train()
+
+
+def inference(custom_args):
+    # Config task and trainer parameters for Isaac Gym environments
+    task, motion_file = custom_args.task.split('_')
+    args_overrides = ["task={}".format(task),
+                      "train=Humanoid{}RofuncRL".format(custom_args.agent.upper()),
+                      "device_id={}".format(custom_args.sim_device),
+                      "rl_device=cuda:{}".format(custom_args.rl_device),
+                      "headless={}".format(False),
+                      "num_envs={}".format(16)]
+    cfg = get_config('./learning/rl', 'config', args=args_overrides)
+    cfg_dict = omegaconf_to_dict(cfg.task)
+
+    set_seed(cfg.train.Trainer.seed)
+
+    # Instantiate the Isaac Gym environment
+    infer_env = Tasks().task_map[task](cfg=cfg_dict,
+                                       rl_device=cfg.rl_device,
+                                       sim_device=f'cuda:{cfg.device_id}',
+                                       graphics_device_id=cfg.device_id,
+                                       headless=cfg.headless,
+                                       virtual_screen_capture=cfg.capture_video,  # TODO: check
+                                       force_render=cfg.force_render)
+
+    # Instantiate the RL trainer
+    trainer = Trainers().trainer_map[custom_args.agent](cfg=cfg,
+                                                        env=infer_env,
+                                                        device=cfg.rl_device,
+                                                        env_name=custom_args.task,
+                                                        inference=True)
+
+    # load checkpoint
+    if custom_args.ckpt_path is None:
+        custom_args.ckpt_path = model_zoo(name=f"{custom_args.task}.pth")
+    trainer.agent.load_ckpt(custom_args.ckpt_path)
+
+    # Start inference
+    trainer.inference()
+
+
+if __name__ == '__main__':
+    gpu_id = 1
+
+    parser = argparse.ArgumentParser()
+    # Available tasks: HumanoidHOTU
+    parser.add_argument("--task", type=str, default="HumanoidHOTU")
+    parser.add_argument("--agent", type=str, default="hotu")  # Available agent: amp
+    parser.add_argument("--motion_file", type=str, default="examples/data/hotu2/test_data_01_optitrack2hotu.npy")
+    parser.add_argument("--object_motion_file", type=str, default="examples/data/hotu2/test_data_01_optitrack.csv")
+    parser.add_argument("--object_asset_names", type=str, default=["box"])
+    parser.add_argument("--object_asset_files", type=str, default=["mjcf/objects/lab_box.xml"])
+    parser.add_argument("--object_asset_sizes", type=str, default=[[1, 1, 1]])
+    parser.add_argument("--num_envs", type=int, default=1024)
+    parser.add_argument("--sim_device", type=int, default=1)
+    parser.add_argument("--rl_device", type=int, default=gpu_id)
+    parser.add_argument("--headless", type=str, default="True")
+    parser.add_argument("--inference", action="store_true", help="turn to inference mode while adding this argument")
+    parser.add_argument("--ckpt_path", type=str, default=None)
+    custom_args = parser.parse_args()
+
+    if not custom_args.inference:
+        train(custom_args)
+    else:
+        inference(custom_args)

examples/learning_rl/IsaacGym_RofuncRL/example_HumanoidHOTU_ViewMotion.py

@@ -25,6 +25,7 @@ def inference(custom_args):
     cfg = get_config("./learning/rl", "config", args=args_overrides)
     cfg.task.env.motion_file = custom_args.motion_file
     cfg.task.env.object_motion_file = custom_args.object_motion_file
+    cfg.task.env.object_asset.assetFileName = custom_args.object_asset_files
     cfg.task.env.object_asset.assetName = custom_args.object_asset_names
     cfg.task.env.object_asset.assetSize = custom_args.object_asset_sizes
 
@@ -61,12 +62,14 @@ def inference(custom_args):
     #  1. test data provided by rofunc: `examples/data/hotu/*.npy`
     #  2. custom motion file with absolute path
     parser.add_argument("--motion_file", type=str,
-                        default="examples/data/hotu2/test_data_01_optitrack2hotu.npy")
+                        default="examples/data/hotu2/test_data_04_optitrack2hotu.npy")
     parser.add_argument("--object_motion_file", type=str,
-                        default="examples/data/hotu2/test_data_01_optitrack.csv")
-    parser.add_argument("--object_asset_names", type=str, default=["box:marker 001", "box:marker 002", "box:marker 003", "box:marker 004"])
-    # parser.add_argument("--object_asset_files", type=str, default=["Box.urdf"])
-    parser.add_argument("--object_asset_sizes", type=str, default=[[0.05, 0.05, 0.05], [0.05, 0.05, 0.05], [0.05, 0.05, 0.05], [0.05, 0.05, 0.05]])
+                        default="examples/data/hotu2/test_data_04_optitrack.csv")
+    # parser.add_argument("--object_asset_names", type=str, default=["box:marker 001", "box:marker 002", "box:marker 003", "box:marker 004"])
+    parser.add_argument("--object_asset_names", type=str, default=["box"])
+    parser.add_argument("--object_asset_files", type=str, default=["mjcf/objects/lab_box.xml"])
+    # parser.add_argument("--object_asset_sizes", type=str, default=[[0.05, 0.05, 0.05], [0.05, 0.05, 0.05], [0.05, 0.05, 0.05], [0.05, 0.05, 0.05]])
+    parser.add_argument("--object_asset_sizes", type=str, default=[[1, 1, 1]])
     custom_args = parser.parse_args()
 
     inference(custom_args)

examples/learning_rl/IsaacGym_RofuncRL/example_HumanoidPhysHOI_RofuncRL.py

@@ -17,27 +17,26 @@
 
 def train(custom_args):
     # Config task and trainer parameters for Isaac Gym environments
-    task, motion_file = custom_args.task.split('_')
-    args_overrides = ["task={}".format(task),
-                      "train={}{}RofuncRL".format(task, custom_args.agent.upper()),
+    args_overrides = ["task={}".format(custom_args.task),
+                      "train=BaseTask{}RofuncRL".format(custom_args.agent.upper()),
                       "device_id={}".format(custom_args.sim_device),
                       "rl_device=cuda:{}".format(custom_args.rl_device),
                       "headless={}".format(custom_args.headless),
                       "num_envs={}".format(custom_args.num_envs)]
     cfg = get_config('./learning/rl', 'config', args=args_overrides)
-    cfg.task.env.motion_file = f'amp_humanoid_{motion_file}.npy'
+    cfg.task.env.motion_file = custom_args.motion_file
     cfg_dict = omegaconf_to_dict(cfg.task)
 
     set_seed(cfg.train.Trainer.seed)
 
     # Instantiate the Isaac Gym environment
-    env = Tasks().task_map[task](cfg=cfg_dict,
-                                 rl_device=cfg.rl_device,
-                                 sim_device=f'cuda:{cfg.device_id}',
-                                 graphics_device_id=cfg.device_id,
-                                 headless=cfg.headless,
-                                 virtual_screen_capture=cfg.capture_video,  # TODO: check
-                                 force_render=cfg.force_render)
+    env = Tasks().task_map[custom_args.task](cfg=cfg_dict,
+                                             rl_device=cfg.rl_device,
+                                             sim_device=f'cuda:{cfg.device_id}',
+                                             graphics_device_id=cfg.device_id,
+                                             headless=cfg.headless,
+                                             virtual_screen_capture=cfg.capture_video,  # TODO: check
+                                             force_render=cfg.force_render)
 
     # Instantiate the RL trainer
     trainer = Trainers().trainer_map[custom_args.agent](cfg=cfg,
@@ -91,10 +90,13 @@ def inference(custom_args):
     gpu_id = 0
 
     parser = argparse.ArgumentParser()
-    # Available tasks: HumanoidAMP_backflip, HumanoidAMP_walk, HumanoidAMP_run, HumanoidAMP_dance, HumanoidAMP_hop
-    parser.add_argument("--task", type=str, default="HumanoidAMP_hop")
-    parser.add_argument("--agent", type=str, default="amp")  # Available agent: amp
-    parser.add_argument("--num_envs", type=int, default=4096)
+    # Available tasks: HumanoidPhyshoi
+    parser.add_argument("--task", type=str, default="HumanoidPhysHOI")
+    # Available motion files: backdribble, backspin, changeleg, fingerspin, pass, rebound, toss, walkpick
+    parser.add_argument("--motion_file", type=str, default="examples/data/ballplay/toss.pt")
+    # Available agent: physhoi
+    parser.add_argument("--agent", type=str, default="physhoi")
+    parser.add_argument("--num_envs", type=int, default=2048)
     parser.add_argument("--sim_device", type=int, default=0)
     parser.add_argument("--rl_device", type=int, default=gpu_id)
     parser.add_argument("--headless", type=str, default="True")

examples/learning_rl/IsaacGym_RofuncRL/example_HumanoidPhysHOI_ViewMotion.py

@@ -0,0 +1,62 @@
+"""
+HumanoidPhysHOI Motion View (RofuncRL)
+===========================
+
+Preview the motion of the digital humanoid
+"""
+
+import isaacgym
+import argparse
+
+from rofunc.config.utils import omegaconf_to_dict, get_config
+from rofunc.learning.RofuncRL.tasks import Tasks
+from rofunc.learning.RofuncRL.trainers import Trainers
+
+
+def inference(custom_args):
+    # Config task and trainer parameters for Isaac Gym environments
+    args_overrides = [
+        f"task=HumanoidPhysHOIViewMotion",
+        "train=BaseTaskPHYSHOIRofuncRL",
+        f"device_id=0",
+        f"rl_device=cuda:{gpu_id}",
+        "headless={}".format(False),
+        "num_envs={}".format(1),
+    ]
+    cfg = get_config("./learning/rl", "config", args=args_overrides)
+    cfg.task.env.motion_file = custom_args.motion_file
+    cfg.task.env.playdataset = True
+
+    cfg_dict = omegaconf_to_dict(cfg.task)
+
+    # Instantiate the Isaac Gym environment
+    infer_env = Tasks().task_map[custom_args.task](cfg=cfg_dict,
+                                                   rl_device=cfg.rl_device,
+                                                   sim_device=f'cuda:{cfg.device_id}',
+                                                   graphics_device_id=cfg.device_id,
+                                                   headless=cfg.headless,
+                                                   virtual_screen_capture=cfg.capture_video,  # TODO: check
+                                                   force_render=cfg.force_render)
+
+    # Instantiate the RL trainer
+    trainer = Trainers().trainer_map["physhoi"](cfg=cfg,
+                                                env=infer_env,
+                                                device=cfg.rl_device,
+                                                env_name=custom_args.task,
+                                                inference=True)
+
+    # Start inference
+    trainer.inference()
+
+
+if __name__ == "__main__":
+    gpu_id = 0
+
+    parser = argparse.ArgumentParser()
+    # Available tasks: HumanoidPhyshoi
+    parser.add_argument("--task", type=str, default="HumanoidPhysHOI")
+    # Available motion files: backdribble, backspin, changeleg, fingerspin, pass, rebound, toss, walkpick
+    parser.add_argument("--motion_file", type=str, default="examples/data/ballplay/backdribble.pt")
+    custom_args = parser.parse_args()
+
+    inference(custom_args)

examples/simulator/example_LLM_robotics.py

@@ -10,12 +10,12 @@
 
 import rofunc as rf
 
-traj_l = np.load('../data/LQT_LQR/taichi_1l.npy')
-traj_l[:, 0] += 0.2
-traj_l[:, 1] = -traj_l[:, 1]
-traj_l[:, 3:] = [1, 0, 0, 0]
-# target_pose = traj_l[10]
-traj_l = traj_l[::20]
+# traj_l = np.load('../data/LQT_LQR/taichi_1l.npy')
+# traj_l[:, 0] += 0.2
+# traj_l[:, 1] = -traj_l[:, 1]
+# traj_l[:, 3:] = [1, 0, 0, 0]
+# # target_pose = traj_l[10]
+# traj_l = traj_l[::20]
 
 args = rf.config.get_sim_config("CURI")
 CURIsim = rf.sim.CURISim(args)