🚀 [RofuncRL] Update

examples/learning_rl/IsaacGym_RofuncRL/example_HumanoidPhysHOI_RofuncRL.py

@@ -0,0 +1,108 @@
+"""
+HumanoidPhysHOI (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
+    task, motion_file = custom_args.task.split('_')
+    args_overrides = ["task={}".format(task),
+                      "train={}{}RofuncRL".format(task, 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_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)
+
+    # 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={}{}RofuncRL".format(task, 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 = 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)
+    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")
+    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)

rofunc/config/learning/rl/task/HumanoidPhysHOI.yaml

@@ -0,0 +1,68 @@
+# if given, will override the device setting in gym. 
+env: 
+  numEnvs: 2048
+  envSpacing: 5
+  episodeLength: 40
+  isFlagrun: False
+  enableDebugVis: False
+  playdataset: False
+  projtype: "None"
+  saveImages: False
+  initVel: False
+
+  pdControl: True
+  powerScale: 1.0
+  controlFrequencyInv: 2 # 30 Hz
+  stateInit: "Start" #Random
+  hybridInitProb: 0.5
+  dataFPS: 25
+  dataFramesScale: 1.2 # 25->30fps
+  ballSize: 1.
+
+  localRootObs: False
+  keyBodies: ["Head", "L_Knee", "R_Knee", "L_Elbow", "R_Elbow", "L_Ankle", "R_Ankle", "L_Index3", "L_Middle3", "L_Pinky3", "L_Ring3","L_Thumb3","R_Index3", "R_Middle3", "R_Pinky3", "R_Ring3","R_Thumb3"] # #["L_Hip", "L_Knee", "left_foot", "L_Toe", "R_Hip", "R_Knee", "right_foot", "R_Toe", "Torso", "Spine", "Chest", "Neck", "Head", "L_Thorax", "L_Shoulder", "L_Elbow", "left_hand", "R_Thorax", "R_Shoulder", "R_Elbow", "right_hand"] #["right_hand", "left_hand", "right_foot", "left_foot"]
+  contactBodies: ["L_Index3", "L_Middle3", "L_Pinky3", "L_Ring3","L_Thumb3","R_Index3", "R_Middle3", "R_Pinky3", "R_Ring3","R_Thumb3"] #["right_foot", "left_foot"]
+  terminationHeight: 0.15
+  enableEarlyTermination: True
+
+  asset:
+    assetRoot: "physhoi/data/assets"
+    assetFileName: "smplx/smplx_capsule.xml"
+
+  plane:
+    staticFriction: 1.0
+    dynamicFriction: 1.0
+    restitution: 1.6
+
+  rewardWeights:
+    p: 50.
+    r: 50.
+    pv: 0.
+    rv: 0.
+
+    op: 1.
+    or: 0.
+    opv: 0.
+    orv: 0.
+
+    ig: 20.
+
+    cg1: 5.
+    cg2: 5.
+
+sim:
+  substeps: 2
+  physx:
+    num_threads: 4
+    solver_type: 1  # 0: pgs, 1: tgs
+    num_position_iterations: 4
+    num_velocity_iterations: 0
+    contact_offset: 0.02
+    rest_offset: 0.0
+    bounce_threshold_velocity: 0.2
+    max_depenetration_velocity: 10.0
+    default_buffer_size_multiplier: 10.0
+
+  flex:
+    num_inner_iterations: 10
+    warm_start: 0.25

rofunc/config/learning/rl/task/HumanoidPhyshoi60hz.yaml

@@ -0,0 +1,68 @@
+# if given, will override the device setting in gym. 
+env: 
+  numEnvs: 2048
+  envSpacing: 5
+  episodeLength: 40
+  isFlagrun: False
+  enableDebugVis: False
+  playdataset: False
+  projtype: "None"
+  saveImages: False
+  initVel: False
+
+  pdControl: True
+  powerScale: 1.0
+  controlFrequencyInv: 1 # 60 Hz
+  stateInit: "Start" #Random
+  hybridInitProb: 0.5
+  dataFPS: 25
+  dataFramesScale: 2.4 # 25->60fps
+  ballSize: 1.
+
+  localRootObs: False
+  keyBodies: ["Head", "L_Knee", "R_Knee", "L_Elbow", "R_Elbow", "L_Ankle", "R_Ankle", "L_Index3", "L_Middle3", "L_Pinky3", "L_Ring3","L_Thumb3","R_Index3", "R_Middle3", "R_Pinky3", "R_Ring3","R_Thumb3"] # #["L_Hip", "L_Knee", "left_foot", "L_Toe", "R_Hip", "R_Knee", "right_foot", "R_Toe", "Torso", "Spine", "Chest", "Neck", "Head", "L_Thorax", "L_Shoulder", "L_Elbow", "left_hand", "R_Thorax", "R_Shoulder", "R_Elbow", "right_hand"] #["right_hand", "left_hand", "right_foot", "left_foot"]
+  contactBodies: ["L_Index3", "L_Middle3", "L_Pinky3", "L_Ring3","L_Thumb3","R_Index3", "R_Middle3", "R_Pinky3", "R_Ring3","R_Thumb3"] #["right_foot", "left_foot"]
+  terminationHeight: 0.15
+  enableEarlyTermination: True
+
+  asset:
+    assetRoot: "physhoi/data/assets"
+    assetFileName: "smplx/smplx_capsule.xml"
+
+  plane:
+    staticFriction: 1.0
+    dynamicFriction: 1.0
+    restitution: 1.6
+
+  rewardWeights:
+    p: 50.
+    r: 50.
+    pv: 0.
+    rv: 0.
+
+    op: 1.
+    or: 0.
+    opv: 0.
+    orv: 0.
+
+    ig: 20.
+
+    cg1: 5.
+    cg2: 1.
+
+sim:
+  substeps: 2
+  physx:
+    num_threads: 4
+    solver_type: 1  # 0: pgs, 1: tgs
+    num_position_iterations: 4
+    num_velocity_iterations: 0
+    contact_offset: 0.02
+    rest_offset: 0.0
+    bounce_threshold_velocity: 0.2
+    max_depenetration_velocity: 10.0
+    default_buffer_size_multiplier: 10.0
+
+  flex:
+    num_inner_iterations: 10
+    warm_start: 0.25

rofunc/learning/RofuncRL/agents/mixline/physhoi_agent.py

@@ -0,0 +1,32 @@
+# Copyright 2023, Junjia LIU, [email protected]
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import gym
+import gymnasium
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from omegaconf import DictConfig
+from typing import Callable, Union, Tuple, Optional
+
+import rofunc as rf
+from rofunc.learning.RofuncRL.agents.base_agent import BaseAgent
+from rofunc.learning.RofuncRL.agents.mixline.amp_agent import AMPAgent
+from rofunc.learning.RofuncRL.models.base_models import BaseMLP
+from rofunc.learning.RofuncRL.utils.memory import Memory
+
+
+class PhysHOIAgent(AMPAgent):
+    def __init__(self):
+        super().__init__()