Add omni installations (#42)

* Add omni installations

* Add initial implementation of isaac

* gitignore

* incorporate isaaclab

---------

Co-authored-by: Chenhao <[email protected]>

.gitignore

@@ -161,4 +161,7 @@ cython_debug/
 
 
 # VSCode
-.vscode/
+.vscode/
+
+# Outputs
+outputs/

README.md

@@ -13,7 +13,6 @@
 3. `source venv/bin/activate`
 4. `pip install -e .`
 
-
 ## Usage
 
 `python train_actsafe.py --help`

actsafe/benchmark_suites/__init__.py

@@ -19,6 +19,10 @@ def make(cfg: DictConfig) -> EnvironmentFactory:
     elif domain_name == "humanoid_bench":
         from actsafe.benchmark_suites.humanoid_bench import make
 
+        make_env = make(cfg)
+    elif domain_name == "isaaclab":
+        from actsafe.benchmark_suites.isaaclab import make
+
         make_env = make(cfg)
     else:
         raise NotImplementedError(f"Environment {domain_name} not implemented")

actsafe/benchmark_suites/isaaclab/__init__.py

@@ -0,0 +1,42 @@
+from omegaconf import DictConfig
+
+
+from actsafe.benchmark_suites.utils import get_domain_and_task
+from actsafe.rl.types import EnvironmentFactory
+
+
+def make(cfg: DictConfig) -> EnvironmentFactory:
+    import argparse
+    from omni.isaac.lab.app import AppLauncher
+
+    # add argparse arguments
+    parser = argparse.ArgumentParser()
+    # append AppLauncher cli args
+    AppLauncher.add_app_launcher_args(parser)
+    args_cli, _ = parser.parse_known_args()
+
+    _, task_cfg = get_domain_and_task(cfg)
+
+    args_cli.task = task_cfg.task
+    args_cli.headless = task_cfg.headless
+
+    # launch omniverse app
+    app_launcher = AppLauncher(args_cli)
+    simulation_app = app_launcher.app
+
+    import torch
+    from omni.isaac.lab_tasks.manager_based.locomotion.velocity.config.anymal_d.flat_env_cfg import AnymalDFlatEnvCfg
+
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.deterministic = False
+    torch.backends.cudnn.benchmark = False
+
+    def make_env():
+        import gymnasium as gym
+
+        env_cfg = AnymalDFlatEnvCfg()
+        env = gym.make(args_cli.task, cfg=env_cfg, render_mode=None)
+        return env
+
+    return make_env

actsafe/configs/config.yaml

@@ -35,6 +35,7 @@ mixed_precision: params=float32,compute=float16,output=float32
 
 training:
   trainer: rl
+  async_parallel: true
   safe: true
   safety_budget: 25
   seed: 0

actsafe/configs/environment/isaaclab.yaml

@@ -0,0 +1,7 @@
+isaaclab:
+  task: Isaac-Velocity-Flat-Anymal-D-v0
+  headless: true
+  image_observation:
+    enabled: false
+    image_size: [64, 64]
+    image_format: "channels_first"

actsafe/configs/experiment/isaaclab.yaml

@@ -0,0 +1,6 @@
+# @package _global_
+defaults:
+  - override /environment: isaaclab
+
+training:
+  async_parallel: false

actsafe/rl/trainer.py

@@ -81,7 +81,7 @@ def __enter__(self):
             self.config.training.parallel_envs,
             self.config.training.time_limit,
             self.config.training.action_repeat,
-        )
+        ) if self.config.training.async_parallel else self.make_env()
         if self.seeds is None:
             self.seeds = PRNGSequence(self.config.training.seed)
         if self.agent is None:

pyproject.toml

@@ -28,7 +28,6 @@ safe-adaptation-gym = { git = "https://[email protected]/lasgroup/safe-adaptation-g
 jmp = { git = "https://github.com/deepmind/jmp" }
 tensorboard = "^2.16.2"
 
-
 [[tool.poetry.source]]
 name = "PyPI"
 priority = "primary"