feat: new codebase architecture (#3)

* wip: completed base classes

* feat: completed first set of Terrain classes

* feat: completed wrapper classes & started implementation for others

* feat: achieved error-free agent creation

* fix: fully customized setup script to reduce files in PYTHONPATH

* fix: fully customized setup script to reduce files in PYTHONPATH

* feat: completed IMU implementation

* wip: idle Newton gait

* feat: laid groundwork for newton's idle task

* feat: refactor reward calculations and environment setup

- Added GPU memory check and stage clearing in environment construction.
- Integrated physics and lighting setup in environment construction.
- Refactored reward calculation logic in NewtonIdleTask.
- Improved observation handling in NewtonIdleTask.
- Enhanced terrain setup in NewtonMultiTerrainEnv.
- Updated IMU data handling in VecIMU.
- Adjusted configuration parameters in world.yaml and newton_idle_task.yaml.

* feat: refactor reward calculations and environment setup

- Added GPU memory check and stage clearing in environment construction.
- Integrated physics and lighting setup in environment construction.
- Refactored reward calculation logic in NewtonIdleTask.
- Improved observation handling in NewtonIdleTask.
- Enhanced terrain setup in NewtonMultiTerrainEnv.
- Updated IMU data handling in VecIMU.
- Adjusted configuration parameters in world.yaml and newton_idle_task.yaml.

* feat: add new configuration files and update parameters for Newton tasks

* feat: adjusted ROMs & GPU buffer allocations

* chore: renamed `Terrain*` to `BaseTerrain*`

* chore: laid groundwork for new domain randomizer!

* feat: added Universe & domain randomizer classes

* fix: unknown args are passed on to `SimulationApp`

* fix: removed physics issues with flying Newtons

* fix: fixed linear & angular acceleration calculations

* fix: changed if to assert in `imu.py`

* feat: created todos GA

.github/workflows/todos.yaml

@@ -0,0 +1,21 @@
+name: "Run TODO to Issue"
+on:
+  push:
+  workflow_dispatch:
+    inputs:
+      MANUAL_COMMIT_REF:
+        description: "The SHA of the commit to get the diff for"
+        required: true
+      MANUAL_BASE_REF:
+        description: "By default, the commit entered above is compared to the one directly before it; to go back further, enter an earlier SHA here"
+        required: false
+jobs:
+  build:
+    runs-on: "ubuntu-latest"
+    steps:
+      - uses: "actions/checkout@v4"
+      - name: "TODO to Issue"
+        uses: "alstr/todo-to-issue-action@v5"
+        with:
+          CLOSE_ISSUES: true
+          AUTO_ASSIGN: true

apps/omni.isaac.sim.newton.kit

@@ -163,7 +163,7 @@ dev_build = false
 fabricDefaultStageFrameHistoryCount = 3 # needed for omni.syntheticdata TODO105 Still True?
 
 [settings.app.window]
-title = "TWIP Sim Python"
+title = "Newton Sim Python"
 hideUi = false
 _iconSize = 256
 iconPath = "${omni.isaac.kit}/data/omni.isaac.sim.png"
@@ -189,9 +189,9 @@ size = 16
 featuredExts = []
 
 [settings]
-log.level = "fatal"
-log.fileLogLevel = "fatal"
-log.outputStreamLevel = "fatal"
+log.level = "info"
+log.fileLogLevel = "info"
+log.outputStreamLevel = "info"
 
 # MGPU is always on, you can turn it from the settings, and force this off to save even more resource if you
 # only want to use a single GPU on your MGPU system

configs/task_twip.yaml → configs/newton_idle_task.yaml

@@ -1,21 +1,25 @@
-task_name: "twip_balancing"
+task_name: "newton_idle"
 device: "cuda"
 policy: "MlpPolicy"
 seed: 14321
 
-n_envs: 128
+n_envs: 96
 timesteps_per_env: 500_000
 base_lr: 3e-4
+episode_length: 5000
+
+newton:
+  inverse_control_frequency: 4  # we're aiming for a total compute budget of 2 times the physics timestep (0.005s * 4 = 0.02s)
 
 delay:
-  enabled: True
+  enabled: False
   obs_delay_range: [1, 8] # sample range of time steps to delay the observations by
   act_delay_range: [1, 2] # sample range of time steps to delay the actions by
   instant_rewards: False # should the rewards be calculated on the brain or remote actor?
 
 ppo:
-  n_steps: 512
-  batch_size: 256
+  n_steps: 24
+  batch_size: 2304
   n_epochs: 5
 
   gamma: 0.99

configs/world.yaml

@@ -1,5 +1,26 @@
-physics_dt: 0.02000
+physics_dt: 0.005000
 stage_units_in_meters: 1.0
 rendering_dt: 0.01667
 backend: torch
-device: cuda
+device: cuda
+
+# the following overrides what is in assets/physics.usd
+sim_params:
+  substeps: 1
+  solver_type: 1 # 0: PGS, 1: TGS
+  enable_stabilization: True # true for a 2nd pass to stabilize
+  enable_scene_query_support: True # true to interact with the scene with the mouse
+
+  use_gpu_pipeline: True # true for GPU dynamics (i.e. deformables)
+
+  # GPU buffers
+  gpu_max_rigid_contact_count: 2097152
+  gpu_max_rigid_patch_count: 655360
+  gpu_found_lost_pairs_capacity: 16776216
+  gpu_found_lost_aggregate_pairs_capacity: 107701464
+  gpu_total_aggregate_pairs_capacity: 4194304
+  gpu_max_soft_body_contacts: 1024 # Soft body contacts are not used for Newton-related sims, so we keep it low
+  gpu_max_particle_contacts: 1024 # Particle contacts are not used for Newton-related sims, so we keep it low
+  gpu_heap_capacity: 134217728
+  gpu_temp_buffer_capacity: 33554432
+  gpu_max_num_partitions: 8

core/agents/__init__.py

@@ -0,0 +1,3 @@
+from .base_agent import *
+from .newton_base_agent import *
+from .newton_vec_agent import *

core/agents/base_agent.py

@@ -0,0 +1,48 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from core.globals import AGENTS_PATH
+from core.types import Observations, Actions
+from core.universe import Universe
+from omni.isaac.core import World
+
+
+class BaseAgent(ABC):
+    def __init__(self, num_agents: int) -> None:
+        self.path: str = ""
+        self.num_agents: int = num_agents
+        self.universe: Optional[Universe] = None
+
+        self._is_constructed: bool = False
+
+    @abstractmethod
+    def construct(self, universe: Universe) -> None:
+        assert (
+            not self._is_constructed
+        ), f"{self.__class__.__name__} already constructed: tried to construct!"
+
+        self.path = AGENTS_PATH
+        self.universe = universe
+
+        from omni.isaac.core.utils.prims import create_prim
+
+        create_prim(
+            prim_path=self.path,
+            prim_type="Scope",
+        )
+
+    @abstractmethod
+    def step(self, actions: Actions) -> Observations:
+        assert (
+            self._is_constructed
+        ), f"{self.__class__.__name__} not constructed: tried to step!"
+
+        return self.get_observations()
+
+    @abstractmethod
+    def get_observations(self) -> Observations:
+        assert (
+            self._is_constructed
+        ), f"{self.__class__.__name__} not constructed: tried to get observations!"
+
+        return {}

core/agents/newton_base_agent.py

@@ -0,0 +1,44 @@
+from abc import abstractmethod
+
+import torch
+from core.agents import BaseAgent
+from core.controllers import VecJointsController
+from core.sensors import VecIMU
+from core.types import Actions, Observations
+from core.universe import Universe
+
+
+class NewtonBaseAgent(BaseAgent):
+    def __init__(
+        self,
+        num_agents: int,
+        imu: VecIMU,
+        joints_controller: VecJointsController,
+    ) -> None:
+        super().__init__(num_agents)
+
+        self.base_path_expr: str = ""
+
+        self.imu: VecIMU = imu
+        self.joints_controller: VecJointsController = joints_controller
+
+    @abstractmethod
+    def construct(self, universe: Universe) -> None:
+        super().construct(universe)
+
+    @abstractmethod
+    def step(self, actions: Actions) -> None:
+        super().step(actions)
+
+        new_joint_positions = torch.from_numpy(actions).to(self.universe.physics_device)
+        self.joints_controller.step(new_joint_positions)
+
+    @abstractmethod
+    def get_observations(self) -> Observations:
+        imu_data_tensor = self.imu.get_data()
+        imu_data_numpy = {}
+
+        for key, value in imu_data_tensor.items():
+            imu_data_numpy[key] = value.cpu().numpy()
+
+        return imu_data_numpy

core/agents/newton_vec_agent.py

@@ -0,0 +1,50 @@
+from core.agents import NewtonBaseAgent
+from core.globals import TERRAINS_PATH, PHYSICS_SCENE_PATH, COLLISION_GROUPS_PATH
+from core.types import Observations, Actions
+from core.universe import Universe
+
+
+class NewtonVecAgent(NewtonBaseAgent):
+    def construct(self, universe: Universe) -> None:
+        super().construct(universe)
+
+        self.base_path_expr = self.path + "/Newton_*/base"
+        self.path = self.path + "/Newton_0"
+
+        import omni.isaac.core.utils.stage as stage_utils
+
+        stage_utils.add_reference_to_stage(
+            "assets/newton/newton.usd",
+            prim_path=self.path,
+        )
+
+        from omni.isaac.cloner import Cloner
+
+        cloner = Cloner()
+        cloner.define_base_env(self.path)
+
+        agent_paths = cloner.generate_paths(self.path[:-2], self.num_agents)
+
+        cloner.filter_collisions(
+            prim_paths=agent_paths,
+            physicsscene_path=PHYSICS_SCENE_PATH,
+            collision_root_path=COLLISION_GROUPS_PATH,
+            global_paths=[TERRAINS_PATH],
+        )
+        cloner.clone(
+            source_prim_path=self.path,
+            prim_paths=agent_paths,
+        )
+
+        self.imu.construct(self.base_path_expr)
+        self.joints_controller.construct(self.base_path_expr)
+
+        self.universe.reset()
+
+        self._is_constructed = True
+
+    def step(self, actions: Actions) -> None:
+        return super().step(actions)
+
+    def get_observations(self) -> Observations:
+        return super().get_observations()

core/animation/__init__.py

@@ -0,0 +1 @@
+from .animation_engine import AnimationEngine

core/animation/animation_engine.py

@@ -0,0 +1,6 @@
+class AnimationEngine:
+    def __init__(self):
+        pass
+
+    def __del__(self):
+        pass

core/controllers/__init__.py

@@ -0,0 +1 @@
+from .joints_controller import VecJointsController