physicausal 0.1.0__py3-none-any.whl

physicausal/__init__.py

@@ -0,0 +1,25 @@
+from .causal.validator import CausalValidator
+from .envs.pendulum import PendulumEnv, PendulumState
+from .envs.simple_push import ObjectState, SimplePushEnv
+from .models.beta_tc_vae import BetaTCVAE
+from .models.vae import BetaVAE, WorldModel
+from .training.agent import CausalLearningAgent, ReplayBuffer, train
+from .utils.tracker import NoOpTracker, Tracker
+
+__all__ = [
+    "BetaTCVAE",
+    "BetaVAE",
+    "CausalLearningAgent",
+    "CausalValidator",
+    "NoOpTracker",
+    "ObjectState",
+    "PendulumEnv",
+    "PendulumState",
+    "ReplayBuffer",
+    "SimplePushEnv",
+    "Tracker",
+    "WorldModel",
+    "train",
+]
+
+__version__ = "0.1.0"

physicausal/benchmarks/__init__.py

@@ -0,0 +1,3 @@
+"""
+基准测试适配器，支持与其他因果世界模型工具的数据格式互转。
+"""

physicausal/benchmarks/causalverse.py

@@ -0,0 +1,142 @@
+"""
+CausalVerse data format adapter.
+
+CausalVerse (NeurIPS 2025 Spotlight) is a comprehensive causal representation learning benchmark,
+containing 10 domains, 1024x1024 images, 300M+ video frames.
+
+Its data structure:
+    - images: (T, H, W, 3) RGB sequence
+    - factors: dict of ground-truth causal factors
+    - actions: (T-1, action_dim) action sequence
+
+This adapter provides:
+    1. Converting CausalVerse video data to physicausal (obs, action, next_obs) format
+    2. Converting physicausal evaluation results to CausalVerse standard metrics (MCC, R2)
+
+Note: CausalVerse requires additional installation (not a physicausal core dependency).
+"""
+
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+
+
+def convert_causalverse_sequence(
+    images: np.ndarray,
+    actions: np.ndarray,
+    factors: dict[str, np.ndarray] | None = None,
+    resize_fn: Callable | None = None,
+) -> list[tuple[dict, np.ndarray, dict]]:
+    """
+    Convert CausalVerse video sequence to physicausal triplet list.
+
+    Parameters
+    ----------
+    images : np.ndarray
+        Image sequence, shape (T, H, W, 3).
+    actions : np.ndarray
+        Action sequence, shape (T-1, action_dim).
+    factors : dict, optional
+        Ground-truth causal factors for validation.
+    resize_fn : callable, optional
+        Image resize function for dimensionality reduction.
+
+    Returns
+    -------
+    data : list
+        physicausal-format (obs, action, next_obs) list.
+    """
+    T = images.shape[0]
+    data = []
+
+    for t in range(T - 1):
+        obs_b = _image_to_obs(images[t], factors, t, resize_fn)
+        obs_a = _image_to_obs(images[t + 1], factors, t + 1, resize_fn)
+        action = np.array(actions[t], dtype=np.float32)
+        data.append((obs_b, action, obs_a))
+
+    return data
+
+
+def _image_to_obs(
+    image: np.ndarray,
+    factors: dict | None,
+    frame_idx: int,
+    resize_fn: Callable | None,
+) -> dict:
+    """Convert single frame to physicausal observation format."""
+    if resize_fn is not None:
+        image = resize_fn(image)
+
+    visual = image.flatten().astype(np.float32) / 255.0
+
+    obs: dict[str, Any] = {"visual": visual}
+
+    if factors is not None:
+        # Extract factor values for current frame
+        for key, values in factors.items():
+            if frame_idx < len(values):
+                obs[key] = values[frame_idx]
+
+    return obs
+
+
+def compute_causalverse_metrics(
+    zs: np.ndarray,
+    ground_truth_factors: np.ndarray,
+) -> dict[str, float]:
+    """
+    Compute CausalVerse standard metrics.
+
+    Parameters
+    ----------
+    zs : np.ndarray
+        Latent variables, shape (n_samples, latent_dim).
+    ground_truth_factors : np.ndarray
+        Ground-truth causal factors, shape (n_samples, n_factors).
+
+    Returns
+    -------
+    metrics : dict
+        Contains MCC (Mean Correlation Coefficient) and R2.
+    """
+    from scipy import stats
+
+    n_factors = ground_truth_factors.shape[1]
+    n_dims = zs.shape[1]
+
+    # Compute correlation matrix for all (factor, latent_dim) pairs
+    corr_matrix = np.zeros((n_factors, n_dims))
+    for i in range(n_factors):
+        for j in range(n_dims):
+            r, _ = stats.pearsonr(ground_truth_factors[:, i], zs[:, j])
+            corr_matrix[i, j] = abs(r)
+
+    # MCC: optimal one-to-one matching between factors and latent dims
+    # (Hungarian algorithm on the |correlation| matrix, standard in CRL literature)
+    from scipy.optimize import linear_sum_assignment
+
+    rows, cols = linear_sum_assignment(-corr_matrix)
+    matched_corrs = [corr_matrix[i, j] for i, j in zip(rows, cols)]
+
+    mcc = np.mean(matched_corrs) if matched_corrs else 0.0
+
+    # R2: linear regression fit
+    from sklearn.linear_model import LinearRegression
+    from sklearn.metrics import r2_score
+
+    r2s = []
+    for i in range(n_factors):
+        reg = LinearRegression()
+        reg.fit(zs, ground_truth_factors[:, i])
+        pred = reg.predict(zs)
+        r2s.append(r2_score(ground_truth_factors[:, i], pred))
+
+    avg_r2 = np.mean(r2s) if r2s else 0.0
+
+    return {
+        "mcc": float(mcc),
+        "r2": float(avg_r2),
+        "matched_corrs": matched_corrs,
+    }

