Skip to content

Commit

Permalink
🚀 [RofuncRL] Update OmniIsaacGymWrapper
Browse files Browse the repository at this point in the history
  • Loading branch information
@Skylark0924
Skylark0924 committed Oct 24, 2023
1 parent fb506df commit fb1fbde
Show file tree
Hide file tree
Showing 14 changed files with 276 additions and 255 deletions.
12 changes: 4 additions & 8 deletions examples/learning_rl/example_AntOmni_RofuncRL.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,6 @@
from rofunc.learning.RofuncRL.tasks import Tasks
from rofunc.learning.pre_trained_models.download import model_zoo
from rofunc.learning.utils.utils import set_seed
from omniisaacgymenvs.envs.vec_env_rlgames import VecEnvRLGames


def train(custom_args):
Expand All @@ -25,22 +24,19 @@ def train(custom_args):
"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)
cfg_dict = omegaconf_to_dict(cfg)

set_seed(cfg.train.Trainer.seed)

# Startup IsaacSim simulator
env_omni = VecEnvRLGames(headless=cfg.headless, sim_device=cfg.device_id)
from omni.isaac.gym.vec_env import VecEnvBase
env_omni = VecEnvBase(headless=cfg.headless, sim_device=cfg.device_id)

# Instantiate the Isaac Gym environment
from omniisaacgymenvs.utils.config_utils.sim_config import SimConfig
sim_config = SimConfig(cfg_dict)
env = Tasks("omniisaacgym").task_map[custom_args.task](name=cfg.task_name,
sim_config=sim_config,
env=env_omni)
init_sim = True
env_omni.set_task(task=env, sim_params=sim_config.get_physics_params(), backend="torch", init_sim=init_sim)
env = Tasks("omniisaacgym").task_map[custom_args.task](name=cfg.task_name, sim_config=sim_config, env=env_omni)
env_omni.set_task(task=env, sim_params=sim_config.get_physics_params(), backend="torch", init_sim=True)

