[WIP] Load extension entrypoints lazily

src/holosoma_inference/docker/run.sh

@@ -17,6 +17,11 @@ PROJECT_ROOT="$( cd "$SCRIPT_DIR/../../.." && pwd )"
 CONTAINER_NAME="holosoma-inference-container"
 IMAGE_NAME="holosoma-inference"
 
+# Mount bash history in the host filesystem in order to preserve history between containers
+HOLOSOMA_DEPS_DIR="$HOME/.holosoma_deps"
+mkdir -p "$HOLOSOMA_DEPS_DIR"
+touch "$HOLOSOMA_DEPS_DIR/bash_history"
+
 # Function to create a new container
 create_container() {
     local mounts=(
@@ -25,6 +30,8 @@ create_container() {
         -v /tmp/.X11-unix:/tmp/.X11-unix
         -v $HOME/.Xauthority:/root/.Xauthority:ro
         -v /dev/shm:/dev/shm
+        -v /dev/input:/dev/input
+        -v "$HOLOSOMA_DEPS_DIR/bash_history":/root/.bash_history
     )
     [[ -d "$EXT_DIR" ]] && mounts+=(-v "$EXT_DIR":/workspace/holosoma-extension) # optionally mount extension repo
 

src/holosoma_inference/holosoma_inference/config/config_values/inference.py

@@ -52,19 +52,56 @@
     task=task.wbt,
 )
 
+# Core defaults - no extension imports at module load time
 DEFAULTS = {
     "g1-29dof-loco": g1_29dof_loco,
     "t1-29dof-loco": t1_29dof_loco,
     "g1-29dof-wbt": g1_29dof_wbt,
 }
 
-# Auto-discover inference configs from installed extensions
-for ep in entry_points(group="holosoma.config.inference"):
-    DEFAULTS[ep.name] = ep.load()
+# Track whether extensions have been loaded
+_extensions_loaded = False
 
-AnnotatedInferenceConfig = Annotated[
-    InferenceConfig,
-    tyro.conf.arg(
-        constructor=tyro.extras.subcommand_type_from_defaults({f"inference:{k}": v for k, v in DEFAULTS.items()})
+
+def _load_extensions() -> None:
+    """Lazily load extension configs from entry points.
+
+    This is deferred to avoid circular imports when extensions import
+    from holosoma_inference.config at module load time.
+    """
+    global _extensions_loaded  # noqa: PLW0603
+    if _extensions_loaded:
+        return
+    _extensions_loaded = True
+    for ep in entry_points(group="holosoma.config.inference"):
+        DEFAULTS[ep.name] = ep.load()
+
+
+def get_annotated_inference_config() -> type:
+    """Build the annotated InferenceConfig type with all discovered configs.
+
+    This function loads extension configs lazily and returns a tyro-compatible
+    annotated type for CLI subcommand generation.
+
+    Returns:
+        Annotated type suitable for use with tyro.cli()
+    """
+    _load_extensions()
+    return Annotated[
+        InferenceConfig,
+        tyro.conf.arg(
+            constructor=tyro.extras.subcommand_type_from_defaults(
+                {f"inference:{k}": v for k, v in DEFAULTS.items()}
+            )
         ),
-]
+    ]
+
+
+def get_defaults() -> dict:
+    """Get all inference config defaults, including extensions.
+
+    Returns:
+        Dictionary mapping config names to InferenceConfig instances.
+    """
+    _load_extensions()
+    return DEFAULTS

src/holosoma_inference/holosoma_inference/config/config_values/robot.py

@@ -192,6 +192,7 @@
 # Default Configurations Dictionary
 # =============================================================================
 
+# Core defaults - no extension imports at module load time
 DEFAULTS = {
     "g1-29dof": g1_29dof,
     "t1-29dof": t1_29dof,
@@ -201,6 +202,29 @@
 Keys use hyphen-case naming convention for CLI compatibility.
 """
 
-# Auto-discover robot configs from installed extensions
-for ep in entry_points(group="holosoma.config.robot"):
-    DEFAULTS[ep.name] = ep.load()
+# Track whether extensions have been loaded
+_extensions_loaded = False
+
+
+def _load_extensions() -> None:
+    """Lazily load extension configs from entry points.
+
+    This is deferred to avoid circular imports when extensions import
+    from holosoma_inference.config at module load time.
+    """
+    global _extensions_loaded  # noqa: PLW0603
+    if _extensions_loaded:
+        return
+    _extensions_loaded = True
+    for ep in entry_points(group="holosoma.config.robot"):
+        DEFAULTS[ep.name] = ep.load()
+
+
+def get_defaults() -> dict:
+    """Get all robot config defaults, including extensions.
+
+    Returns:
+        Dictionary mapping config names to RobotConfig instances.
+    """
+    _load_extensions()
+    return DEFAULTS

src/holosoma_inference/holosoma_inference/policies/wbt.py

@@ -251,6 +251,9 @@ def rl_inference(self, robot_state_data):
             self.curr_motion_timestep = self.timestep_util.timestep
 
         obs = self.prepare_obs_for_rl(robot_state_data)
+        if self.config.task.print_observations:
+            self._print_observations(obs)
+
         input_feed = {"time_step": np.array([[self.curr_motion_timestep]], dtype=np.float32), "obs": obs["actor_obs"]}
         policy_action, self.motion_command_t, self.ref_quat_xyzw_t = self.policy(input_feed)
 

src/holosoma_inference/holosoma_inference/run_policy.py

@@ -17,7 +17,7 @@
 from loguru import logger
 
 from holosoma_inference.config.config_types.inference import InferenceConfig
-from holosoma_inference.config.config_values.inference import AnnotatedInferenceConfig
+from holosoma_inference.config.config_values.inference import get_annotated_inference_config
 from holosoma_inference.config.utils import TYRO_CONFIG
 from holosoma_inference.policies.locomotion import LocomotionPolicy
 from holosoma_inference.policies.wbt import WholeBodyTrackingPolicy
@@ -120,7 +120,8 @@ def run_policy(config: InferenceConfig):
 def main(annotated_config=None):
     """Main entry point. Extensions can pass their own AnnotatedInferenceConfig."""
     if annotated_config is None:
-        annotated_config = AnnotatedInferenceConfig
+        # Use factory function to lazily load extension configs
+        annotated_config = get_annotated_inference_config()
     config = tyro.cli(annotated_config, config=TYRO_CONFIG)
     run_policy(config)