-
Notifications
You must be signed in to change notification settings - Fork 50
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
8c07767
commit 3031a95
Showing
91 changed files
with
14,072 additions
and
487 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
105 changes: 105 additions & 0 deletions
105
examples/learning_rl/IsaacLab_RofuncRL/example_isaaclab_env.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,105 @@ | ||
| """ | ||
| Ant (RofuncRL) | ||
| =========================== | ||
| Ant RL using RofuncRL | ||
| """ | ||
|
|
||
| import sys,os | ||
|
|
||
| sys.path.append("/home/ubuntu/Github/Rofunc") | ||
|
|
||
| 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 | ||
| from rofunc.learning.RofuncRL.tasks.utils.env_loaders import load_isaaclab_env | ||
| from rofunc.learning.utils.env_wrappers import wrap_env | ||
|
|
||
|
|
||
| def train(custom_args): | ||
| # Config task and trainer parameters for Isaac Gym environments | ||
| custom_args.num_envs = 64 if custom_args.agent.upper() in ["SAC", "TD3"] else custom_args.num_envs | ||
|
|
||
| args_overrides = ["task={}".format(custom_args.task), | ||
| "train={}{}RofuncRL".format(custom_args.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_dict = omegaconf_to_dict(cfg.task) | ||
|
|
||
| set_seed(cfg.train.Trainer.seed) | ||
|
|
||
| # Instantiate the Isaac Gym environment | ||
| env = load_isaaclab_env(task_name="Isaac-Cartpole-v0", headless=True, num_envs=custom_args.num_envs) | ||
|
|
||
| # 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 | ||
| args_overrides = ["task={}".format(custom_args.task), | ||
| "train={}{}RofuncRL".format(custom_args.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[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=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="AntRofuncRLPPO.pth") | ||
| trainer.agent.load_ckpt(custom_args.ckpt_path) | ||
|
|
||
| # Start inference | ||
| trainer.inference() | ||
|
|
||
|
|
||
| if __name__ == '__main__': | ||
| gpu_id = 0 | ||
|
|
||
| parser = argparse.ArgumentParser() | ||
| parser.add_argument("--task", type=str, default="Cartpole") | ||
| parser.add_argument("--agent", type=str, default="ppo") # Available agents: ppo, sac, td3, a2c | ||
| 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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,72 @@ | ||
| """ | ||
| FD from models | ||
| ======================== | ||
| Forward dynamics from URDF or MuJoCo XML files. | ||
| """ | ||
|
|
||
| import pprint | ||
| import math | ||
|
|
||
| import rofunc as rf | ||
|
|
||
| rf.logger.beauty_print("########## Forward kinematics from URDF or MuJoCo XML files with RobotModel class ##########") | ||
| rf.logger.beauty_print("---------- Forward kinematics for Franka Panda using URDF file ----------") | ||
| model_path = "../../rofunc/simulator/assets/urdf/franka_description/robots/franka_panda.urdf" | ||
|
|
||
| joint_value = [[0.0, -math.pi / 4.0, 0.0, math.pi / 2.0, 0.0, math.pi / 4.0, 0.0, 0.0, 0.0], | ||
| [0.0, -math.pi / 4.0, 0.0, math.pi / 2.0, 0.0, math.pi / 4.0, 0.0, 0.0, 0.0]] | ||
| export_link = "panda_hand" | ||
|
|
||
| # # Build the robot model with kinpy | ||
| # # Deprecated: kinpy is not supported anymore, just for checking the results!!!! Please use pytorch_kinematics instead. | ||
| # robot = rf.robolab.RobotModel(model_path, solve_engine="kinpy", verbose=False) | ||
| # # Show the robot chain and joint names, can also be done by verbose=True | ||
| # robot.show_chain() | ||
| # # Get the forward kinematics of export_link | ||
| # pos, rot, ret = robot.get_fk(joint_value, export_link) | ||
| # | ||
| # # Convert the orientation representation and print the results | ||
| # rot = rf.robolab.convert_quat_order(rot, "wxyz", "xyzw") | ||
| # rf.logger.beauty_print(f"Position of {export_link}: {pos}") | ||
| # rf.logger.beauty_print(f"Rotation of {export_link}: {rot}") | ||
| # pprint.pprint(ret, width=1) | ||
|
|
||
| # Try the same thing with pytorch_kinematics | ||
| robot = rf.robolab.RobotModel(model_path, solve_engine="pytorch_kinematics", verbose=False) | ||
| pos, rot, ret = robot.get_fk(joint_value, export_link) | ||
| rf.logger.beauty_print(f"Position of {export_link}: {pos}") | ||
| rf.logger.beauty_print(f"Rotation of {export_link}: {rot}") | ||
| pprint.pprint(ret, width=1) | ||
|
|
||
| # rf.logger.beauty_print("---------- Forward kinematics for Bruce Humanoid Robot using MJCF file ----------") | ||
| model_path = "../../rofunc/simulator/assets/mjcf/bruce/bruce.xml" | ||
| joint_value = [0.0 for _ in range(16)] | ||
|
|
||
| export_link = "elbow_pitch_link_r" | ||
|
|
||
| # # Build the robot model with pytorch_kinematics, kinpy is not supported for MJCF files | ||
| robot = rf.robolab.RobotModel(model_path, solve_engine="pytorch_kinematics", verbose=True) | ||
| # Get the forward kinematics of export_link | ||
| pos, rot, ret = robot.get_fk(joint_value, export_link) | ||
| # | ||
| # # Print the results | ||
| # rf.logger.beauty_print(f"Position of {export_link}: {pos}") | ||
| # rf.logger.beauty_print(f"Rotation of {export_link}: {rot}") | ||
| # pprint.pprint(ret, width=1) | ||
|
|
||
|
|
||
| model_path = "../../rofunc/simulator/assets/mjcf/hotu/hotu_humanoid.xml" | ||
| joint_value = [0.1 for _ in range(34)] | ||
|
|
||
| export_link = "left_hand_link_2" | ||
|
|
||
| # # Build the robot model with pytorch_kinematics, kinpy is not supported for MJCF files | ||
| robot = rf.robolab.RobotModel(model_path, solve_engine="pytorch_kinematics", verbose=True) | ||
| # Get the forward kinematics of export_link | ||
| pos, rot, ret = robot.get_fk(joint_value, export_link) | ||
|
|
||
| # # Print the results | ||
| rf.logger.beauty_print(f"Position of {export_link}: {pos}") | ||
| rf.logger.beauty_print(f"Rotation of {export_link}: {rot}") | ||
| pprint.pprint(ret, width=1) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.