physicausal/benchmarks/causalworld.py

@@ -0,0 +1,122 @@
+"""
+CausalWorld data format adapter.
+
+CausalWorld (https://github.com/rr-learning/CausalWorld) uses PyBullet
+for robot manipulation simulation. Its data structure is:
+    - observations: dict with 'goal', 'mask', 'rgb', 'depth'
+    - actions: continuous action vector
+    - rewards: scalar reward
+
+This adapter provides:
+    1. Converting CausalWorld data to physicausal format
+    2. Converting physicausal evaluation results to CausalWorld metrics
+
+Note: CausalWorld requires additional installation (not a physicausal core dependency).
+"""
+
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+
+
+def convert_causalworld_triplet(
+    obs_before: dict,
+    action: np.ndarray,
+    obs_after: dict,
+    extract_visual_fn: Callable | None = None,
+) -> tuple[dict, np.ndarray, dict]:
+    """
+    Convert one CausalWorld (obs, action, next_obs) to physicausal format.
+
+    Parameters
+    ----------
+    obs_before : dict
+        CausalWorld observation containing 'rgb', 'depth', 'mask', etc.
+    action : np.ndarray
+        Action vector.
+    obs_after : dict
+        Next observation.
+    extract_visual_fn : callable, optional
+        Custom visual feature extraction function. Defaults to flattening RGB.
+
+    Returns
+    -------
+    (obs_before_fmt, action_fmt, obs_after_fmt) : tuple
+        physicausal-format triplet.
+    """
+    if extract_visual_fn is None:
+        # Default: flatten RGB image as visual feature
+        def extract_visual_fn(obs: dict) -> np.ndarray:
+            rgb = obs.get("rgb", obs.get("observation", {}).get("rgb"))
+            if rgb is None:
+                raise ValueError("Cannot find 'rgb' in observation")
+            return rgb.flatten().astype(np.float32) / 255.0
+
+    obs_b_fmt = {
+        "visual": extract_visual_fn(obs_before),
+        "position": obs_before.get("position", np.zeros(3)),
+    }
+    obs_a_fmt = {
+        "visual": extract_visual_fn(obs_after),
+        "position": obs_after.get("position", np.zeros(3)),
+    }
+
+    return obs_b_fmt, np.array(action, dtype=np.float32), obs_a_fmt
+
+
+def causalworld_to_physicausal(
+    dataset: list[dict],
+    extract_visual_fn: Callable | None = None,
+) -> list[tuple[dict, np.ndarray, dict]]:
+    """
+    Batch convert CausalWorld dataset to physicausal format.
+
+    Parameters
+    ----------
+    dataset : list[dict]
+        CausalWorld dataset, each item is {"obs": ..., "action": ..., "next_obs": ...}.
+
+    Returns
+    -------
+    data : list
+        physicausal-format data list.
+    """
+    data = []
+    for item in dataset:
+        obs_b = item.get("obs", item.get("observation"))
+        action = item.get("action")
+        obs_a = item.get("next_obs", item.get("next_observation"))
+        triplet = convert_causalworld_triplet(obs_b, action, obs_a, extract_visual_fn)
+        data.append(triplet)
+    return data
+
+
+def compute_causalworld_metrics(
+    validator_results: dict[str, Any],
+) -> dict[str, float]:
+    """
+    Convert physicausal validation results to CausalWorld-style metrics.
+
+    CausalWorld uses the following metrics:
+        - success_rate: task success rate
+        - intervention_success: intervention success rate
+        - disentanglement_score: disentanglement score
+
+    This function provides approximate mapping.
+    """
+    results = validator_results.get("results", {})
+
+    # Compute average causal score
+    causal_scores = [
+        res["best_abs_corr"] for res in results.values() if res.get("is_causal")
+    ]
+
+    avg_corr = np.mean(causal_scores) if causal_scores else 0.0
+    n_causal = sum(1 for r in results.values() if r.get("is_causal"))
+
+    return {
+        "disentanglement_score": avg_corr,
+        "n_causal_properties": n_causal,
+        "total_properties": len(results),
+    }

