holosoma 0.0.1__py3-none-any.whl

holosoma/agents/base_algo/base_algo.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable
+
+from holosoma.config_types.algo import AlgoInitConfig
+from holosoma.envs.base_task.base_task import BaseTask
+from holosoma.utils.safe_torch_import import torch
+
+if TYPE_CHECKING:
+    from holosoma.config_types.experiment import ExperimentConfig
+
+
+class BaseAlgo:
+    def __init__(self, env: BaseTask, config: AlgoInitConfig, device, multi_gpu_cfg=None):
+        self.env = env
+        self.config = config
+        self.device = device
+
+        self.is_multi_gpu = multi_gpu_cfg is not None
+        if multi_gpu_cfg is not None:
+            self.gpu_global_rank = multi_gpu_cfg["global_rank"]
+            self.gpu_local_rank = multi_gpu_cfg["local_rank"]
+            self.gpu_world_size = multi_gpu_cfg["world_size"]
+        else:
+            self.gpu_global_rank = 0
+            self.gpu_local_rank = 0
+            self.gpu_world_size = 1
+        self.is_main_process = self.gpu_global_rank == 0
+        self._experiment_config: ExperimentConfig | None = None
+        self._wandb_run_path: str | None = None
+
+    def setup(self):
+        return NotImplementedError
+
+    def learn(self):
+        return NotImplementedError
+
+    def load(self, path):
+        return NotImplementedError
+
+    @property
+    def inference_model(self):
+        return NotImplementedError
+
+    @property
+    def actor_onnx_wrapper(self):
+        return NotImplementedError
+
+    def env_step(self, actions, extra_info=None):
+        obs_dict, rewards, dones, extras = self.env.step(actions, extra_info)
+        return obs_dict, rewards, dones, extras
+
+    def attach_checkpoint_metadata(
+        self,
+        experiment_config: ExperimentConfig,
+        wandb_run_path: str | None = None,
+    ) -> None:
+        """Attach metadata that should be saved with checkpoints."""
+
+        self._experiment_config = experiment_config
+        self._wandb_run_path = wandb_run_path
+
+    def _checkpoint_metadata(self, iteration: int | None = None) -> dict[str, Any]:
+        if self._experiment_config is None:
+            raise RuntimeError("Experiment config metadata missing. Call attach_checkpoint_metadata() before saving.")
+
+        metadata: dict[str, Any] = {"experiment_config": self._experiment_config.to_serializable_dict()}
+        if self._wandb_run_path:
+            metadata["wandb_run_path"] = self._wandb_run_path
+        if iteration is not None:
+            metadata["iteration"] = int(iteration)
+        return metadata
+
+    def has_curricula_enabled(self) -> bool:
+        """Check if any curricula are enabled in the environment.
+
+        This helper method checks for the presence of various curriculum flags
+        to determine if any curriculum learning is active. This is commonly used
+        for multi-GPU synchronization and logging purposes.
+
+        Returns
+        -------
+        bool
+            True if any curriculum is enabled, False otherwise.
+        """
+        return getattr(self.env, "use_reward_penalty_curriculum", False) or getattr(
+            self.env, "use_domain_rand_scale_curriculum", False
+        )
+
+    def _synchronize_curriculum_metrics(self):
+        """Synchronize curriculum-related metrics across all GPUs."""
+        # Check if any curricula are enabled before synchronizing
+        if not self.has_curricula_enabled():
+            return
+
+        env = self._unwrap_env()
+        env.synchronize_curriculum_state(device=self.device, world_size=self.gpu_world_size)
+
+    def get_inference_policy(self, device: str | None = None) -> Callable[[dict[str, torch.Tensor]], torch.Tensor]:
+        """Get a callable policy function for inference.
+
+        This method returns a function that takes observations as input and returns
+        actions. The policy function is configured to run on the specified device.
+
+        Parameters
+        ----------
+        device : str | None, optional
+            The device to run the policy on (e.g., 'cuda', 'cpu').
+            If None, uses the default device.
+
+        Returns
+        -------
+        Callable[[torch.Tensor], torch.Tensor]
+            A function that takes observations as input and returns actions.
+            The function expects input observations as a torch.Tensor and
+            returns actions as a torch.Tensor. Both input and output tensors
+            should be on the specified device.
+
+        Notes
+        -----
+        This is an abstract method that should be implemented by subclasses.
+        The returned policy function should:
+        - Run the policy network on the specified device
+        - Return actions in the expected format
+        """
+        raise NotImplementedError
+
+    @torch.no_grad()
+    def evaluate_policy(self, max_eval_steps: int | None = None):
+        raise NotImplementedError
+
+    def save(self, path=None, name="last.ckpt"):
+        raise NotImplementedError
+
+    def _unwrap_env(self) -> BaseTask | Any:
+        """Return the underlying environment.
+
+        Algorithms that wrap the task (e.g. ``FastSACEnv``) keep a reference to the
+        original environment as ``unwrapped_env`` during construction, so we simply
+        return that when it is present.
+        """
+        return getattr(self, "unwrapped_env", self.env)
+
+    def _collect_env_state(self) -> dict[str, torch.Tensor | float]:
+        """Collect environment state for checkpointing via the environment interface."""
+        env = self._unwrap_env()
+        state = env.get_checkpoint_state()
+        return state or {}
+
+    def _restore_env_state(self, env_state: dict[str, torch.Tensor | float] | None) -> None:
+        """Restore environment state from checkpoint via the environment interface."""
+        if not env_state:
+            return
+        env = self._unwrap_env()
+        env.load_checkpoint_state(env_state)

