""" Robot brain — the main think–act orchestrator. Timing ------ HAL thread : 50 Hz — steps MuJoCo, buffers observations Brain thread : 5 Hz — world model + policy → action chunk Language : on demand — only re-encodes when instruction changes Data flow each think cycle -------------------------- obs_history (8 frames) → PerceptionModule.process() → world_model.encode(frames) latent world state → language_model.encode(instr, frame) task embedding (cached) → policy.act(world_emb, task_emb) action chunk (10, 21) → hal.push_actions(chunk) """ from __future__ import annotations import logging import threading import time from typing import Optional import numpy as np from brain.perception import PerceptionModule from hal.hal import HAL from models.base import BaseDiffusionPlanner, BaseLanguageModel, BasePolicy, BaseWorldModel log = logging.getLogger(__name__) class Brain: def __init__( self, hal: HAL, world_model: BaseWorldModel, policy: BasePolicy, language_model: BaseLanguageModel, config: dict, path_planner: Optional[BaseDiffusionPlanner] = None, ): self.hal = hal self.world_model = world_model self.policy = policy self.language_model = language_model self.path_planner = path_planner # diffusion policy path planner (optional) self._cfg = config # Inject the planner into the policy so the pick can follow a planned # approach arc (the planner is otherwise only used for the world-model # trajectory). No-op for policies that don't support it. if path_planner is not None and hasattr(self.policy, "set_path_planner"): self.policy.set_path_planner(path_planner) self._perception = PerceptionModule() self._think_period = 1.0 / config.get("think_freq", 5) self._context_window = config.get("context_window", 8) self._instruction: str = config.get("default_task", "maintain balance") self._task_embedding: Optional[np.ndarray] = None self._instruction_dirty = True # force encode on first think() self._running = False self._thread: Optional[threading.Thread] = None # Metrics self._think_count = 0 self._last_think_ms = 0.0 # ────────────────────────────────────────────────── public API def set_task(self, instruction: str): """Change the active task instruction (thread-safe).""" if instruction != self._instruction: self._instruction = instruction self._instruction_dirty = True log.info("Task set to: '%s'", instruction) # Notify IK-capable policy immediately so it can start planning if hasattr(self.policy, 'set_task_text'): self.policy.set_task_text(instruction) def think(self): """One think cycle: perceive → model → plan → push.""" t0 = time.monotonic() obs_history = self.hal.get_recent_obs(self._context_window) if not obs_history: return percept = self._perception.process(obs_history) latest_frame_uint8 = (percept.frames[-1] * 255).astype(np.uint8) # Re-encode task embedding when instruction changes if self._instruction_dirty or self._task_embedding is None: self._task_embedding = self.language_model.encode( self._instruction, latest_frame_uint8 ) self._instruction_dirty = False log.debug("Task re-encoded: '%s'", self._instruction) # World model: compress video context → latent world_embedding = self.world_model.encode(percept.frames) # Multi-view images: HAL renders them on its thread at vision cadence. # Pull the most recent obs that has multi_view populated. multi_view = None for obs in reversed(obs_history): if obs.multi_view is not None: multi_view = obs.multi_view break # Diffusion Policy: plan a trajectory from current state toward the goal. # The planned path is passed to the VLA policy as conditioning so it # tracks the diffusion-generated waypoints rather than reasoning from scratch. planned_trajectory: Optional[np.ndarray] = None if self.path_planner is not None: planned_trajectory = self.path_planner.plan( world_embedding, self._task_embedding, percept.proprio_now, ) log.debug("Diffusion plan: shape=%s", planned_trajectory.shape) # Policy (VLA): world state + task embedding + planned path → action chunk. # Always pass text and proprio so IK-capable policies can use them. policy_kwargs: dict = { 'text': self._instruction, 'proprio': percept.proprio_now, } # Perception: live object poses (de-hardcodes the targets). Grounding: # language→symbol (verb + target colour). Together they let the policy # bind the instruction to the object's CURRENT position. env = getattr(self.hal, "env", None) if env is not None and hasattr(env, "perceive_objects"): try: policy_kwargs['objects'] = env.perceive_objects() except Exception as e: log.debug("perceive_objects failed: %s", e) if hasattr(self.language_model, "ground"): policy_kwargs['grounding'] = self.language_model.ground(self._instruction) if multi_view is not None: import inspect sig = inspect.signature(self.policy.act) if 'images' in sig.parameters: policy_kwargs['images'] = multi_view if planned_trajectory is not None: policy_kwargs['planned_trajectory'] = planned_trajectory actions = self.policy.act(world_embedding, self._task_embedding, **policy_kwargs) # Push to HAL action buffer self.hal.push_actions(actions) self._think_count += 1 self._last_think_ms = (time.monotonic() - t0) * 1000 log.debug( "think #%d %.1f ms queue=%d", self._think_count, self._last_think_ms, len(self.hal._action_queue), ) # ────────────────────────────────────────────────── headless run def run(self, task: Optional[str] = None, duration: float = float("inf")): """ Run headless (no viewer window). HAL steps on its own background thread; brain loops on the calling thread. """ if task: self.set_task(task) self.hal.start() t_end = time.monotonic() + duration log.info("Brain running headless for %.1f s task='%s'", duration, self._instruction) try: while time.monotonic() < t_end: t0 = time.monotonic() self.think() elapsed = time.monotonic() - t0 sleep = max(0.0, self._think_period - elapsed) time.sleep(sleep) except KeyboardInterrupt: log.info("Brain interrupted.") finally: self.hal.stop() log.info("Brain stopped after %d think cycles.", self._think_count) # ────────────────────────────────────────────────── viewer run def run_with_viewer(self, task: Optional[str] = None, duration: float = 30.0): """ Run with MuJoCo's passive viewer (requires a display). Viewer owns the main thread; brain runs on a background thread. """ import mujoco.viewer if task: self.set_task(task) self._running = True self._thread = threading.Thread( target=self._think_loop, args=(duration,), daemon=True, name="Brain", ) self._thread.start() env = self.hal.env vision_every = env.config.get("vision", {}).get("update_every_n_steps", 5) sync_rate = env.config.get("hal", {}).get("viewer_sync_rate", 60) sync_dt = 1.0 / sync_rate log.info("Launching MuJoCo viewer task='%s'", self._instruction) env.reset() with mujoco.viewer.launch_passive(env.model, env.data) as viewer: t_end = time.monotonic() + duration while viewer.is_running() and time.monotonic() < t_end: t0 = time.monotonic() with viewer.lock(): self.hal.step(vision_every=vision_every) viewer.sync() sleep = max(0.0, sync_dt - (time.monotonic() - t0)) time.sleep(sleep) self._running = False self._thread.join(timeout=2.0) log.info("Viewer closed. Think cycles: %d", self._think_count) # ────────────────────────────────────────────────── internals def _think_loop(self, duration: float): t_end = time.monotonic() + duration # Give HAL a head-start to fill the obs buffer time.sleep(0.5) while self._running and time.monotonic() < t_end: t0 = time.monotonic() self.think() elapsed = time.monotonic() - t0 time.sleep(max(0.0, self._think_period - elapsed)) # ────────────────────────────────────────────────── diagnostics @property def stats(self) -> dict: return { "think_count": self._think_count, "last_think_ms": round(self._last_think_ms, 2), "action_queue_depth": len(self.hal._action_queue), "obs_buffer_len": len(self.hal.obs_buffer), "task": self._instruction, "path_planner": self.path_planner is not None, }