physicausal/causal/__init__.py

@@ -0,0 +1,4 @@
+from .intervention import bootstrap_ci
+from .validator import CausalValidator
+
+__all__ = ["CausalValidator", "bootstrap_ci"]

physicausal/causal/intervention.py

@@ -0,0 +1,47 @@
+"""
+因果干预与评估工具。
+
+此模块保留向后兼容的 API，主要功能已迁移到 CausalValidator。
+"""
+
+from collections.abc import Callable
+
+import numpy as np
+
+
+def bootstrap_ci(
+    data: np.ndarray,
+    statistic_fn: Callable[[np.ndarray], float],
+    n_bootstrap: int = 1000,
+    alpha: float = 0.05,
+    seed: int = 42,
+) -> tuple[float, float]:
+    """
+    Bootstrap 置信区间估计。
+
+    Parameters
+    ----------
+    data : np.ndarray
+        原始样本数据。
+    statistic_fn : callable
+        统计量计算函数，输入数组返回标量。
+    n_bootstrap : int
+        Bootstrap 重采样次数。
+    alpha : float
+        显著性水平。
+    seed : int
+        随机种子。
+
+    Returns
+    -------
+    (lower, upper) : 置信区间上下界。
+    """
+    rng = np.random.RandomState(seed)
+    n = len(data)
+    stats_ = np.zeros(n_bootstrap)
+    for i in range(n_bootstrap):
+        sample = data[rng.randint(0, n, n)]
+        stats_[i] = statistic_fn(sample)
+    lower = np.percentile(stats_, 100 * alpha / 2)
+    upper = np.percentile(stats_, 100 * (1 - alpha / 2))
+    return float(lower), float(upper)

physicausal/causal/validator.py