# Instantiate the RL trainer
trainer = Trainers().trainer_map[custom_args.agent](cfg=cfg.train,
Expand Down
48 changes: 24 additions & 24 deletions examples/learning_rl/example_Ant_RofuncRL.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@
import argparse

from rofunc.config.utils import omegaconf_to_dict, get_config
from rofunc.learning.RofuncRL.tasks import task_map
from rofunc.learning.RofuncRL.trainers import trainer_map
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

Expand All @@ -31,19 +31,19 @@ def train(custom_args):
set_seed(cfg.train.Trainer.seed)

# Instantiate the Isaac Gym environment
env = task_map[custom_args.task](cfg=cfg_dict,
rl_device=cfg.rl_device,
sim_device=cfg.sim_device,
graphics_device_id=cfg.graphics_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=cfg.sim_device,
graphics_device_id=cfg.graphics_device_id,
headless=cfg.headless,
virtual_screen_capture=cfg.capture_video, # TODO: check
force_render=cfg.force_render)

# Instantiate the RL trainer
trainer = trainer_map[custom_args.agent](cfg=cfg.train,
env=env,
device=cfg.rl_device,
env_name=custom_args.task)
trainer = Trainers().trainer_map[custom_args.agent](cfg=cfg.train,
env=env,
device=cfg.rl_device,
env_name=custom_args.task)

# Start training
trainer.train()
Expand All @@ -64,19 +64,19 @@ def inference(custom_args):
set_seed(cfg.train.Trainer.seed)

# Instantiate the Isaac Gym environment
infer_env = task_map[custom_args.task](cfg=cfg_dict,
rl_device=cfg.rl_device,
sim_device=cfg.sim_device,
graphics_device_id=cfg.graphics_device_id,
headless=cfg.headless,
virtual_screen_capture=cfg.capture_video, # TODO: check
force_render=cfg.force_render)
infer_env = Tasks().task_map[custom_args.task](cfg=cfg_dict,
rl_device=cfg.rl_device,
sim_device=cfg.sim_device,
graphics_device_id=cfg.graphics_device_id,
headless=cfg.headless,
virtual_screen_capture=cfg.capture_video, # TODO: check
force_render=cfg.force_render)

# Instantiate the RL trainer
trainer = trainer_map[custom_args.agent](cfg=cfg.train,
env=infer_env,
device=cfg.rl_device,
env_name=custom_args.task)
trainer = Trainers().trainer_map[custom_args.agent](cfg=cfg.train,
env=infer_env,
device=cfg.rl_device,
env_name=custom_args.task)
# load checkpoint
if custom_args.ckpt_path is None:
custom_args.ckpt_path = model_zoo(name="AntRofuncRLPPO.pth")
Expand Down
48 changes: 24 additions & 24 deletions examples/learning_rl/example_CURICabinetImage_RofuncRL.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@
import argparse

from rofunc.config.utils import omegaconf_to_dict, get_config
from rofunc.learning.RofuncRL.tasks import task_map
from rofunc.learning.RofuncRL.trainers import trainer_map
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

Expand All @@ -31,19 +31,19 @@ def train(custom_args):
set_seed(cfg.train.Trainer.seed)

# Instantiate the Isaac Gym environment
env = task_map[custom_args.task](cfg=cfg_dict,
rl_device=cfg.rl_device,
sim_device=cfg.sim_device,
graphics_device_id=cfg.graphics_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=cfg.sim_device,
graphics_device_id=cfg.graphics_device_id,
headless=cfg.headless,
virtual_screen_capture=cfg.capture_video, # TODO: check
force_render=cfg.force_render)

# Instantiate the RL trainer
trainer = trainer_map[custom_args.agent](cfg=cfg.train,
env=env,
device=cfg.rl_device,
env_name=custom_args.task)
trainer = Trainers().trainer_map[custom_args.agent](cfg=cfg.train,
env=env,
device=cfg.rl_device,
env_name=custom_args.task)

# Start training
trainer.train()
Expand All @@ -64,19 +64,19 @@ def inference(custom_args):
set_seed(cfg.train.Trainer.seed)

# Instantiate the Isaac Gym environment
infer_env = task_map[custom_args.task](cfg=cfg_dict,
rl_device=cfg.rl_device,
sim_device=cfg.sim_device,
graphics_device_id=cfg.graphics_device_id,
headless=cfg.headless,
virtual_screen_capture=cfg.capture_video, # TODO: check
force_render=cfg.force_render)
infer_env = Tasks().task_map[custom_args.task](cfg=cfg_dict,
rl_device=cfg.rl_device,
sim_device=cfg.sim_device,
graphics_device_id=cfg.graphics_device_id,
headless=cfg.headless,
virtual_screen_capture=cfg.capture_video, # TODO: check
force_render=cfg.force_render)

# Instantiate the RL trainer
trainer = trainer_map[custom_args.agent](cfg=cfg.train,
env=infer_env,
device=cfg.rl_device,
env_name=custom_args.task)
trainer = Trainers().trainer_map[custom_args.agent](cfg=cfg.train,
env=infer_env,
device=cfg.rl_device,
env_name=custom_args.task)
# load checkpoint
if custom_args.ckpt_path is None:
custom_args.ckpt_path = model_zoo(name="CURICabinetRofuncRLPPO_left_arm.pth") # TODO: Need to update the ckpt
Expand Down
23 changes: 10 additions & 13 deletions examples/learning_rl/example_CURICabinet_RLlib.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,19 +11,15 @@
# If you use ray==2.2.0, you need to comment out line 23 in ray.rllib.agents.sac.__init__.py

import argparse
import sys

import isaacgym

import os
import sys

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

from rofunc.config.utils import get_config, omegaconf_to_dict
from rofunc.learning.pre_trained_models import model_zoo
from rofunc.utils.logger.beauty_logger import beauty_print
from rofunc.learning.RofuncRL.utils.rllib_utils import RLlibIsaacGymVecEnvWrapper
from rofunc.learning.RofuncRL.tasks import task_map
from rofunc.learning.RofuncRL.tasks import Tasks

from hydra._internal.utils import get_args_parser
from tqdm.auto import tqdm
Expand Down Expand Up @@ -100,13 +96,14 @@ def test(custom_args, ckpt_path):
beauty_print("Start testing")

agent, env_config = setup(custom_args, eval_mode=True)
env = task_map[env_config["task_name"]](cfg=env_config["task_cfg_dict"],
rl_device=env_config["cfg"].rl_device,
sim_device=env_config["cfg"].sim_device,
graphics_device_id=env_config["cfg"].graphics_device_id,
headless=env_config["cfg"].headless,
virtual_screen_capture=env_config["cfg"].capture_video, # TODO: check
force_render=env_config["cfg"].force_render)
env = Tasks().task_map[env_config["task_name"]](cfg=env_config["task_cfg_dict"],
rl_device=env_config["cfg"].rl_device,
sim_device=env_config["cfg"].sim_device,
graphics_device_id=env_config["cfg"].graphics_device_id,
headless=env_config["cfg"].headless,
virtual_screen_capture=env_config["cfg"].capture_video,
# TODO: check
force_render=env_config["cfg"].force_render)
agent.restore(ckpt_path)

done = True
Expand Down
48 changes: 24 additions & 24 deletions examples/learning_rl/example_CURICabinet_RofuncRL.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,8 @@
import argparse

from rofunc.config.utils import omegaconf_to_dict, get_config
from rofunc.learning.RofuncRL.tasks import task_map
from rofunc.learning.RofuncRL.trainers import trainer_map
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

Expand All @@ -31,19 +31,19 @@ def train(custom_args):
set_seed(cfg.train.Trainer.seed)

# Instantiate the Isaac Gym environment
env = task_map[custom_args.task](cfg=cfg_dict,
rl_device=cfg.rl_device,
sim_device=cfg.sim_device,
graphics_device_id=cfg.graphics_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=cfg.sim_device,
graphics_device_id=cfg.graphics_device_id,
headless=cfg.headless,
virtual_screen_capture=cfg.capture_video, # TODO: check
force_render=cfg.force_render)

# Instantiate the RL trainer
trainer = trainer_map[custom_args.agent](cfg=cfg.train,
env=env,
device=cfg.rl_device,
env_name=custom_args.task)
trainer = Trainers().trainer_map[custom_args.agent](cfg=cfg.train,
env=env,
device=cfg.rl_device,
env_name=custom_args.task)

# Start training
trainer.train()
Expand All @@ -64,19 +64,19 @@ def inference(custom_args):
set_seed(cfg.train.Trainer.seed)

# Instantiate the Isaac Gym environment
infer_env = task_map[custom_args.task](cfg=cfg_dict,
rl_device=cfg.rl_device,
sim_device=cfg.sim_device,
graphics_device_id=cfg.graphics_device_id,
headless=cfg.headless,
virtual_screen_capture=cfg.capture_video, # TODO: check
force_render=cfg.force_render)
infer_env = Tasks().task_map[custom_args.task](cfg=cfg_dict,
rl_device=cfg.rl_device,
sim_device=cfg.sim_device,
graphics_device_id=cfg.graphics_device_id,
headless=cfg.headless,
virtual_screen_capture=cfg.capture_video, # TODO: check
force_render=cfg.force_render)

# Instantiate the RL trainer
trainer = trainer_map[custom_args.agent](cfg=cfg.train,
env=infer_env,
device=cfg.rl_device,
env_name=custom_args.task)
trainer = Trainers().trainer_map[custom_args.agent](cfg=cfg.train,
env=infer_env,
device=cfg.rl_device,
env_name=custom_args.task)
# load checkpoint
if custom_args.ckpt_path is None:
custom_args.ckpt_path = model_zoo(name="CURICabinetRofuncRLPPO_left_arm.pth")
Expand Down
34 changes: 17 additions & 17 deletions examples/learning_rl/example_D4RL_RofuncRL.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,8 @@
import gymnasium as gym

from rofunc.config.utils import omegaconf_to_dict, get_config
from rofunc.learning.RofuncRL.tasks import task_map
from rofunc.learning.RofuncRL.trainers import trainer_map
from rofunc.learning.RofuncRL.tasks import Tasks
from rofunc.learning.RofuncRL.trainers import Trainers
from rofunc.learning.utils.download_datasets import download_d4rl_dataset
from rofunc.learning.utils.utils import set_seed

Expand All @@ -42,10 +42,10 @@ def train(custom_args):
env = gym.make(f'{custom_args.task}-v3')

# Instantiate the RL trainer
trainer = trainer_map[custom_args.agent](cfg=cfg.train,
env=env,
device=cfg.rl_device,
env_name=custom_args.task)
trainer = Trainers().trainer_map[custom_args.agent](cfg=cfg.train,
env=env,
device=cfg.rl_device,
env_name=custom_args.task)

# Start training
trainer.train()
Expand All @@ -66,19 +66,19 @@ def inference(custom_args):
set_seed(cfg.train.Trainer.seed)

# Instantiate the Isaac Gym environment
infer_env = task_map[custom_args.task](cfg=cfg_dict,
rl_device=cfg.rl_device,
sim_device=cfg.sim_device,
graphics_device_id=cfg.graphics_device_id,
headless=cfg.headless,
virtual_screen_capture=cfg.capture_video, # TODO: check
force_render=cfg.force_render)
infer_env = Tasks().task_map[custom_args.task](cfg=cfg_dict,
rl_device=cfg.rl_device,
sim_device=cfg.sim_device,
graphics_device_id=cfg.graphics_device_id,
headless=cfg.headless,
virtual_screen_capture=cfg.capture_video, # TODO: check
force_render=cfg.force_render)

# Instantiate the RL trainer
trainer = trainer_map[custom_args.agent](cfg=cfg.train,
env=infer_env,
device=cfg.rl_device,
env_name=custom_args.task)
trainer = Trainers().trainer_map[custom_args.agent](cfg=cfg.train,
env=infer_env,
device=cfg.rl_device,
env_name=custom_args.task)
# load checkpoint
if custom_args.ckpt_path is None:
raise ValueError("Please specify the checkpoint path for inference.")
Expand Down
Loading

0 comments on commit fb1fbde

Please sign in to comment.