holosoma/agents/callbacks/base_callback.py

@@ -0,0 +1,21 @@
+from torch.nn import Module
+
+
+class RLEvalCallback(Module):
+    def __init__(self, config, training_loop):
+        super().__init__()
+        self.config = config
+        self.training_loop = training_loop
+        self.device = self.training_loop.device
+
+    def on_pre_evaluate_policy(self):
+        pass
+
+    def on_pre_eval_env_step(self, actor_state):
+        return actor_state
+
+    def on_post_eval_env_step(self, actor_state):
+        return actor_state
+
+    def on_post_evaluate_policy(self):
+        pass

holosoma/agents/callbacks/payload.py

@@ -0,0 +1,172 @@
+"""Eval callback that applies constant downward forces on wrists to simulate a held payload.
+
+Applies F = mass * g / num_bodies on each resolved wrist/elbow link
+as a sustained downward (-Z) force throughout the evaluation episode. Uses the same
+IsaacLab set_external_force_and_torque API as the push callback, with substep-level
+force application via monkey-patching.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+from loguru import logger
+
+from holosoma.agents.callbacks.base_callback import RLEvalCallback
+from holosoma.config_types.eval_callback import PayloadConfig
+from holosoma.utils.safe_torch_import import torch
+
+GRAVITY = 9.81
+
+
+class EvalPayloadCallback(RLEvalCallback):
+    """Applies constant downward forces on wrist links to simulate holding a payload.
+
+    The total payload weight (mass * g) is split evenly across all resolved
+    wrist bodies. Forces are applied at every physics substep via monkey-patch.
+
+    Recorded arrays (per env step):
+        payload_force_per_body [T]     force magnitude per body in N
+        payload_body_pos_w     [T, N, 3] world positions of payload bodies
+    """
+
+    def __init__(
+        self,
+        config: PayloadConfig,
+        training_loop: Any = None,
+    ):
+        super().__init__(config, training_loop)
+        self.env_id = config.env_id
+        self.candidate_body_names = [s.strip() for s in config.body_names.split(",")]
+        self._record_buffers: dict[str, list[np.ndarray]] | None = None
+        self._record_metadata: dict[str, Any] | None = None
+
+        # Resolved at on_pre_evaluate_policy
+        self._resolved_bodies: list[tuple[str, int, int]] = []
+        self._force_per_body: float = 0.0
+        self._sim: Any = None
+        self._original_apply_force: Any = None
+        self._step_count: int = 0
+
+    def _get_env(self):
+        return self.training_loop._unwrap_env()
+
+    def _find_recording_callback(self):
+        from holosoma.agents.callbacks.recording import EvalRecordingCallback
+
+        for cb in self.training_loop.eval_callbacks:
+            if isinstance(cb, EvalRecordingCallback):
+                return cb
+        return None
+
+    def on_pre_evaluate_policy(self) -> None:
+        # Discover recording callback for shared buffer/metadata access
+        recording_cb = self._find_recording_callback()
+        if recording_cb is not None:
+            self._record_buffers = recording_cb._buffers
+            self._record_metadata = recording_cb._metadata
+
+        env = self._get_env()
+        self._sim = env.simulator
+        if self._sim.simulator_config.name != "isaacsim":
+            raise RuntimeError(f"EvalPayloadCallback requires IsaacSim, got '{self._sim.simulator_config.name}'")
+
+        # Resolve body names
+        self._resolved_bodies = []
+        available = list(getattr(self._sim, "body_names", []))
+        for name in self.candidate_body_names:
+            body_names_idx = self._sim.find_rigid_body_indice(name)
+            if isinstance(body_names_idx, int) and body_names_idx >= 0:
+                isaac_body_id = self._sim.body_ids[body_names_idx]
+                self._resolved_bodies.append((name, body_names_idx, isaac_body_id))
+            else:
+                raise ValueError(f"EvalPayloadCallback: body '{name}' not found. Available: {available}")
+
+        total_force = self.config.mass_kg * GRAVITY
+        self._force_per_body = total_force / len(self._resolved_bodies)
+
+        resolved_names = [name for name, _, _ in self._resolved_bodies]
+        logger.info(
+            f"EvalPayloadCallback: {len(self._resolved_bodies)} payload bodies: {resolved_names}, "
+            f"mass={self.config.mass_kg}kg, total_force={total_force:.1f}N, "
+            f"force_per_body={self._force_per_body:.1f}N"
+        )
+
+        # Store config in recording metadata
+        if self._record_metadata is not None:
+            self._record_metadata["payload_config"] = {
+                "body_names": resolved_names,
+                "mass_kg": self.config.mass_kg,
+                "total_force_n": total_force,
+                "force_per_body_n": self._force_per_body,
+            }
+
+        # Initialize recording buffers
+        buffers = self._get_buffers()
+        if buffers is not None:
+            buffers["payload_force_per_body"] = []
+            buffers["payload_body_pos_w"] = []
+
+        # Monkey-patch substep force application
+        self._original_apply_force = env._apply_force_in_physics_step
+
+        def _patched_apply_force():
+            self._original_apply_force()
+            self._apply_payload_forces()
+
+        env._apply_force_in_physics_step = _patched_apply_force
+
+    def on_post_eval_env_step(self, actor_state: dict) -> dict:
+        buffers = self._get_buffers()
+        if buffers is not None:
+            eid = self.env_id
+            body_positions = []
+            for _name, _body_names_idx, isaac_body_id in self._resolved_bodies:
+                pos = self._sim._robot.data.body_pos_w[eid, isaac_body_id].detach().cpu().numpy().copy()
+                body_positions.append(pos)
+
+            buffers["payload_force_per_body"].append(np.array(self._force_per_body, dtype=np.float32))
+            buffers["payload_body_pos_w"].append(np.stack(body_positions, axis=0).astype(np.float32))
+
+        self._step_count += 1
+        return actor_state
+
+    def on_post_evaluate_policy(self) -> None:
+        self._get_env()._apply_force_in_physics_step = self._original_apply_force
+        logger.info(
+            f"EvalPayloadCallback: completed {self._step_count} steps, "
+            f"payload={self.config.mass_kg}kg on {len(self._resolved_bodies)} bodies"
+        )
+
+    def _apply_payload_forces(self) -> None:
+        """Apply downward force on all resolved wrist bodies (called each substep)."""
+        from isaaclab.utils.math import quat_apply_inverse
+
+        eid = self.env_id
+        device = self._sim.sim_device
+
+        # Force in world frame: straight down (-Z)
+        force_world = torch.tensor(
+            [0.0, 0.0, -self._force_per_body],
+            dtype=torch.float32,
+            device=device,
+        )
+
+        for _name, _body_names_idx, isaac_body_id in self._resolved_bodies:
+            # Transform world-frame force to body-local frame
+            body_quat_w = self._sim._robot.data.body_quat_w[eid, isaac_body_id]
+            force_body = quat_apply_inverse(body_quat_w, force_world)
+
+            forces = force_body.unsqueeze(0).unsqueeze(0)  # [1, 1, 3]
+            torques = torch.zeros_like(forces)
+
+            self._sim._robot.set_external_force_and_torque(
+                forces=forces,
+                torques=torques,
+                env_ids=torch.tensor([eid], device=device),
+                body_ids=torch.tensor([isaac_body_id], device=device),
+            )
+
+    def _get_buffers(self) -> dict[str, list[np.ndarray]] | None:
+        return self._record_buffers

holosoma/agents/callbacks/push.py

@@ -0,0 +1,263 @@
+"""Eval callback that applies random force pulses to test policy robustness.
+
+Periodically pushes the robot at random body locations with configurable
+force magnitude, duration, and interval. Uses IsaacLab's
+set_external_force_and_torque API to apply forces in the physics simulation.
+
+Push events are optionally recorded for visualization in viser by sharing
+buffer/metadata dicts with EvalRecordingCallback.
+"""
+
+from __future__ import annotations
+
+import math
+import random
+from typing import Any
+
+import numpy as np
+from loguru import logger
+
+from holosoma.agents.callbacks.base_callback import RLEvalCallback
+from holosoma.config_types.eval_callback import PushConfig
+from holosoma.utils.safe_torch_import import torch
+
+
+class EvalPushCallback(RLEvalCallback):
+    """Applies random force pulses during evaluation and optionally records them.
+
+    Forces are applied at the physics substep level by patching the env's
+    _apply_force_in_physics_step method, ensuring the force is active for the
+    full duration across all substeps (control decimation).
+
+    Recorded arrays (per env step):
+        push_active    [T]    1.0 when pushing, 0.0 otherwise
+        push_force_w   [T, 3] force vector in world frame (zeros when inactive)
+        push_body_pos_w[T, 3] world position of pushed body (zeros when inactive)
+        push_body_idx  [T]    index into body_names (-1 when inactive)
+    """
+
+    def __init__(
+        self,
+        config: PushConfig,
+        training_loop: Any = None,
+    ):
+        super().__init__(config, training_loop)
+        self.env_id = config.env_id
+        self.candidate_body_names = [s.strip() for s in config.body_names.split(",")]
+        self._record_buffers: dict[str, list[np.ndarray]] | None = None
+        self._record_metadata: dict[str, Any] | None = None
+
+        # Resolved at on_pre_evaluate_policy
+        self._resolved_bodies: list[tuple[str, int, int]] = []
+        self._dt: float = 0.0
+        self._sim: Any = None
+        self._original_apply_force: Any = None
+
+        # Push state machine
+        self._push_active: bool = False
+        self._push_steps_remaining: int = 0
+        self._idle_steps_remaining: int = 0
+        self._current_force_w: torch.Tensor | None = None
+        self._current_body_names_idx: int = -1
+        self._current_isaac_body_id: int = -1
+
+        self._step_count: int = 0
+
+    def _get_env(self):
+        return self.training_loop._unwrap_env()
+
+    def _find_recording_callback(self):
+        from holosoma.agents.callbacks.recording import EvalRecordingCallback
+
+        for cb in self.training_loop.eval_callbacks:
+            if isinstance(cb, EvalRecordingCallback):
+                return cb
+        return None
+
+    def on_pre_evaluate_policy(self) -> None:
+        # Discover recording callback for shared buffer/metadata access
+        recording_cb = self._find_recording_callback()
+        if recording_cb is not None:
+            self._record_buffers = recording_cb._buffers
+            self._record_metadata = recording_cb._metadata
+
+        env = self._get_env()
+        self._sim = env.simulator
+        if self._sim.simulator_config.name != "isaacsim":
+            raise RuntimeError(f"EvalPushCallback requires IsaacSim, got '{self._sim.simulator_config.name}'")
+        self._dt = float(env.dt)
+
+        # Resolve body names to simulator indices
+        self._resolved_bodies = []
+        available = list(getattr(self._sim, "body_names", []))
+        for name in self.candidate_body_names:
+            body_names_idx = self._sim.find_rigid_body_indice(name)
+            if isinstance(body_names_idx, int) and body_names_idx >= 0:
+                isaac_body_id = self._sim.body_ids[body_names_idx]
+                self._resolved_bodies.append((name, body_names_idx, isaac_body_id))
+            else:
+                raise ValueError(f"EvalPushCallback: body '{name}' not found. Available: {available}")
+
+        resolved_names = [name for name, _, _ in self._resolved_bodies]
+        logger.info(
+            f"EvalPushCallback: {len(self._resolved_bodies)} push bodies: {resolved_names}, "
+            f"force={self.config.force_range}N, duration={self.config.duration_s}s, interval={self.config.interval_s}s"
+        )
+
+        # Store push config in recording metadata
+        if self._record_metadata is not None:
+            self._record_metadata["push_config"] = {
+                "body_names": resolved_names,
+                "force_range": list(self.config.force_range),
+                "duration_s": list(self.config.duration_s),
+                "interval_s": list(self.config.interval_s),
+            }
+
+        # Initialize recording buffers
+        buffers = self._get_buffers()
+        if buffers is not None:
+            for key in ("push_active", "push_force_w", "push_body_pos_w", "push_body_idx"):
+                buffers[key] = []
+
+        # Start with a random idle period
+        self._idle_steps_remaining = self._sample_interval_steps()
+        self._push_active = False
+
+        # Patch _apply_force_in_physics_step to inject push forces at every substep
+        self._original_apply_force = env._apply_force_in_physics_step
+
+        def _patched_apply_force():
+            self._original_apply_force()
+            if self._push_active:
+                self._apply_push_force()
+
+        env._apply_force_in_physics_step = _patched_apply_force
+
+    def on_post_eval_env_step(self, actor_state: dict) -> dict:
+        # Manage push schedule
+        if self._push_active:
+            self._push_steps_remaining -= 1
+            if self._push_steps_remaining <= 0:
+                self._deactivate_push()
+                self._idle_steps_remaining = self._sample_interval_steps()
+        else:
+            self._idle_steps_remaining -= 1
+            if self._idle_steps_remaining <= 0:
+                self._activate_push()
+
+        buffers = self._get_buffers()
+        eid = self.env_id
+
+        if self._push_active:
+            assert self._current_force_w is not None
+            force_w = self._current_force_w.detach().cpu().numpy().copy()
+            body_pos = self._sim._robot.data.body_pos_w[eid, self._current_isaac_body_id].detach().cpu().numpy().copy()
+            body_idx = self._current_body_names_idx
+            active = 1.0
+        else:
+            force_w = np.zeros(3, dtype=np.float32)
+            body_pos = np.zeros(3, dtype=np.float32)
+            body_idx = -1
+            active = 0.0
+
+        if buffers is not None:
+            buffers["push_active"].append(np.array(active, dtype=np.float32))
+            buffers["push_force_w"].append(force_w.astype(np.float32))
+            buffers["push_body_pos_w"].append(body_pos.astype(np.float32))
+            buffers["push_body_idx"].append(np.array(body_idx, dtype=np.int32))
+
+        self._step_count += 1
+        return actor_state
+
+    def on_post_evaluate_policy(self) -> None:
+        if self._push_active:
+            self._deactivate_push()
+        self._get_env()._apply_force_in_physics_step = self._original_apply_force
+        logger.info(f"EvalPushCallback: completed {self._step_count} steps")
+
+    # ------------------------------------------------------------------
+    # Push lifecycle
+    # ------------------------------------------------------------------
+
+    def _activate_push(self) -> None:
+        name, body_names_idx, isaac_body_id = random.choice(self._resolved_bodies)
+        self._current_body_names_idx = body_names_idx
+        self._current_isaac_body_id = isaac_body_id
+
+        # Random force magnitude
+        mag = random.uniform(*self.config.force_range)
+
+        # Random horizontal direction (XY plane)
+        angle = random.uniform(0, 2 * math.pi)
+        direction = torch.tensor(
+            [math.cos(angle), math.sin(angle), 0.0],
+            dtype=torch.float32,
+            device=self._sim.sim_device,
+        )
+        self._current_force_w = direction * mag
+
+        # Duration in env steps
+        dur_s = random.uniform(*self.config.duration_s)
+        self._push_steps_remaining = max(1, round(dur_s / self._dt))
+
+        self._push_active = True
+        logger.info(
+            f"PUSH START body='{name}' force={mag:.0f}N "
+            f"dir=[{direction[0].item():.2f},{direction[1].item():.2f}] "
+            f"duration={self._push_steps_remaining} steps ({dur_s:.2f}s)"
+        )
+
+    def _deactivate_push(self) -> None:
+        self._clear_push_force()
+        self._push_active = False
+        self._current_force_w = None
+        self._current_body_names_idx = -1
+        self._current_isaac_body_id = -1
+        logger.info("PUSH END")
+
+    # ------------------------------------------------------------------
+    # Force application (called each physics substep via monkey-patch)
+    # ------------------------------------------------------------------
+
+    def _apply_push_force(self) -> None:
+        from isaaclab.utils.math import quat_apply_inverse
+
+        eid = self.env_id
+        isaac_body_id = self._current_isaac_body_id
+
+        # Transform world-frame force to body-local frame
+        # (IsaacLab 2.1 hardcodes is_global=False)
+        body_quat_w = self._sim._robot.data.body_quat_w[eid, isaac_body_id]
+        force_body = quat_apply_inverse(body_quat_w, self._current_force_w)
+
+        forces = force_body.unsqueeze(0).unsqueeze(0)  # [1, 1, 3]
+        torques = torch.zeros_like(forces)  # [1, 1, 3]
+
+        self._sim._robot.set_external_force_and_torque(
+            forces=forces,
+            torques=torques,
+            env_ids=torch.tensor([eid], device=self._sim.sim_device),
+            body_ids=torch.tensor([isaac_body_id], device=self._sim.sim_device),
+        )
+
+    def _clear_push_force(self) -> None:
+        if self._current_isaac_body_id < 0:
+            return
+        zero = torch.zeros(1, 1, 3, device=self._sim.sim_device)
+        self._sim._robot.set_external_force_and_torque(
+            forces=zero,
+            torques=zero,
+            env_ids=torch.tensor([self.env_id], device=self._sim.sim_device),
+            body_ids=torch.tensor([self._current_isaac_body_id], device=self._sim.sim_device),
+        )
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _sample_interval_steps(self) -> int:
+        s = random.uniform(*self.config.interval_s)
+        return max(1, round(s / self._dt))
+
+    def _get_buffers(self) -> dict[str, list[np.ndarray]] | None:
+        return self._record_buffers