@@ -0,0 +1,315 @@
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+from scipy import stats
+from sklearn.feature_selection import mutual_info_regression
+
+if TYPE_CHECKING:
+    from ..models.base import CausalWorldModel
+
+
+class CausalValidator:
+    """
+    Causal validator: evaluates whether a world model's latent variables capture real physical causal properties.
+
+    Supports the following validation methods:
+    1. Pearson correlation + permutation test (statistical significance)
+    2. Mutual information (non-linear association)
+    3. Do-operator intervention (causal consistency)
+    4. Sensitivity analysis (gradient attribution)
+    5. DCS (Disentanglement Completeness Score)
+    """
+
+    def __init__(self, model: "CausalWorldModel", env):
+        self.model = model
+        self.env = env
+
+    # ------------------------------------------------------------------
+    # Encoding
+    # ------------------------------------------------------------------
+
+    def encode_objects(
+        self, objects: list, batch_size: int = 64
+    ) -> np.ndarray:
+        """
+        Encode a list of objects into a latent variable matrix.
+
+        Each object is probed with a standardized interaction (``env.probe``) whose
+        response is a deterministic function of the object's physical properties.
+        The probe observations are encoded and averaged into a single latent vector.
+
+        For probabilistic models the posterior mean (``encode_dist``) is used instead
+        of a sampled ``z``, avoiding sampling noise in the correlation estimates.
+
+        Returns
+        -------
+        zs : np.ndarray
+            shape (n_objects, latent_dim)
+        """
+        device = next(self.model.parameters()).device
+        zs: list[np.ndarray] = []
+
+        for i in range(0, len(objects), batch_size):
+            batch = objects[i : i + batch_size]
+            batch_zs: list[np.ndarray] = []
+            for obj in batch:
+                obj_idx = self.env.objects.index(obj)
+                obs_list = self.env.probe(obj_idx)
+                obs_t = torch.tensor(
+                    np.stack([o["visual"] for o in obs_list]),
+                    dtype=torch.float32,
+                    device=device,
+                )
+                with torch.no_grad():
+                    if hasattr(self.model, "encode_dist"):
+                        mu, _ = self.model.encode_dist(obs_t)
+                        z_obj = mu.mean(dim=0)
+                    else:
+                        z_obj = self.model.encode(obs_t).mean(dim=0)
+                    batch_zs.append(z_obj.cpu().numpy())
+            zs.append(np.stack(batch_zs, axis=0))
+
+        return np.concatenate(zs, axis=0)
+
+    # ------------------------------------------------------------------
+    # Single property validation
+    # ------------------------------------------------------------------
+
+    def test_property(
+        self,
+        objects: list,
+        property_name: str = "mass",
+        batch_size: int = 64,
+        n_perm: int = 5000,
+    ) -> dict[str, Any]:
+        """
+        Test association between latent variables and a physical property.
+
+        Returns
+        -------
+        dict containing:
+            - best_dim: index of most correlated latent dimension
+            - best_corr: Pearson r
+            - best_abs_corr: |r|
+            - best_pvalue: two-tailed p
+            - permutation_pvalue: permutation test p
+            - best_mi: mutual information
+            - is_causal: passes causal threshold (|r| > 0.5 and perm_p < 0.05)
+            - all_corrs: list of correlations for all dimensions
+        """
+        zs = self.encode_objects(objects, batch_size=batch_size)
+        properties = np.array(
+            [getattr(obj, property_name) for obj in objects], dtype=np.float64
+        )
+
+        n_dims = zs.shape[1]
+        correlations = []
+        for i in range(n_dims):
+            r, p = stats.pearsonr(zs[:, i], properties)
+            correlations.append((i, float(r), float(abs(r)), float(p)))
+
+        # Sort by |r|
+        correlations.sort(key=lambda x: x[2], reverse=True)
+        best = correlations[0]
+
+        perm_p = self._permutation_test(zs[:, best[0]], properties, n_perm=n_perm)
+
+        n_neighbors = min(3, max(1, len(objects) - 1))
+        mi = mutual_info_regression(
+            zs[:, best[0] : best[0] + 1], properties, n_neighbors=n_neighbors
+        )[0]
+
+        return {
+            "property": property_name,
+            "best_dim": best[0],
+            "best_corr": best[1],
+            "best_abs_corr": best[2],
+            "best_pvalue": best[3],
+            "permutation_pvalue": perm_p,
+            "best_mi": float(mi),
+            "is_causal": best[2] > 0.5 and perm_p < 0.05,
+            "n_objects": len(objects),
+            "n_dims": n_dims,
+            "all_corrs": correlations,
+        }
+
+    def test_multiple_properties(
+        self,
+        objects: list,
+        properties: list[str],
+        batch_size: int = 64,
+        n_perm: int = 5000,
+    ) -> dict[str, Any]:
+        """Validate multiple physical properties in batch."""
+        results = {}
+        for prop in properties:
+            results[prop] = self.test_property(
+                objects, prop, batch_size, n_perm
+            )
+        return {"results": results, "properties": properties}
+
+    # ------------------------------------------------------------------
+    # Do-operator intervention
+    # ------------------------------------------------------------------
+
+    def intervene(
+        self,
+        obj,
+        latent_idx: int,
+        value: float,
+        action: list[float],
+    ) -> dict[str, Any]:
+        """
+        Perform do-operator intervention on latent variable.
+
+        Parameters
+        ----------
+        obj : ObjectState
+            Target object.
+        latent_idx : int
+            Latent dimension index to intervene on (note: not latent_dim size).
+        value : float
+            Intervention value.
+        action : list[float]
+            Action vector.
+
+        Returns
+        -------
+        dict containing predicted displacement, original/intervened latents, etc.
+        """
+        device = next(self.model.parameters()).device
+
+        with torch.no_grad():
+            obs = self.env._get_observation(obj)
+            obs_t = torch.tensor(
+                obs["visual"], dtype=torch.float32, device=device
+            ).unsqueeze(0)
+            z_orig = self.model.encode(obs_t).squeeze(0).cpu().numpy()
+
+        z_int = z_orig.copy()
+        z_int[latent_idx] = float(value)
+
+        with torch.no_grad():
+            z_t = torch.tensor(
+                z_int, dtype=torch.float32, device=device
+            ).unsqueeze(0)
+            a_t = torch.tensor(
+                action, dtype=torch.float32, device=device
+            ).unsqueeze(0)
+            pred_obs = self.model.decode(z_t, a_t)
+
+        pred_disp = float(
+            torch.norm(
+                pred_obs.squeeze()[:3]
+                - torch.tensor(obj.position, device=device)
+            ).item()
+        )
+
+        return {
+            "predicted_displacement": pred_disp,
+            "z_orig": z_orig,
+            "z_intervened": z_int,
+            "intervention_dim": latent_idx,
+            "intervention_value": value,
+        }
+
+    def sweep_interventions(
+        self, obj, action: list[float], values: list[float] | None = None
+    ) -> list[dict]:
+        """
+        Sweep interventions across all latent dimensions.
+
+        Parameters
+        ----------
+        values : list[float], optional
+            Intervention value for each dimension. Defaults to 2.0 for all.
+        """
+        n_dims = self.model.latent_dim
+        if values is None:
+            values = [2.0] * n_dims
+        return [
+            self.intervene(obj, d, values[d], action) for d in range(n_dims)
+        ]
+
+    # ------------------------------------------------------------------
+    # Sensitivity analysis
+    # ------------------------------------------------------------------
+
+    def sensitivity_analysis(self, obj) -> np.ndarray:
+        """
+        Compute gradient sensitivity of each latent dimension to predicted displacement.
+
+        Returns
+        -------
+        grads : np.ndarray
+            shape (latent_dim,), mean absolute gradient per dimension.
+        """
+        device = next(self.model.parameters()).device
+        obs = self.env._get_observation(obj)
+        obs_t = torch.tensor(
+            obs["visual"], dtype=torch.float32, device=device
+        ).unsqueeze(0)
+
+        with torch.no_grad():
+            z = self.model.encode(obs_t).clone().detach().requires_grad_(True)
+
+        action = torch.zeros(1, self.model.action_dim, device=device)
+        pred = self.model.decode(z, action)
+        loss = torch.norm(pred[:, :3] - obs_t[:, :3])
+        loss.backward()
+
+        return z.grad.abs().squeeze(0).cpu().numpy()
+
+    # ------------------------------------------------------------------
+    # DCS (Disentanglement Completeness Score)
+    # ------------------------------------------------------------------
+
+    def compute_dcs(
+        self, objects: list, property_name: str = "mass"
+    ) -> float:
+        """
+        Compute Disentanglement Completeness Score (DCS).
+
+        DCS = 1 - (second_best_correlation / best_correlation)
+
+        Value closer to 1 means the property is more concentrated in a single latent dimension.
+        This is a lightweight custom proxy for disentanglement (not the literature-standard DCI).
+        """
+        result = self.test_property(objects, property_name)
+        all_corrs = result["all_corrs"]
+        if len(all_corrs) < 2:
+            return 0.0
+
+        # Top two by |r|
+        best_abs = all_corrs[0][2]
+        second_abs = all_corrs[1][2]
+
+        if best_abs < 1e-8:
+            return 0.0
+
+        dcs = 1.0 - (second_abs / best_abs)
+        return float(max(0.0, dcs))
+
+    # ------------------------------------------------------------------
+    # Internal utilities
+    # ------------------------------------------------------------------
+
+    def _permutation_test(
+        self, z_col: np.ndarray, props: np.ndarray, n_perm: int = 5000
+    ) -> float:
+        """Permutation test for Pearson correlation significance."""
+        observed_r = abs(float(np.corrcoef(z_col, props)[0, 1]))
+        if np.isnan(observed_r):
+            return 1.0
+
+        rng = np.random.RandomState(42)
+        null = np.array(
+            [
+                abs(
+                    float(np.corrcoef(z_col, rng.permutation(props))[0, 1])
+                )
+                for _ in range(n_perm)
+            ]
+        )
+        return float(max((null >= observed_r).mean(), 1.0 / n_perm))

physicausal/envs/__init__.py

@@ -0,0 +1,4 @@
+from .pendulum import PendulumEnv, PendulumState
+from .simple_push import ObjectState, SimplePushEnv
+
+__all__ = ["ObjectState", "PendulumEnv", "PendulumState", "SimplePushEnv"]