Refactor: Improve organization

Signed-off-by: Andrej Orsula <[email protected]>

space_robotics_bench/assets/light/sky.py

@@ -19,6 +19,9 @@ def sky_from_env_cfg(
 ) -> Optional[AssetBaseCfg]:
     texture_file: Optional[str] = None
 
+    # Enable rendering effects
+    rtx_settings.auto_exposure()
+
     match env_cfg.scenario:
         case env_utils.Scenario.EARTH:
             texture_file = (

space_robotics_bench/envs/__init__.py

@@ -1,3 +1,3 @@
-from .aerial_robotics import *  # noqa: F403
+from .aerial import *  # noqa: F403
 from .manipulation import *  # noqa: F403
-from .mobile_robotics import *  # noqa: F403
+from .mobile import *  # noqa: F403

space_robotics_bench/envs/aerial_robotics/cfg.py → space_robotics_bench/envs/aerial/cfg.py

@@ -101,6 +101,9 @@ class BaseAerialRoboticsEnvCfg(BaseEnvCfg):
             gpu_collision_stack_size=2 ** (26 - 4),
             gpu_max_num_partitions=8,
         ),
+        render=sim_utils.RenderCfg(
+            enable_reflections=True,
+        ),
         physics_material=sim_utils.RigidBodyMaterialCfg(
             static_friction=1.0,
             dynamic_friction=1.0,

space_robotics_bench/envs/manipulation/cfg.py

@@ -87,6 +87,9 @@ class BaseManipulationEnvCfg(BaseEnvCfg):
             gpu_collision_stack_size=2 ** (26 - 3),
             gpu_max_num_partitions=8,
         ),
+        render=sim_utils.RenderCfg(
+            enable_reflections=True,
+        ),
         physics_material=sim_utils.RigidBodyMaterialCfg(
             static_friction=1.0,
             dynamic_friction=1.0,

space_robotics_bench/envs/manipulation/impl.py

@@ -1,7 +1,4 @@
-from omni.isaac.lab.sensors import ContactSensor
-from omni.isaac.lab.sensors.frame_transformer.frame_transformer_cfg import (
-    FrameTransformer,
-)
+from omni.isaac.lab.sensors import ContactSensor, FrameTransformer
 
 from space_robotics_bench.core.assets import Articulation
 from space_robotics_bench.core.envs import BaseEnv

space_robotics_bench/envs/mobile_robotics/cfg.py → space_robotics_bench/envs/mobile/cfg.py

@@ -85,6 +85,10 @@ class BaseMobileRoboticsEnvCfg(BaseEnvCfg):
             gpu_collision_stack_size=2 ** (26 - 4),
             gpu_max_num_partitions=8,
         ),
+        render=sim_utils.RenderCfg(
+            enable_reflections=True,
+            enable_translucency=True,
+        ),
         physics_material=sim_utils.RigidBodyMaterialCfg(
             static_friction=1.0,
             dynamic_friction=1.0,

space_robotics_bench/tasks/debris_capture/task.py

@@ -296,7 +296,7 @@ def _compute_intermediate_state(
         joint_pos_normalized[:, robot_hand_joint_indices],
     )
 
-    # End-effector pose (position and '6D' rotation)
+    # End-effector '6D' rotation
     robot_ee_rotmat_wrt_base = math_utils.matrix_from_quat(robot_ee_quat_wrt_base)
 
     # Transformation | Object origin -> Object CoM

space_robotics_bench/tasks/peg_in_hole/task.py

@@ -359,7 +359,7 @@ def _compute_intermediate_state(
         joint_pos_normalized[:, robot_hand_joint_indices],
     )
 
-    # End-effector pose (position and '6D' rotation)
+    # End-effector '6D' rotation
     robot_ee_rotmat_wrt_base = math_utils.matrix_from_quat(robot_ee_quat_wrt_base)
 
     # Transformation | Object origin -> Object CoM

space_robotics_bench/tasks/peg_in_hole_multi/task.py

@@ -1,5 +1,5 @@
 from dataclasses import MISSING as DELAYED_CFG
-from typing import Dict, List, Optional, Sequence, Tuple
+from typing import Dict, List, Sequence, Tuple
 
 import torch
 from omni.isaac.core.prims.xform_prim_view import XFormPrimView
@@ -392,7 +392,7 @@ def _compute_intermediate_state(
         joint_pos_normalized[:, robot_hand_joint_indices],
     )
 
-    # End-effector pose (position and '6D' rotation)
+    # End-effector '6D' rotation
     robot_ee_rotmat_wrt_base = math_utils.matrix_from_quat(robot_ee_quat_wrt_base)
 
     # Transformation | Object origin -> Object CoM

space_robotics_bench/tasks/sample_collection/task.py

@@ -350,7 +350,7 @@ def _compute_intermediate_state(
         joint_pos_normalized[:, robot_hand_joint_indices],
     )
 
-    # End-effector pose (position and '6D' rotation)
+    # End-effector '6D' rotation
     robot_ee_rotmat_wrt_base = math_utils.matrix_from_quat(robot_ee_quat_wrt_base)
 
     # Transformation | Object origin -> Object CoM

space_robotics_bench/tasks/sample_collection_multi/task.py

@@ -1,5 +1,5 @@
 from dataclasses import MISSING as DELAYED_CFG
-from typing import Dict, List, Optional, Sequence, Tuple
+from typing import Dict, List, Sequence, Tuple
 
 import torch
 from omni.isaac.lab.managers import EventTermCfg, SceneEntityCfg
@@ -340,7 +340,7 @@ def _compute_intermediate_state(
         joint_pos_normalized[:, robot_hand_joint_indices],
     )
 
-    # End-effector pose (position and '6D' rotation)
+    # End-effector '6D' rotation
     robot_ee_rotmat_wrt_base = math_utils.matrix_from_quat(robot_ee_quat_wrt_base)
 
     # Transformation | Object origin -> Object CoM

space_robotics_bench/tasks/solar_panel_assembly/task.py

@@ -1,4 +1,3 @@
-import os
 import sys
 from typing import Any, Dict, List, Optional, Sequence, Tuple
 
@@ -11,18 +10,12 @@
 import space_robotics_bench.core.sim as sim_utils
 import space_robotics_bench.utils.math as math_utils
 from space_robotics_bench import assets
-from space_robotics_bench.core.assets import (
-    AssetBaseCfg,
-    AssetCfg,
-    RigidObject,
-    RigidObjectCfg,
-)
+from space_robotics_bench.core.assets import AssetCfg, RigidObject, RigidObjectCfg
 from space_robotics_bench.core.managers import EventTermCfg, SceneEntityCfg
 from space_robotics_bench.core.markers import (
     VisualizationMarkers,
     VisualizationMarkersCfg,
 )
-from space_robotics_bench.core.sim.schemas import RigidBodyPropertiesCfg
 from space_robotics_bench.envs import (
     BaseManipulationEnv,
     BaseManipulationEnvCfg,
@@ -514,7 +507,7 @@ def _compute_intermediate_state(
         joint_pos_normalized[:, robot_hand_joint_indices],
     )
 
-    # End-effector pose (position and '6D' rotation)
+    # End-effector '6D' rotation
     robot_ee_rotmat_wrt_base = math_utils.matrix_from_quat(robot_ee_quat_wrt_base)
 
     # Transformation | Object origin -> Object CoM