Repository: DoubleML/doubleml-for-py (language: python)

Issue:
[Feature Request]: Support for multiple discrete treatments in `doubleml.DoubleMLIRM` 
### Describe the feature you want to propose or implement

Would be nice to extend the `DoubleMLIRM` class to accommodate multiple discrete-valued treatments $w \in \mathcal{W}$, where the estimation primatives are the marginal counterfactual mean $Y^{(w)}$ under each treatment, and one can construct treatment effects as contrasts between them. The marginal mean construction is the same as in the binary treatment case

$$
\phi_i^{(w)} = \mu^{(w)}(x) + \frac{1(W_i = w)}{\pi^{w}(x)} \left(Y_i - \mu^{(w)}(x)\right)
$$

which can be constructed at the unit level, and the subsequent analysis can be performed as usual (the binary influence function would be a special case where one constructs $\psi_i = \phi_i^{(1)} - \phi_i^{(0)} - \theta$ and takes the mean and variance for point estimate and standard error. 

---
Nuances:
 - ATT isn't a well defined concept for multiple-treatments (or at the very least one has to stipulate the subpopulation that is allowed to contribute to the estimator of $E[Y^({0}) \mid W = 1]$, typically the 'pure control' (0-valued by convention) subpopulation might be reasonable)


### Propose a possible solution or implementation


This approach is implemented in [Ed Kennedy's package](https://github.com/ehkennedy/npcausal/blob/master/R/ate.R), and I have a minimal implementation in [this](https://github.com/apoorvalal/aipyw) library. Wondering if there's interest in accommodating this use case (since this will involve changing `IIRM` from [`check_data`](https://github.com/DoubleML/doubleml-for-py/blob/main/doubleml/irm/irm.py#L246) onwards.


### Did you consider alternatives to the proposed solution. If yes, please describe

_No response_

### Comments, context or references

_No response_


Discussion / hints:
Thanks @apoorvalal  for adding this feature request. We are working on an extension to the multi-treatment IRM model. Also thanks for the hint for the ATET in multi-treatment, makes sense 😄 

We are currently working on various extensions including a change of the DoubleML interface to make it possible to better combine various DoubleML objects... The development is already pretty far but we are facing a couple of issues regarding parallel changes (new model, also docu updates etc). Maybe @SvenKlaassen can add more info to what extent multi-level treatment IRM is already possible with the dev version. Anyways, I think the refactoring will be completed in the next months.

We'll add an update and an example here
Yes, i think this would be a great addition. 
The with the new `DoubleMLFramework` class it should be possible to substract two potential outcome estimates, such that an implementation for average potential outcomes $\mathbb{E}[Y(a)]$ for action/treatment $a\in A$, would be a first step. Then a simple class for multiple levels (like `DoubleMLQTE`, see https://github.com/DoubleML/doubleml-for-py/blob/main/doubleml/irm/qte.py) could handle the comparisons on different levels. 

Fix patch:
diff --git a/doubleml/__init__.py b/doubleml/__init__.py
index 700991c7..c97bddf7 100644
--- a/doubleml/__init__.py
+++ b/doubleml/__init__.py
@@ -5,6 +5,8 @@
 from .plm.plr import DoubleMLPLR
 from .plm.pliv import DoubleMLPLIV
 from .irm.irm import DoubleMLIRM
+from .irm.apo import DoubleMLAPO
+from .irm.apos import DoubleMLAPOS
 from .irm.iivm import DoubleMLIIVM
 from .double_ml_data import DoubleMLData, DoubleMLClusterData
 from .did.did import DoubleMLDID
@@ -18,22 +20,26 @@
 from .utils.blp import DoubleMLBLP
 from .utils.policytree import DoubleMLPolicyTree
 
-__all__ = ['concat',
-           'DoubleMLFramework',
-           'DoubleMLPLR',
-           'DoubleMLPLIV',
-           'DoubleMLIRM',
-           'DoubleMLIIVM',
-           'DoubleMLData',
-           'DoubleMLClusterData',
-           'DoubleMLDID',
-           'DoubleMLDIDCS',
-           'DoubleMLPQ',
-           'DoubleMLQTE',
-           'DoubleMLLPQ',
-           'DoubleMLCVAR',
-           'DoubleMLBLP',
-           'DoubleMLPolicyTree',
-           'DoubleMLSSM']
+__all__ = [
+    'concat',
+    'DoubleMLFramework',
+    'DoubleMLPLR',
+    'DoubleMLPLIV',
+    'DoubleMLIRM',
+    'DoubleMLAPO',
+    'DoubleMLAPOS',
+    'DoubleMLIIVM',
+    'DoubleMLData',
+    'DoubleMLClusterData',
+    'DoubleMLDID',
+    'DoubleMLDIDCS',
+    'DoubleMLPQ',
+    'DoubleMLQTE',
+    'DoubleMLLPQ',
+    'DoubleMLCVAR',
+    'DoubleMLBLP',
+    'DoubleMLPolicyTree',
+    'DoubleMLSSM'
+]
 
 __version__ = importlib.metadata.version('doubleml')
diff --git a/doubleml/datasets.py b/doubleml/datasets.py
index 21585f15..a3be46ce 100644
--- a/doubleml/datasets.py
+++ b/doubleml/datasets.py
@@ -1485,3 +1485,162 @@ def make_ssm_data(n_obs=8000, dim_x=100, theta=1, mar=True, return_type='DoubleM
             return DoubleMLData(data, 'y', 'd', x_cols, 'z', None, 's')
     else:
         raise ValueError('Invalid return_type.')
+
+
+def make_irm_data_discrete_treatments(n_obs=200, n_levels=3, linear=False, random_state=None, **kwargs):
+    """
+    Generates data from a interactive regression (IRM) model with multiple treatment levels (based on an
+    underlying continous treatment).
+
+    The data generating process is defined as follows (similar to the Monte Carlo simulation used
+    in Sant'Anna and Zhao (2020)).
+
+    Let :math:`X= (X_1, X_2, X_3, X_4, X_5)^T \\sim \\mathcal{N}(0, \\Sigma)`, where  :math:`\\Sigma` corresponds
+    to the identity matrix.
+    Further, define :math:`Z_j = (\\tilde{Z_j} - \\mathbb{E}[\\tilde{Z}_j]) / \\sqrt{\\text{Var}(\\tilde{Z}_j)}`,
+    where
+
+    .. math::
+
+            \\tilde{Z}_1 &= \\exp(0.5 \\cdot X_1)
+
+            \\tilde{Z}_2 &= 10 + X_2/(1 + \\exp(X_1))
+
+            \\tilde{Z}_3 &= (0.6 + X_1 \\cdot X_3 / 25)^3
+
+            \\tilde{Z}_4 &= (20 + X_2 + X_4)^2
+
+            \\tilde{Z}_5 &= X_5.
+
+    A continuous treatment :math:`D_{\\text{cont}}` is generated as
+
+    .. math::
+
+        D_{\\text{cont}} = \\xi (-Z_1 + 0.5 Z_2 - 0.25 Z_3 - 0.1 Z_4) + \\varepsilon_D,
+
+    where :math:`\\varepsilon_D \\sim \\mathcal{N}(0,1)` and :math:`\\xi=0.3`. The corresponding treatment
+    effect is defined as
+
+    .. math::
+
+        \\text{\\theta}(d) = 0.1 \\exp(d) + 10 \\sin(0.7 d) + 2 d - 0.2 d^2.
+
+    Based on the continous treatment, a discrete treatment :math:`D` is generated as with a baseline level of
+    :math:`D=0` and additional levels based on the quantiles of :math:`D_{\\text{cont}}`. The number of levels
+    is defined by :math:`n_{\\text{levels}}`. Each level is chosen to have the same probability of being selected.
+
+    The potential outcomes are defined as
+
+    .. math::
+
+            Y(0) &= 210 + 27.4 Z_1 + 13.7 (Z_2 + Z_3 + Z_4) + \\varepsilon_Y
+
+            Y(1) &= \\text{\\theta}(D_{\\text{cont}}) 1\\{D_{\\text{cont}} > 0\\} + Y(0),
+
+    where :math:`\\varepsilon_Y \\sim \\mathcal{N}(0,5)`. Further, the observed outcome is defined as
+
+    .. math::
+
+        Y = Y(1) 1\\{D > 0\\} + Y(0) 1\\{D = 0\\}.
+
+    The data is returned as a dictionary with the entries ``x``, ``y``, ``d`` and ``oracle_values``.
+
+    Parameters
+    ----------
+    n_obs : int
+        The number of observations to simulate.
+        Default is ``200``.
+
+    n_levels : int
+        The number of treatment levels.
+        Default is ``3``.
+
+    linear : bool
+        Indicates whether the true underlying regression is linear.
+        Default is ``False``.
+
+    random_state : int
+        Random seed for reproducibility.
+        Default is ``42``.
+
+    Returns
+    -------
+    res_dict : dictionary
+       Dictionary with entries ``x``, ``y``, ``d`` and ``oracle_values``.
+
+    """
+    if random_state is not None:
+        np.random.seed(random_state)
+    xi = kwargs.get('xi', 0.3)
+    c = kwargs.get('c', 0.0)
+    dim_x = kwargs.get('dim_x', 5)
+
+    if not isinstance(n_levels, int):
+        raise ValueError('n_levels must be an integer.')
+    if n_levels < 2:
+        raise ValueError('n_levels must be at least 2.')
+
+    # observed covariates
+    cov_mat = toeplitz([np.power(c, k) for k in range(dim_x)])
+    x = np.random.multivariate_normal(np.zeros(dim_x), cov_mat, size=[n_obs, ])
+
+    def f_reg(w):
+        res = 210 + 27.4*w[:, 0] + 13.7*(w[:, 1] + w[:, 2] + w[:, 3])
+        return res
+
+    def f_treatment(w, xi):
+        res = xi * (-w[:, 0] + 0.5*w[:, 1] - 0.25*w[:, 2] - 0.1*w[:, 3])
+        return res
+
+    def treatment_effect(d, scale=15):
+        return scale * (1 / (1 + np.exp(-d - 1.2 * np.cos(d)))) - 2
+
+    z_tilde_1 = np.exp(0.5 * x[:, 0])
+    z_tilde_2 = 10 + x[:, 1] / (1 + np.exp(x[:, 0]))
+    z_tilde_3 = (0.6 + x[:, 0] * x[:, 2]/25)**3
+    z_tilde_4 = (20 + x[:, 1] + x[:, 3])**2
+
+    z_tilde = np.column_stack((z_tilde_1, z_tilde_2, z_tilde_3, z_tilde_4, x[:, 4:]))
+    z = (z_tilde - np.mean(z_tilde, axis=0)) / np.std(z_tilde, axis=0)
+
+    # error terms
+    var_eps_y = 5
+    eps_y = np.random.normal(loc=0, scale=np.sqrt(var_eps_y), size=n_obs)
+    var_eps_d = 1
+    eps_d = np.random.normal(loc=0, scale=np.sqrt(var_eps_d), size=n_obs)
+
+    if linear:
+        g = f_reg(x)
+        m = f_treatment(x, xi)
+    else:
+        assert not linear
+        g = f_reg(z)
+        m = f_treatment(z, xi)
+
+    cont_d = m + eps_d
+    level_bounds = np.quantile(cont_d, q=np.linspace(0, 1, n_levels + 1))
+    potential_level = sum([1.0 * (cont_d >= bound) for bound in level_bounds[1:-1]]) + 1
+    eta = np.random.uniform(0, 1, size=n_obs)
+    d = 1.0 * (eta >= 1/n_levels) * potential_level
+
+    ite = treatment_effect(cont_d)
+    y0 = g + eps_y
+    # only treated for d > 0 compared to the baseline
+    y = ite * (d > 0) + y0
+
+    oracle_values = {
+        'cont_d': cont_d,
+        'level_bounds': level_bounds,
+        'potential_level': potential_level,
+        'ite': ite,
+        'y0': y0,
+    }
+
+    resul_dict = {
+        'x': x,
+        'y': y,
+        'd': d,
+        'oracle_values': oracle_values
+    }
+
+    return resul_dict
diff --git a/doubleml/double_ml.py b/doubleml/double_ml.py
index c8bb7604..71f8b441 100644
--- a/doubleml/double_ml.py
+++ b/doubleml/double_ml.py
@@ -1219,6 +1219,8 @@ def set_sample_splitting(self, all_smpls, all_smpls_cluster=None):
         >>> ml_m = learner
         >>> obj_dml_data = make_plr_CCDDHNR2018(n_obs=10, alpha=0.5)
         >>> dml_plr_obj = dml.DoubleMLPLR(obj_dml_data, ml_g, ml_m)
+        >>> # simple sample splitting with two folds and without cross-fitting
+        >>> smpls = ([0, 1, 2, 3, 4], [5, 6, 7, 8, 9])
         >>> dml_plr_obj.set_sample_splitting(smpls)
         >>> # sample splitting with two folds and cross-fitting
         >>> smpls = [([0, 1, 2, 3, 4], [5, 6, 7, 8, 9]),
@@ -1434,44 +1436,11 @@ def sensitivity_summary(self):
         res : str
             Summary for the sensitivity analysis.
         """
-        header = '================== Sensitivity Analysis ==================\n'
-        if self.sensitivity_params is None:
-            res = header + 'Apply sensitivity_analysis() to generate sensitivity_summary.'
+        if self._framework is None:
+            raise ValueError('Apply sensitivity_analysis() before sensitivity_summary.')
         else:
-            sig_level = f'Significance Level: level={self.sensitivity_params["input"]["level"]}\n'
-            scenario_params = f'Sensitivity parameters: cf_y={self.sensitivity_params["input"]["cf_y"]}; ' \
-                              f'cf_d={self.sensitivity_params["input"]["cf_d"]}, ' \
-                              f'rho={self.sensitivity_params["input"]["rho"]}'
-
-            theta_and_ci_col_names = ['CI lower', 'theta lower', ' theta', 'theta upper', 'CI upper']
-            theta_and_ci = np.transpose(np.vstack((self.sensitivity_params['ci']['lower'],
-                                                   self.sensitivity_params['theta']['lower'],
-                                                   self.coef,
-                                                   self.sensitivity_params['theta']['upper'],
-                                                   self.sensitivity_params['ci']['upper'])))
-            df_theta_and_ci = pd.DataFrame(theta_and_ci,
-                                           columns=theta_and_ci_col_names,
-                                           index=self._dml_data.d_cols)
-            theta_and_ci_summary = str(df_theta_and_ci)
-
-            rvs_col_names = ['H_0', 'RV (%)', 'RVa (%)']
-            rvs = np.transpose(np.vstack((self.sensitivity_params['rv'],
-                                          self.sensitivity_params['rva']))) * 100
-
-            df_rvs = pd.DataFrame(np.column_stack((self.sensitivity_params["input"]["null_hypothesis"], rvs)),
-                                  columns=rvs_col_names,
-                                  index=self._dml_data.d_cols)
-            rvs_summary = str(df_rvs)
-
-            res = header + \
-                '\n------------------ Scenario          ------------------\n' + \
-                sig_level + scenario_params + '\n' + \
-                '\n------------------ Bounds with CI    ------------------\n' + \
-                theta_and_ci_summary + '\n' + \
-                '\n------------------ Robustness Values ------------------\n' + \
-                rvs_summary
-
-        return res
+            sensitivity_summary = self._framework.sensitivity_summary
+        return sensitivity_summary
 
     def sensitivity_plot(self, idx_treatment=0, value='theta', rho=1.0, level=0.95, null_hypothesis=0.0,
                          include_scenario=True, benchmarks=None, fill=True, grid_bounds=(0.15, 0.15), grid_size=100):
diff --git a/doubleml/double_ml_framework.py b/doubleml/double_ml_framework.py
index feedee41..b70a1623 100644
--- a/doubleml/double_ml_framework.py
+++ b/doubleml/double_ml_framework.py
@@ -9,6 +9,7 @@
 from .utils._estimation import _draw_weights, _aggregate_coefs_and_ses, _var_est
 from .utils._checks import _check_bootstrap, _check_framework_compatibility, _check_in_zero_one, \
     _check_float, _check_integer, _check_bool, _check_benchmarks
+from .utils._descriptive import generate_summary
 from .utils._plots import _sensitivity_contour_plot
 
 
@@ -17,10 +18,10 @@ class DoubleMLFramework():
 
     Parameters
     ----------
-   doubleml_dict : :dict
+    doubleml_dict : :dict
         A dictionary providing the estimated parameters and normalized scores. Keys have to be 'thetas', 'ses',
-         'all_thetas', 'all_ses', 'var_scaling_factors' and 'scaled_psi'.
-          Values have to be numpy arrays with the corresponding shapes.
+        'all_thetas', 'all_ses', 'var_scaling_factors' and 'scaled_psi'.
+        Values have to be numpy arrays with the corresponding shapes.
 
     """
 
@@ -57,6 +58,12 @@ def __init__(
 
         # check if all sizes match
         self._check_framework_shapes()
+
+        self._treatment_names = None
+        if 'treatment_names' in doubleml_dict.keys():
+            self._check_treatment_names(doubleml_dict['treatment_names'])
+            self._treatment_names = doubleml_dict['treatment_names']
+
         # initialize bootstrap distribution
         self._boot_t_stat = None
         self._boot_method = None
@@ -196,6 +203,77 @@ def sensitivity_params(self):
         """
         return self._sensitivity_params
 
+    @property
+    def treatment_names(self):
+        """
+        Names of the treatments.
+        """
+        return self._treatment_names
+
+    @treatment_names.setter
+    def treatment_names(self, value):
+        self._check_treatment_names(value)
+        self._treatment_names = value
+
+    @property
+    def summary(self):
+        """
+        A summary for the estimated causal parameters ``thetas``.
+        """
+        ci = self.confint()
+        df_summary = generate_summary(self.thetas, self.ses, self.t_stats,
+                                      self.pvals, ci, self._treatment_names)
+        return df_summary
+
+    @property
+    def sensitivity_summary(self):
+        """
+        Returns a summary for the sensitivity analysis after calling :meth:`sensitivity_analysis`.
+
+        Returns
+        -------
+        res : str
+            Summary for the sensitivity analysis.
+        """
+        header = '================== Sensitivity Analysis ==================\n'
+        if self.sensitivity_params is None:
+            res = header + 'Apply sensitivity_analysis() to generate sensitivity_summary.'
+        else:
+            sig_level = f'Significance Level: level={self.sensitivity_params["input"]["level"]}\n'
+            scenario_params = f'Sensitivity parameters: cf_y={self.sensitivity_params["input"]["cf_y"]}; ' \
+                              f'cf_d={self.sensitivity_params["input"]["cf_d"]}, ' \
+                              f'rho={self.sensitivity_params["input"]["rho"]}'
+
+            theta_and_ci_col_names = ['CI lower', 'theta lower', ' theta', 'theta upper', 'CI upper']
+            theta_and_ci = np.transpose(np.vstack((self.sensitivity_params['ci']['lower'],
+                                                   self.sensitivity_params['theta']['lower'],
+                                                   self.thetas,
+                                                   self.sensitivity_params['theta']['upper'],
+                                                   self.sensitivity_params['ci']['upper'])))
+            df_theta_and_ci = pd.DataFrame(theta_and_ci,
+                                           columns=theta_and_ci_col_names,
+                                           index=self.treatment_names)
+            theta_and_ci_summary = str(df_theta_and_ci)
+
+            rvs_col_names = ['H_0', 'RV (%)', 'RVa (%)']
+            rvs = np.transpose(np.vstack((self.sensitivity_params['rv'],
+                                          self.sensitivity_params['rva']))) * 100
+
+            df_rvs = pd.DataFrame(np.column_stack((self.sensitivity_params["input"]["null_hypothesis"], rvs)),
+                                  columns=rvs_col_names,
+                                  index=self.treatment_names)
+            rvs_summary = str(df_rvs)
+
+            res = header + \
+                '\n------------------ Scenario          ------------------\n' + \
+                sig_level + scenario_params + '\n' + \
+                '\n------------------ Bounds with CI    ------------------\n' + \
+                theta_and_ci_summary + '\n' + \
+                '\n------------------ Robustness Values ------------------\n' + \
+                rvs_summary
+
+        return res
+
     def __add__(self, other):
 
         if isinstance(other, DoubleMLFramework):
@@ -612,8 +690,11 @@ def confint(self, joint=False, level=0.95):
              self.all_thetas + self.all_ses * critical_values),
             axis=1)
         ci = np.median(self._all_cis, axis=2)
-        # TODO: add treatment names
         df_ci = pd.DataFrame(ci, columns=['{:.1f} %'.format(i * 100) for i in percentages])
+
+        if self._treatment_names is not None:
+            df_ci.set_index(pd.Index(self._treatment_names), inplace=True)
+
         return df_ci
 
     def bootstrap(self, method='normal', n_rep_boot=500):
@@ -944,6 +1025,19 @@ def _check_framework_shapes(self):
 
         return None
 
+    def _check_treatment_names(self, treatment_names):
+        if not isinstance(treatment_names, list):
+            raise TypeError('treatment_names must be a list. '
+                            f'Got {str(treatment_names)} of type {str(type(treatment_names))}.')
+        is_str = [isinstance(name, str) for name in treatment_names]
+        if not all(is_str):
+            raise TypeError('treatment_names must be a list of strings. '
+                            f'At least one element is not a string: {str(treatment_names)}.')
+        if len(treatment_names) != self._n_thetas:
+            raise ValueError('The length of treatment_names does not match the number of treatments. '
+                             f'Got {self._n_thetas} treatments and {len(treatment_names)} treatment names.')
+        return None
+
 
 def concat(objs):
     """
diff --git a/doubleml/irm/apo.py b/doubleml/irm/apo.py
new file mode 100644
index 00000000..93c3c0df
--- /dev/null
+++ b/doubleml/irm/apo.py
@@ -0,0 +1,457 @@
+import numpy as np
+import pandas as pd
+import warnings
+
+from sklearn.utils import check_X_y
+
+from ..double_ml import DoubleML
+
+from ..utils.blp import DoubleMLBLP
+from ..double_ml_score_mixins import LinearScoreMixin
+
+from ..utils._estimation import _dml_cv_predict, _dml_tune, _get_cond_smpls, _cond_targets, _trimm, \
+    _normalize_ipw
+from ..utils._checks import _check_score, _check_trimming, _check_weights, _check_finite_predictions, \
+    _check_is_propensity, _check_binary_predictions
+
+
+class DoubleMLAPO(LinearScoreMixin, DoubleML):
+    """Double machine learning average potential outcomes for interactive regression models.
+
+    Parameters
+    ----------
+    obj_dml_data : :class:`DoubleMLData` object
+        The :class:`DoubleMLData` object providing the data and specifying the variables for the causal model.
+
+    ml_g : estimator implementing ``fit()`` and ``predict()``
+        A machine learner implementing ``fit()`` and ``predict()`` methods (e.g.
+        :py:class:`sklearn.ensemble.RandomForestRegressor`) for the nuisance function :math:`g_0(D,X) = E[Y|X,D]`.
+        For a binary outcome variable :math:`Y` (with values 0 and 1), a classifier implementing ``fit()`` and
+        ``predict_proba()`` can also be specified. If :py:func:`sklearn.base.is_classifier` returns ``True``,
+        ``predict_proba()`` is used otherwise ``predict()``.
+
+    ml_m : classifier implementing ``fit()`` and ``predict_proba()``
+        A machine learner implementing ``fit()`` and ``predict_proba()`` methods (e.g.
+        :py:class:`sklearn.ensemble.RandomForestClassifier`) for the nuisance function :math:`m_0(X) = E[D|X]`.
+
+    treatment_level : int or float
+        Chosen treatment level for average potential outcomes.
+
+    n_folds : int
+        Number of folds.
+        Default is ``5``.
+
+    n_rep : int
+        Number of repetitons for the sample splitting.
+        Default is ``1``.
+
+    score : str or callable
+        A str (``'APO'``) specifying the score function.
+        Default is ``'APO'``.
+
+    weights : array, dict or None
+        An numpy array of weights for each individual observation. If None, then the ``'APO'`` score
+        is applied (corresponds to weights equal to 1).
+        An array has to be of shape ``(n,)``, where ``n`` is the number of observations.
+        A dictionary can be used to specify weights which depend on the treatment variable.
+        In this case, the dictionary has to contain two keys ``weights`` and ``weights_bar``, where the values
+        have to be arrays of shape ``(n,)`` and ``(n, n_rep)``.
+        Default is ``None``.
+
+    normalize_ipw : bool
+        Indicates whether the inverse probability weights are normalized.
+        Default is ``False``.
+
+    trimming_rule : str
+        A str (``'truncate'`` is the only choice) specifying the trimming approach.
+        Default is ``'truncate'``.
+
+    trimming_threshold : float
+        The threshold used for trimming.
+        Default is ``1e-2``.
+
+    draw_sample_splitting : bool
+        Indicates whether the sample splitting should be drawn during initialization of the object.
+        Default is ``True``.
+
+    """
+    def __init__(self,
+                 obj_dml_data,
+                 ml_g,
+                 ml_m,
+                 treatment_level,
+                 n_folds=5,
+                 n_rep=1,
+                 score='APO',
+                 weights=None,
+                 normalize_ipw=False,
+                 trimming_rule='truncate',
+                 trimming_threshold=1e-2,
+                 draw_sample_splitting=True):
+        super().__init__(obj_dml_data,
+                         n_folds,
+                         n_rep,
+                         score,
+                         draw_sample_splitting)
+
+        # set up treatment level and check data
+        self._treatment_level = treatment_level
+        self._treated = self._dml_data.d == self._treatment_level
+
+        self._check_data(self._dml_data)
+        valid_scores = ['APO']
+        _check_score(self.score, valid_scores, allow_callable=False)
+
+        # set stratication for resampling
+        self._strata = self._dml_data.d
+        if draw_sample_splitting:
+            self.draw_sample_splitting()
+
+        ml_g_is_classifier = self._check_learner(ml_g, 'ml_g', regressor=True, classifier=True)
+        _ = self._check_learner(ml_m, 'ml_m', regressor=False, classifier=True)
+        self._learner = {'ml_g': ml_g, 'ml_m': ml_m}
+        if ml_g_is_classifier:
+            if obj_dml_data.binary_outcome:
+                self._predict_method = {'ml_g': 'predict_proba', 'ml_m': 'predict_proba'}
+            else:
+                raise ValueError(f'The ml_g learner {str(ml_g)} was identified as classifier '
+                                 'but the outcome variable is not binary with values 0 and 1.')
+        else:
+            self._predict_method = {'ml_g': 'predict', 'ml_m': 'predict_proba'}
+        self._initialize_ml_nuisance_params()
+
+        self._normalize_ipw = normalize_ipw
+        if not isinstance(self.normalize_ipw, bool):
+            raise TypeError('Normalization indicator has to be boolean. ' +
+                            f'Object of type {str(type(self.normalize_ipw))} passed.')
+        self._trimming_rule = trimming_rule
+        self._trimming_threshold = trimming_threshold
+        _check_trimming(self._trimming_rule, self._trimming_threshold)
+
+        self._sensitivity_implemented = True
+        self._external_predictions_implemented = True
+
+        # APO weights
+        _check_weights(weights, score="ATE", n_obs=obj_dml_data.n_obs, n_rep=self.n_rep)
+        self._initialize_weights(weights)
+
+    @property
+    def treatment_level(self):
+        """
+        Chosen treatment level for average potential outcomes.
+        """
+        return self._treatment_level
+
+    @property
+    def treated(self):
+        """
+        Indicator for treated observations (with the corresponding treatment level).
+        """
+        return self._treated
+
+    @property
+    def normalize_ipw(self):
+        """
+        Indicates whether the inverse probability weights are normalized.
+        """
+        return self._normalize_ipw
+
+    @property
+    def trimming_rule(self):
+        """
+        Specifies the used trimming rule.
+        """
+        return self._trimming_rule
+
+    @property
+    def trimming_threshold(self):
+        """
+        Specifies the used trimming threshold.
+        """
+        return self._trimming_threshold
+
+    @property
+    def weights(self):
+        """
+        Specifies the weights for a weighted average potential outcome.
+        """
+        return self._weights
+
+    def _initialize_ml_nuisance_params(self):
+        valid_learner = ['ml_g0', 'ml_g1', 'ml_m']
+        self._params = {learner: {key: [None] * self.n_rep for key in self._dml_data.d_cols}
+                        for learner in valid_learner}
+
+    def _initialize_weights(self, weights):
+        if weights is None:
+            weights = np.ones(self._dml_data.n_obs)
+        if isinstance(weights, np.ndarray):
+            self._weights = {'weights': weights}
+        else:
+            assert isinstance(weights, dict)
+            self._weights = weights
+
+    def _get_weights(self):
+        # standard case for APO/ATE
+        weights = self._weights['weights']
+        if 'weights_bar' not in self._weights.keys():
+            weights_bar = self._weights['weights']
+        else:
+            weights_bar = self._weights['weights_bar'][:, self._i_rep]
+
+        return weights, weights_bar
+
+    def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=False):
+        x, y = check_X_y(self._dml_data.x, self._dml_data.y,
+                         force_all_finite=False)
+        # use the treated indicator to get the correct sample splits
+        x, treated = check_X_y(x, self.treated,
+                               force_all_finite=False)
+
+        # get train indices for d == treatment_level
+        smpls_d0, smpls_d1 = _get_cond_smpls(smpls, treated)
+        g0_external = external_predictions['ml_g0'] is not None
+        g1_external = external_predictions['ml_g1'] is not None
+        m_external = external_predictions['ml_m'] is not None
+
+        # nuisance g (g0 only relevant for sensitivity analysis)
+        if g0_external:
+            # use external predictions
+            g_hat0 = {'preds': external_predictions['ml_g0'],
+                      'targets': _cond_targets(y, cond_sample=(treated == 0)),
+                      'models': None}
+        else:
+            g_hat0 = _dml_cv_predict(self._learner['ml_g'], x, y, smpls=smpls_d0, n_jobs=n_jobs_cv,
+                                     est_params=self._get_params('ml_g0'), method=self._predict_method['ml_g'],
+                                     return_models=return_models)
+            _check_finite_predictions(g_hat0['preds'], self._learner['ml_g'], 'ml_g', smpls)
+            g_hat0['targets'] = _cond_targets(g_hat0['targets'], cond_sample=(treated == 0))
+
+        if self._dml_data.binary_outcome:
+            _check_binary_predictions(g_hat0['preds'], self._learner['ml_g'], 'ml_g', self._dml_data.y_col)
+
+        if g1_external:
+            # use external predictions
+            g_hat1 = {'preds': external_predictions['ml_g1'],
+                      'targets': _cond_targets(y, cond_sample=(treated == 1)),
+                      'models': None}
+        else:
+            g_hat1 = _dml_cv_predict(self._learner['ml_g'], x, y, smpls=smpls_d1, n_jobs=n_jobs_cv,
+                                     est_params=self._get_params('ml_g1'), method=self._predict_method['ml_g'],
+                                     return_models=return_models)
+            _check_finite_predictions(g_hat1['preds'], self._learner['ml_g'], 'ml_g', smpls)
+            # adjust target values to consider only compatible subsamples
+            g_hat1['targets'] = _cond_targets(g_hat1['targets'], cond_sample=(treated == 1))
+
+        if self._dml_data.binary_outcome:
+            _check_binary_predictions(g_hat1['preds'], self._learner['ml_g'], 'ml_g', self._dml_data.y_col)
+
+        # nuisance m
+        if m_external:
+            # use external predictions
+            m_hat = {'preds': external_predictions['ml_m'],
+                     'targets': treated,
+                     'models': None}
+        else:
+            m_hat = _dml_cv_predict(self._learner['ml_m'], x, treated, smpls=smpls, n_jobs=n_jobs_cv,
+                                    est_params=self._get_params('ml_m'), method=self._predict_method['ml_m'],
+                                    return_models=return_models)
+            _check_finite_predictions(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls)
+            _check_is_propensity(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls, eps=1e-12)
+
+        # also trimm external predictions
+        m_hat['preds'] = _trimm(m_hat['preds'], self.trimming_rule, self.trimming_threshold)
+
+        psi_a, psi_b = self._score_elements(y, treated, g_hat0['preds'], g_hat1['preds'],
+                                            m_hat['preds'], smpls)
+        psi_elements = {'psi_a': psi_a,
+                        'psi_b': psi_b}
+
+        preds = {'predictions': {'ml_g0': g_hat0['preds'],
+                                 'ml_g1': g_hat1['preds'],
+                                 'ml_m': m_hat['preds']},
+                 'targets': {'ml_g0': g_hat0['targets'],
+                             'ml_g1': g_hat1['targets'],
+                             'ml_m': m_hat['targets']},
+                 'models': {'ml_g0': g_hat0['models'],
+                            'ml_g1': g_hat1['models'],
+                            'ml_m': m_hat['models']}
+                 }
+        return psi_elements, preds
+
+    def _score_elements(self, y, treated, g_hat0, g_hat1, m_hat, smpls):
+        if self.normalize_ipw:
+            m_hat_adj = _normalize_ipw(m_hat, treated)
+        else:
+            m_hat_adj = m_hat
+
+        u_hat = y - g_hat1
+        weights, weights_bar = self._get_weights()
+        psi_b = weights * g_hat1 + weights_bar * np.divide(np.multiply(treated, u_hat), m_hat_adj)
+        psi_a = np.full_like(m_hat_adj, -1.0)
+
+        return psi_a, psi_b
+
+    def _sensitivity_element_est(self, preds):
+        # set elments for readability
+        y = self._dml_data.y
+        treated = self.treated
+
+        m_hat = preds['predictions']['ml_m']
+        g_hat0 = preds['predictions']['ml_g0']
+        g_hat1 = preds['predictions']['ml_g1']
+
+        weights, weights_bar = self._get_weights()
+
+        sigma2_score_element = np.square(y - np.multiply(treated, g_hat1) - np.multiply(1.0-treated, g_hat0))
+        sigma2 = np.mean(sigma2_score_element)
+        psi_sigma2 = sigma2_score_element - sigma2
+
+        # calc m(W,alpha) and Riesz representer
+        m_alpha = np.multiply(weights, np.multiply(weights_bar, np.divide(1.0, m_hat)))
+        rr = np.multiply(weights_bar, np.divide(treated, m_hat))
+
+        nu2_score_element = np.multiply(2.0, m_alpha) - np.square(rr)
+        nu2 = np.mean(nu2_score_element)
+        psi_nu2 = nu2_score_element - nu2
+
+        element_dict = {'sigma2': sigma2,
+                        'nu2': nu2,
+                        'psi_sigma2': psi_sigma2,
+                        'psi_nu2': psi_nu2,
+                        'riesz_rep': rr,
+                        }
+        return element_dict
+
+    def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv,
+                         search_mode, n_iter_randomized_search):
+        x, y = check_X_y(self._dml_data.x, self._dml_data.y,
+                         force_all_finite=False)
+        x, treated = check_X_y(x, self.treated,
+                               force_all_finite=False)
+        # get train indices for d == 0 and d == 1
+        smpls_d0, smpls_d1 = _get_cond_smpls(smpls, treated)
+
+        if scoring_methods is None:
+            scoring_methods = {'ml_g': None,
+                               'ml_m': None}
+
+        train_inds = [train_index for (train_index, _) in smpls]
+        train_inds_d0 = [train_index for (train_index, _) in smpls_d0]
+        train_inds_d1 = [train_index for (train_index, _) in smpls_d1]
+        g0_tune_res = _dml_tune(y, x, train_inds_d0,
+                                self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'],
+                                n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search)
+        g1_tune_res = _dml_tune(y, x, train_inds_d1,
+                                self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'],
+                                n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search)
+
+        m_tune_res = _dml_tune(treated, x, train_inds,
+                               self._learner['ml_m'], param_grids['ml_m'], scoring_methods['ml_m'],
+                               n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search)
+
+        g0_best_params = [xx.best_params_ for xx in g0_tune_res]
+        g1_best_params = [xx.best_params_ for xx in g1_tune_res]
+        m_best_params = [xx.best_params_ for xx in m_tune_res]
+
+        params = {'ml_g0': g0_best_params,
+                  'ml_g1': g1_best_params,
+                  'ml_m': m_best_params}
+        tune_res = {'g0_tune': g0_tune_res,
+                    'g1_tune': g1_tune_res,
+                    'm_tune': m_tune_res}
+
+        res = {'params': params,
+               'tune_res': tune_res}
+
+        return res
+
+    def _check_data(self, obj_dml_data):
+        if len(obj_dml_data.d_cols) > 1:
+            raise ValueError('Only one treatment variable is allowed. ' +
+                             f'Got {len(obj_dml_data.d_cols)} treatment variables.')
+
+        if obj_dml_data.z_cols is not None:
+            raise ValueError('Incompatible data. ' +
+                             ' and '.join(obj_dml_data.z_cols) +
+                             ' have been set as instrumental variable(s).')
+
+        # check if treatment level is valid
+        if np.sum(self.treated) < 5:
+            raise ValueError(
+                'The number of treated observations is less than 5. ' +
+                f'Number of treated observations: {np.sum(self.treated)} for treatment level {self.treatment_level}.'
+            )
+
+        if np.mean(self.treated) <= 0.05:
+            warnings.warn(f'The proportion of observations with treatment level {self.treatment_level} is less than 5%.'
+                          f' Got {np.mean(self.treated) * 100:.2f}%.')
+
+        return
+
+    def capo(self, basis, is_gate=False):
+        """
+        Calculate conditional average potential outcomes (CAPO) for a given basis.
+
+        Parameters
+        ----------
+        basis : :class:`pandas.DataFrame`
+            The basis for estimating the best linear predictor. Has to have the shape ``(n_obs, d)``,
+            where ``n_obs`` is the number of observations and ``d`` is the number of predictors.
+        is_gate : bool
+            Indicates whether the basis is constructed for GATE/GAPOs (dummy-basis).
+            Default is ``False``.
+
+        Returns
+        -------
+        model : :class:`doubleML.DoubleMLBLP`
+            Best linear Predictor model.
+        """
+        valid_score = ['APO']
+        if self.score not in valid_score:
+            raise ValueError('Invalid score ' + self.score + '. ' +
+                             'Valid score ' + ' or '.join(valid_score) + '.')
+
+        if self.n_rep != 1:
+            raise NotImplementedError('Only implemented for one repetition. ' +
+                                      f'Number of repetitions is {str(self.n_rep)}.')
+
+        # define the orthogonal signal
+        orth_signal = self.psi_elements['psi_b'].reshape(-1)
+        # fit the best linear predictor
+        model = DoubleMLBLP(orth_signal, basis=basis, is_gate=is_gate)
+        model.fit()
+        return model
+
+    def gapo(self, groups):
+        """
+        Calculate group average potential outcomes (GAPO) for groups.
+
+        Parameters
+        ----------
+        groups : :class:`pandas.DataFrame`
+            The group indicator for estimating the best linear predictor. Groups should be mutually exclusive.
+            Has to be dummy coded with shape ``(n_obs, d)``, where ``n_obs`` is the number of observations
+            and ``d`` is the number of groups or ``(n_obs, 1)`` and contain the corresponding groups (as str).
+
+        Returns
+        -------
+        model : :class:`doubleML.DoubleMLBLP`
+            Best linear Predictor model for group average potential outcomes.
+        """
+        if not isinstance(groups, pd.DataFrame):
+            raise TypeError('Groups must be of DataFrame type. '
+                            f'Groups of type {str(type(groups))} was passed.')
+
+        if not all(groups.dtypes == bool) or all(groups.dtypes == int):
+            if groups.shape[1] == 1:
+                groups = pd.get_dummies(groups, prefix='Group', prefix_sep='_')
+            else:
+                raise TypeError('Columns of groups must be of bool type or int type (dummy coded). '
+                                'Alternatively, groups should only contain one column.')
+
+        if any(groups.sum(0) <= 5):
+            warnings.warn('At least one group effect is estimated with less than 6 observations.')
+
+        model = self.capo(groups, is_gate=True)
+        return model
diff --git a/doubleml/irm/apos.py b/doubleml/irm/apos.py
new file mode 100644
index 00000000..2a6b5ce1
--- /dev/null
+++ b/doubleml/irm/apos.py
@@ -0,0 +1,847 @@
+import numpy as np
+import pandas as pd
+import copy
+from collections.abc import Iterable
+
+from sklearn.base import clone
+
+from joblib import Parallel, delayed
+
+from ..double_ml import DoubleML
+from ..double_ml_data import DoubleMLData, DoubleMLClusterData
+from .apo import DoubleMLAPO
+from ..double_ml_framework import concat
+
+from ..utils.resampling import DoubleMLResampling
+from ..utils._descriptive import generate_summary
+from ..utils._checks import _check_score, _check_trimming, _check_weights, _check_sample_splitting
+from ..utils.gain_statistics import gain_statistics
+
+
+class DoubleMLAPOS:
+    """Double machine learning for interactive regression models with multiple discrete treatments.
+    """
+    def __init__(self,
+                 obj_dml_data,
+                 ml_g,
+                 ml_m,
+                 treatment_levels,
+                 n_folds=5,
+                 n_rep=1,
+                 score='APO',
+                 weights=None,
+                 normalize_ipw=False,
+                 trimming_rule='truncate',
+                 trimming_threshold=1e-2,
+                 draw_sample_splitting=True):
+
+        self._dml_data = obj_dml_data
+        self._is_cluster_data = isinstance(obj_dml_data, DoubleMLClusterData)
+        self._check_data(self._dml_data)
+
+        self._all_treatment_levels = np.unique(self._dml_data.d)
+
+        self._treatment_levels = self._check_treatment_levels(treatment_levels)
+        self._n_treatment_levels = len(self._treatment_levels)
+        # Check if there are elements in self._all_treatments that are not in self.treatment_levels
+        self._add_treatment_levels = [t for t in self._all_treatment_levels if t not in self._treatment_levels]
+
+        self._normalize_ipw = normalize_ipw
+        self._n_folds = n_folds
+        self._n_rep = n_rep
+
+        # check score
+        self._score = score
+        valid_scores = ['APO']
+        _check_score(self.score, valid_scores, allow_callable=False)
+
+        # initialize framework which is constructed after the fit method is called
+        self._framework = None
+
+        # initialize and check trimming
+        self._trimming_rule = trimming_rule
+        self._trimming_threshold = trimming_threshold
+        _check_trimming(self._trimming_rule, self._trimming_threshold)
+
+        if not isinstance(self.normalize_ipw, bool):
+            raise TypeError('Normalization indicator has to be boolean. ' +
+                            f'Object of type {str(type(self.normalize_ipw))} passed.')
+
+        ml_g_is_classifier = DoubleML._check_learner(ml_g, 'ml_g', regressor=True, classifier=True)
+        _ = DoubleML._check_learner(ml_m, 'ml_m', regressor=False, classifier=True)
+        self._learner = {'ml_g': clone(ml_g), 'ml_m': clone(ml_m)}
+        if ml_g_is_classifier:
+            if obj_dml_data.binary_outcome:
+                self._predict_method = {'ml_g': 'predict_proba', 'ml_m': 'predict_proba'}
+            else:
+                raise ValueError(f'The ml_g learner {str(ml_g)} was identified as classifier '
+                                 'but the outcome variable is not binary with values 0 and 1.')
+        else:
+            self._predict_method = {'ml_g': 'predict', 'ml_m': 'predict_proba'}
+
+        # APO weights
+        _check_weights(weights, score="ATE", n_obs=obj_dml_data.n_obs, n_rep=self.n_rep)
+        self._initialize_weights(weights)
+
+        # perform sample splitting
+        self._smpls = None
+        if draw_sample_splitting:
+            self.draw_sample_splitting()
+
+            # initialize all models if splits are known
+            self._modellist = self._initialize_models()
+
+    def __str__(self):
+        class_name = self.__class__.__name__
+        header = f'================== {class_name} Object ==================\n'
+        fit_summary = str(self.summary)
+        res = header + \
+            '\n------------------ Fit summary       ------------------\n' + fit_summary
+        return res
+
+    @property
+    def score(self):
+        """
+        The score function.
+        """
+        return self._score
+
+    @property
+    def n_treatment_levels(self):
+        """
+        The number of treatment levels.
+        """
+        return self._n_treatment_levels
+
+    @property
+    def treatment_levels(self):
+        """
+        The evaluated treatment levels.
+        """
+        return self._treatment_levels
+
+    @property
+    def normalize_ipw(self):
+        """
+        Indicates whether the inverse probability weights are normalized.
+        """
+        return self._normalize_ipw
+
+    @property
+    def trimming_rule(self):
+        """
+        Specifies the used trimming rule.
+        """
+        return self._trimming_rule
+
+    @property
+    def trimming_threshold(self):
+        """
+        Specifies the used trimming threshold.
+        """
+        return self._trimming_threshold
+
+    @property
+    def weights(self):
+        """
+        Specifies the weights for a weighted average potential outcome.
+        """
+        return self._weights
+
+    @property
+    def n_folds(self):
+        """
+        Number of folds.
+        """
+        return self._n_folds
+
+    @property
+    def n_rep(self):
+        """
+        Number of repetitions for the sample splitting.
+        """
+        return self._n_rep
+
+    @property
+    def n_rep_boot(self):
+        """
+        The number of bootstrap replications.
+        """
+        if self._framework is None:
+            n_rep_boot = None
+        else:
+            n_rep_boot = self._framework.n_rep_boot
+        return n_rep_boot
+
+    @property
+    def boot_method(self):
+        """
+        The method to construct the bootstrap replications.
+        """
+        if self._framework is None:
+            method = None
+        else:
+            method = self._framework.boot_method
+        return method
+
+    @property
+    def coef(self):
+        """
+        Estimates for the causal parameter(s) after calling :meth:`fit` (shape (``n_treatment_levels``,)).
+        """
+        if self._framework is None:
+            coef = None
+        else:
+            coef = self.framework.thetas
+        return coef
+
+    @property
+    def all_coef(self):
+        """
+        Estimates of the causal parameter(s) for the ``n_rep`` different sample splits after calling :meth:`fit`
+         (shape (``n_treatment_levels``, ``n_rep``)).
+        """
+        if self._framework is None:
+            all_coef = None
+        else:
+            all_coef = self.framework.all_thetas
+        return all_coef
+
+    @property
+    def se(self):
+        """
+        Standard errors for the causal parameter(s) after calling :meth:`fit` (shape (``n_treatment_levels``,)).
+        """
+        if self._framework is None:
+            se = None
+        else:
+            se = self.framework.ses
+        return se
+
+    @property
+    def all_se(self):
+        """
+        Standard errors of the causal parameter(s) for the ``n_rep`` different sample splits after calling :meth:`fit`
+         (shape (``n_treatment_levels``, ``n_rep``)).
+        """
+        if self._framework is None:
+            all_se = None
+        else:
+            all_se = self.framework.all_ses
+        return all_se
+
+    @property
+    def t_stat(self):
+        """
+        t-statistics for the causal parameter(s) after calling :meth:`fit` (shape (``n_treatment_levels``,)).
+        """
+        if self._framework is None:
+            t_stats = None
+        else:
+            t_stats = self.framework.t_stats
+        return t_stats
+
+    @property
+    def pval(self):
+        """
+        p-values for the causal parameter(s) (shape (``n_treatment_levels``,)).
+        """
+        if self._framework is None:
+            pvals = None
+        else:
+            pvals = self.framework.pvals
+        return pvals
+
+    @property
+    def smpls(self):
+        """
+        The partition used for cross-fitting.
+        """
+        if self._smpls is None:
+            err_msg = ('Sample splitting not specified. Draw samples via .draw_sample splitting(). ' +
+                       'External samples not implemented yet.')
+            raise ValueError(err_msg)
+        return self._smpls
+
+    @property
+    def framework(self):
+        """
+        The corresponding :class:`doubleml.DoubleMLFramework` object.
+        """
+        return self._framework
+
+    @property
+    def boot_t_stat(self):
+        """
+        Bootstrapped t-statistics for the causal parameter(s) after calling :meth:`fit` and :meth:`bootstrap`
+         (shape (``n_rep_boot``, ``n_treatment_levels``, ``n_rep``)).
+        """
+        if self._framework is None:
+            boot_t_stat = None
+        else:
+            boot_t_stat = self._framework.boot_t_stat
+        return boot_t_stat
+
+    @property
+    def modellist(self):
+        """
+        The list of models for each level.
+        """
+        return self._modellist
+
+    @property
+    def sensitivity_elements(self):
+        """
+        Values of the sensitivity components after calling :meth:`fit`;
+        If available (e.g., PLR, IRM) a dictionary with entries ``sigma2``, ``nu2``, ``psi_sigma2``, ``psi_nu2``
+        and ``riesz_rep``.
+        """
+        if self._framework is None:
+            sensitivity_elements = None
+        else:
+            sensitivity_elements = self._framework.sensitivity_elements
+        return sensitivity_elements
+
+    @property
+    def sensitivity_params(self):
+        """
+        Values of the sensitivity parameters after calling :meth:`sesitivity_analysis`;
+        If available (e.g., PLR, IRM) a dictionary with entries ``theta``, ``se``, ``ci``, ``rv``
+        and ``rva``.
+        """
+        if self._framework is None:
+            sensitivity_params = None
+        else:
+            sensitivity_params = self._framework.sensitivity_params
+        return sensitivity_params
+
+    @property
+    def summary(self):
+        """
+        A summary for the estimated causal effect after calling :meth:`fit`.
+        """
+        if self.framework is None:
+            col_names = ['coef', 'std err', 't', 'P>|t|']
+            df_summary = pd.DataFrame(columns=col_names)
+        else:
+            ci = self.confint()
+            df_summary = generate_summary(self.coef, self.se, self.t_stat,
+                                          self.pval, ci, self._treatment_levels)
+        return df_summary
+
+    @property
+    def sensitivity_summary(self):
+        """
+        Returns a summary for the sensitivity analysis after calling :meth:`sensitivity_analysis`.
+
+        Returns
+        -------
+        res : str
+            Summary for the sensitivity analysis.
+        """
+        if self._framework is None:
+            raise ValueError('Apply sensitivity_analysis() before sensitivity_summary.')
+        else:
+            sensitivity_summary = self._framework.sensitivity_summary
+        return sensitivity_summary
+
+    def fit(self, n_jobs_models=None, n_jobs_cv=None, store_predictions=True, store_models=False, external_predictions=None):
+        """
+        Estimate DoubleMLAPOS models.
+
+        Parameters
+        ----------
+        n_jobs_models : None or int
+            The number of CPUs to use to fit the treatment_levels. ``None`` means ``1``.
+            Default is ``None``.
+
+        n_jobs_cv : None or int
+            The number of CPUs to use to fit the learners. ``None`` means ``1``.
+            Does not speed up computation for quantile models.
+            Default is ``None``.
+
+        store_predictions : bool
+            Indicates whether the predictions for the nuisance functions should be stored in ``predictions``.
+            Default is ``True``.
+
+        store_models : bool
+            Indicates whether the fitted models for the nuisance functions should be stored in ``models``. This allows
+            to analyze the fitted models or extract information like variable importance.
+            Default is ``False``.
+
+        external_predictions : dict or None
+            A nested dictionary where the keys correspond the the treatment levels and can contain predictions according to
+            each treatment level. The values have to be dictionaries which can contain keys ``'ml_g0'``, ``'ml_g1'``
+            and ``'ml_m'``.
+            Default is `None`.
+
+        Returns
+        -------
+        self : object
+        """
+
+        if external_predictions is not None:
+            self._check_external_predictions(external_predictions)
+            ext_pred_dict = self._rename_external_predictions(external_predictions)
+        else:
+            ext_pred_dict = None
+
+        # parallel estimation of the models
+        parallel = Parallel(n_jobs=n_jobs_models, verbose=0, pre_dispatch='2*n_jobs')
+        fitted_models = parallel(
+            delayed(self._fit_model)(
+                i_level,
+                n_jobs_cv,
+                store_predictions,
+                store_models,
+                ext_pred_dict)
+            for i_level in range(self.n_treatment_levels)
+        )
+
+        # combine the estimates and scores
+        framework_list = [None] * self.n_treatment_levels
+
+        for i_level in range(self.n_treatment_levels):
+            self._modellist[i_level] = fitted_models[i_level]
+            framework_list[i_level] = self._modellist[i_level].framework
+
+        # aggregate all frameworks
+        self._framework = concat(framework_list)
+
+        return self
+
+    def confint(self, joint=False, level=0.95):
+        """
+        Confidence intervals for DoubleML models.
+
+        Parameters
+        ----------
+        joint : bool
+            Indicates whether joint confidence intervals are computed.
+            Default is ``False``
+
+        level : float
+            The confidence level.
+            Default is ``0.95``.
+
+        Returns
+        -------
+        df_ci : pd.DataFrame
+            A data frame with the confidence interval(s).
+        """
+
+        if self.framework is None:
+            raise ValueError('Apply fit() before confint().')
+
+        df_ci = self.framework.confint(joint=joint, level=level)
+        df_ci.set_index(pd.Index(self._treatment_levels), inplace=True)
+
+        return df_ci
+
+    def bootstrap(self, method='normal', n_rep_boot=500):
+        """
+        Multiplier bootstrap for DoubleML models.
+
+        Parameters
+        ----------
+        method : str
+            A str (``'Bayes'``, ``'normal'`` or ``'wild'``) specifying the multiplier bootstrap method.
+            Default is ``'normal'``
+
+        n_rep_boot : int
+            The number of bootstrap replications.
+
+        Returns
+        -------
+        self : object
+        """
+        if self._framework is None:
+            raise ValueError('Apply fit() before bootstrap().')
+        self._framework.bootstrap(method=method, n_rep_boot=n_rep_boot)
+
+        return self
+
+    def sensitivity_analysis(self, cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95, null_hypothesis=0.0):
+        """
+        Performs a sensitivity analysis to account for unobserved confounders.
+
+        The evaluated scenario is stored as a dictionary in the property ``sensitivity_params``.
+
+        Parameters
+        ----------
+        cf_y : float
+            Percentage of the residual variation of the outcome explained by latent/confounding variables.
+            Default is ``0.03``.
+
+        cf_d : float
+            Percentage gains in the variation of the Riesz representer generated by latent/confounding variables.
+            Default is ``0.03``.
+
+        rho : float
+            The correlation between the differences in short and long representations in the main regression and
+            Riesz representer. Has to be in [-1,1]. The absolute value determines the adversarial strength of the
+            confounding (maximizes at 1.0).
+            Default is ``1.0``.
+
+        level : float
+            The confidence level.
+            Default is ``0.95``.
+
+        null_hypothesis : float or numpy.ndarray
+            Null hypothesis for the effect. Determines the robustness values.
+            If it is a single float uses the same null hypothesis for all estimated parameters.
+            Else the array has to be of shape (n_coefs,).
+            Default is ``0.0``.
+
+        Returns
+        -------
+        self : object
+        """
+
+        if self._framework is None:
+            raise ValueError('Apply fit() before sensitivity_analysis().')
+        self._framework.sensitivity_analysis(
+            cf_y=cf_y,
+            cf_d=cf_d,
+            rho=rho,
+            level=level,
+            null_hypothesis=null_hypothesis
+        )
+
+        return self
+
+    def sensitivity_plot(self, idx_treatment=0, value='theta', rho=1.0, level=0.95, null_hypothesis=0.0,
+                         include_scenario=True, benchmarks=None, fill=True, grid_bounds=(0.15, 0.15), grid_size=100):
+        """
+        Contour plot of the sensivity with respect to latent/confounding variables.
+
+        Parameters
+        ----------
+        idx_treatment : int
+            Index of the treatment to perform the sensitivity analysis.
+            Default is ``0``.
+
+        value : str
+            Determines which contours to plot. Valid values are ``'theta'`` (refers to the bounds)
+            and ``'ci'`` (refers to the bounds including statistical uncertainty).
+            Default is ``'theta'``.
+
+        rho: float
+            The correlation between the differences in short and long representations in the main regression and
+            Riesz representer. Has to be in [-1,1]. The absolute value determines the adversarial strength of the
+            confounding (maximizes at 1.0).
+            Default is ``1.0``.
+
+        level : float
+            The confidence level.
+            Default is ``0.95``.
+
+        null_hypothesis : float
+            Null hypothesis for the effect. Determines the direction of the contour lines.
+
+        include_scenario : bool
+            Indicates whether to highlight the scenario from the call of :meth:`sensitivity_analysis`.
+            Default is ``True``.
+
+        benchmarks : dict or None
+            Dictionary of benchmarks to be included in the plot. The keys are ``cf_y``, ``cf_d`` and ``name``.
+            Default is ``None``.
+
+        fill : bool
+            Indicates whether to use a heatmap style or only contour lines.
+            Default is ``True``.
+
+        grid_bounds : tuple
+            Determines the evaluation bounds of the grid for ``cf_d`` and ``cf_y``. Has to contain two floats in [0, 1).
+            Default is ``(0.15, 0.15)``.
+
+        grid_size : int
+            Determines the number of evaluation points of the grid.
+            Default is ``100``.
+
+        Returns
+        -------
+        fig : object
+            Plotly figure of the sensitivity contours.
+        """
+        if self._framework is None:
+            raise ValueError('Apply fit() before sensitivity_plot().')
+        fig = self._framework.sensitivity_plot(
+            idx_treatment=idx_treatment,
+            value=value,
+            rho=rho,
+            level=level,
+            null_hypothesis=null_hypothesis,
+            include_scenario=include_scenario,
+            benchmarks=benchmarks,
+            fill=fill,
+            grid_bounds=grid_bounds,
+            grid_size=grid_size
+        )
+
+        return fig
+
+    def sensitivity_benchmark(self, benchmarking_set, fit_args=None):
+        """
+        Computes a benchmark for a given set of features.
+        Returns a DataFrame containing the corresponding values for cf_y, cf_d, rho and the change in estimates.
+        Returns
+        -------
+        benchmark_results : pandas.DataFrame
+            Benchmark results.
+        """
+        x_list_long = self._dml_data.x_cols
+
+        # input checks
+        if self.sensitivity_elements is None:
+            raise NotImplementedError(f'Sensitivity analysis not yet implemented for {self.__class__.__name__}.')
+        if not isinstance(benchmarking_set, list):
+            raise TypeError('benchmarking_set must be a list. '
+                            f'{str(benchmarking_set)} of type {type(benchmarking_set)} was passed.')
+        if len(benchmarking_set) == 0:
+            raise ValueError('benchmarking_set must not be empty.')
+        if not set(benchmarking_set) <= set(x_list_long):
+            raise ValueError(f"benchmarking_set must be a subset of features {str(self._dml_data.x_cols)}. "
+                             f'{str(benchmarking_set)} was passed.')
+        if fit_args is not None and not isinstance(fit_args, dict):
+            raise TypeError('fit_args must be a dict. '
+                            f'{str(fit_args)} of type {type(fit_args)} was passed.')
+
+        # refit short form of the model
+        x_list_short = [x for x in x_list_long if x not in benchmarking_set]
+        dml_short = copy.deepcopy(self)
+        dml_short._dml_data.x_cols = x_list_short
+        if fit_args is not None:
+            dml_short.fit(**fit_args)
+        else:
+            dml_short.fit()
+
+        benchmark_dict = gain_statistics(dml_long=self, dml_short=dml_short)
+        df_benchmark = pd.DataFrame(benchmark_dict, index=self.treatment_levels)
+        return df_benchmark
+
+    def draw_sample_splitting(self):
+        """
+        Draw sample splitting for DoubleML models.
+
+        The samples are drawn according to the attributes
+        ``n_folds`` and ``n_rep``.
+
+        Returns
+        -------
+        self : object
+        """
+        obj_dml_resampling = DoubleMLResampling(n_folds=self.n_folds,
+                                                n_rep=self.n_rep,
+                                                n_obs=self._dml_data.n_obs,
+                                                stratify=self._dml_data.d)
+        self._smpls = obj_dml_resampling.split_samples()
+
+        return self
+
+    def set_sample_splitting(self, all_smpls, all_smpls_cluster=None):
+        """
+        Set the sample splitting for DoubleML models.
+
+        The  attributes ``n_folds`` and ``n_rep`` are derived from the provided partition.
+
+        Parameters
+        ----------
+        all_smpls : list or tuple
+            If nested list of lists of tuples:
+                The outer list needs to provide an entry per repeated sample splitting (length of list is set as
+                ``n_rep``).
+                The inner list needs to provide a tuple (train_ind, test_ind) per fold (length of list is set as
+                ``n_folds``). test_ind must form a partition for each inner list.
+            If list of tuples:
+                The list needs to provide a tuple (train_ind, test_ind) per fold (length of list is set as
+                ``n_folds``). test_ind must form a partition. ``n_rep=1`` is always set.
+            If tuple:
+                Must be a tuple with two elements train_ind and test_ind. Only viable option is to set
+                train_ind and test_ind to np.arange(n_obs), which corresponds to no sample splitting.
+                ``n_folds=1`` and ``n_rep=1`` is always set.
+
+        Returns
+        -------
+        self : object
+
+        Examples
+        --------
+        >>> import numpy as np
+        >>> import doubleml as dml
+        >>> from doubleml.datasets import make_plr_CCDDHNR2018
+        >>> from sklearn.ensemble import RandomForestRegressor
+        >>> from sklearn.base import clone
+        >>> np.random.seed(3141)
+        >>> learner = RandomForestRegressor(max_depth=2, n_estimators=10)
+        >>> ml_g = learner
+        >>> ml_m = learner
+        >>> obj_dml_data = make_plr_CCDDHNR2018(n_obs=10, alpha=0.5)
+        >>> dml_plr_obj = dml.DoubleMLPLR(obj_dml_data, ml_g, ml_m)
+        >>> # sample splitting with two folds and cross-fitting
+        >>> smpls = [([0, 1, 2, 3, 4], [5, 6, 7, 8, 9]),
+        >>>          ([5, 6, 7, 8, 9], [0, 1, 2, 3, 4])]
+        >>> dml_plr_obj.set_sample_splitting(smpls)
+        >>> # sample splitting with two folds and repeated cross-fitting with n_rep = 2
+        >>> smpls = [[([0, 1, 2, 3, 4], [5, 6, 7, 8, 9]),
+        >>>           ([5, 6, 7, 8, 9], [0, 1, 2, 3, 4])],
+        >>>          [([0, 2, 4, 6, 8], [1, 3, 5, 7, 9]),
+        >>>           ([1, 3, 5, 7, 9], [0, 2, 4, 6, 8])]]
+        >>> dml_plr_obj.set_sample_splitting(smpls)
+        """
+        self._smpls, self._smpls_cluster, self._n_rep, self._n_folds = _check_sample_splitting(
+            all_smpls, all_smpls_cluster, self._dml_data, self._is_cluster_data)
+
+        self._modellist = self._initialize_models()
+
+        return self
+
+    def causal_contrast(self, reference_levels):
+        """
+        Average causal contrasts for DoubleMLAPOS models. Estimates the difference in
+        average potential outcomes between the treatment levels and the reference levels.
+        The reference levels have to be a subset of the treatment levels or a single
+        treatment level.
+
+        Parameters
+        ----------
+        reference_levels :
+            The reference levels for the difference in average potential outcomes.
+            Has to be an element of ``treatment_levels``.
+
+        Returns
+        -------
+        acc : DoubleMLFramework
+            A DoubleMLFramwork class for average causal contrast(s).
+        """
+
+        if self.framework is None:
+            raise ValueError('Apply fit() before causal_contrast().')
+        if self.n_treatment_levels == 1:
+            raise ValueError('Only one treatment level. No causal contrast can be computed.')
+        is_iterable = isinstance(reference_levels, Iterable)
+        if not is_iterable:
+            reference_levels = [reference_levels]
+        is_treatment_level_subset = set(reference_levels).issubset(set(self.treatment_levels))
+        if not is_treatment_level_subset:
+            raise ValueError('Invalid reference_levels. reference_levels has to be an iterable subset of treatment_levels or '
+                             'a single treatment level.')
+
+        skip_index = []
+        all_treatment_names = []
+        all_acc_frameworks = []
+        for ref_lvl in reference_levels:
+            i_ref_lvl = self.treatment_levels.index(ref_lvl)
+            ref_framework = self.modellist[i_ref_lvl].framework
+
+            skip_index += [i_ref_lvl]
+            all_acc_frameworks += [model.framework - ref_framework for i, model in
+                                   enumerate(self.modellist) if i not in skip_index]
+            all_treatment_names += [f"{self.treatment_levels[i]} vs {self.treatment_levels[i_ref_lvl]}" for
+                                    i in range(self.n_treatment_levels) if i not in skip_index]
+
+        acc = concat(all_acc_frameworks)
+        acc.treatment_names = all_treatment_names
+        return acc
+
+    def _fit_model(self, i_level, n_jobs_cv=None, store_predictions=True, store_models=False, external_predictions_dict=None):
+
+        model = self.modellist[i_level]
+        if external_predictions_dict is not None:
+            external_predictions = external_predictions_dict[self.treatment_levels[i_level]]
+        else:
+            external_predictions = None
+        model.fit(n_jobs_cv=n_jobs_cv, store_predictions=store_predictions, store_models=store_models,
+                  external_predictions=external_predictions)
+        return model
+
+    def _check_treatment_levels(self, treatment_levels):
+        is_iterable = isinstance(treatment_levels, Iterable)
+        if not is_iterable:
+            treatment_level_list = [treatment_levels]
+        else:
+            treatment_level_list = [t_lvl for t_lvl in treatment_levels]
+        is_d_subset = set(treatment_level_list).issubset(set(self._all_treatment_levels))
+        if not is_d_subset:
+            raise ValueError('Invalid reference_levels. reference_levels has to be an iterable subset or '
+                             'a single element of the unique treatment levels in the data.')
+        return treatment_level_list
+
+    def _check_data(self, obj_dml_data):
+        if not isinstance(obj_dml_data, DoubleMLData):
+            raise TypeError('The data must be of DoubleMLData or DoubleMLClusterData type.')
+        if obj_dml_data.z is not None:
+            raise ValueError('The data must not contain instrumental variables.')
+        return
+
+    def _check_external_predictions(self, external_predictions):
+        expected_keys = self.treatment_levels
+        if not isinstance(external_predictions, dict):
+            raise TypeError('external_predictions must be a dictionary. ' +
+                            f'Object of type {type(external_predictions)} passed.')
+
+        if not set(external_predictions.keys()).issubset(set(expected_keys)):
+            raise ValueError('external_predictions must be a subset of all treatment levels. ' +
+                             f'Expected keys: {set(expected_keys)}. ' +
+                             f'Passed keys: {set(external_predictions.keys())}.')
+
+        expected_learner_keys = ['ml_g0', 'ml_g1', 'ml_m']
+        for key, value in external_predictions.items():
+            if not isinstance(value, dict):
+                raise TypeError(f'external_predictions[{key}] must be a dictionary. ' +
+                                f'Object of type {type(value)} passed.')
+            if not set(value.keys()).issubset(set(expected_learner_keys)):
+                raise ValueError(f'external_predictions[{key}] must be a subset of {set(expected_learner_keys)}. ' +
+                                 f'Passed keys: {set(value.keys())}.')
+
+        return
+
+    def _rename_external_predictions(self, external_predictions):
+        d_col = self._dml_data.d_cols[0]
+        ext_pred_dict = {treatment_level: {d_col: {}} for treatment_level in self.treatment_levels}
+        for treatment_level in self.treatment_levels:
+            if "ml_g1" in external_predictions[treatment_level]:
+                ext_pred_dict[treatment_level][d_col]['ml_g1'] = external_predictions[treatment_level]['ml_g1']
+            if "ml_m" in external_predictions[treatment_level]:
+                ext_pred_dict[treatment_level][d_col]['ml_m'] = external_predictions[treatment_level]['ml_m']
+            if "ml_g0" in external_predictions[treatment_level]:
+                ext_pred_dict[treatment_level][d_col]['ml_g0'] = external_predictions[treatment_level]['ml_g0']
+
+        return ext_pred_dict
+
+    def _initialize_weights(self, weights):
+        if weights is None:
+            weights = np.ones(self._dml_data.n_obs)
+        if isinstance(weights, np.ndarray):
+            self._weights = weights
+        else:
+            assert isinstance(weights, dict)
+            self._weights = weights
+
+    def _initialize_models(self):
+        modellist = [None] * self.n_treatment_levels
+        kwargs = {
+            'obj_dml_data': self._dml_data,
+            'ml_g': self._learner['ml_g'],
+            'ml_m': self._learner['ml_m'],
+            'score': self.score,
+            'n_folds': self.n_folds,
+            'n_rep': self.n_rep,
+            'weights': self.weights,
+            'trimming_rule': self.trimming_rule,
+            'trimming_threshold': self.trimming_threshold,
+            'normalize_ipw': self.normalize_ipw,
+            'draw_sample_splitting': False
+        }
+        for i_level in range(self.n_treatment_levels):
+            # initialize models for all levels
+            model = DoubleMLAPO(
+                treatment_level=self._treatment_levels[i_level],
+                **kwargs
+            )
+
+            # synchronize the sample splitting
+            model.set_sample_splitting(all_smpls=self.smpls)
+            modellist[i_level] = model
+
+        return modellist
diff --git a/doubleml/irm/iivm.py b/doubleml/irm/iivm.py
index c131449b..c2f85dd4 100644
--- a/doubleml/irm/iivm.py
+++ b/doubleml/irm/iivm.py
@@ -7,7 +7,8 @@
 from ..double_ml_score_mixins import LinearScoreMixin
 
 from ..utils._estimation import _dml_cv_predict, _get_cond_smpls, _dml_tune, _trimm, _normalize_ipw
-from ..utils._checks import _check_score, _check_trimming, _check_finite_predictions, _check_is_propensity
+from ..utils._checks import _check_score, _check_trimming, _check_finite_predictions, _check_is_propensity, \
+    _check_binary_predictions
 
 
 class DoubleMLIIVM(LinearScoreMixin, DoubleML):
@@ -264,15 +265,9 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa
             g_hat0['targets'][z == 1] = np.nan
 
         if self._dml_data.binary_outcome:
-            binary_preds = (type_of_target(g_hat0['preds']) == 'binary')
-            zero_one_preds = np.all((np.power(g_hat0['preds'], 2) - g_hat0['preds']) == 0)
-            if binary_preds & zero_one_preds:
-                raise ValueError(f'For the binary outcome variable {self._dml_data.y_col}, '
-                                 f'predictions obtained with the ml_g learner {str(self._learner["ml_g"])} are also '
-                                 'observed to be binary with values 0 and 1. Make sure that for classifiers '
-                                 'probabilities and not labels are predicted.')
-
+            _check_binary_predictions(g_hat0['preds'], self._learner['ml_g'], 'ml_g', self._dml_data.y_col)
             _check_is_propensity(g_hat0['preds'], self._learner['ml_g'], 'ml_g', smpls, eps=1e-12)
+
         if external_predictions['ml_g1'] is not None:
             g_hat1 = {'preds': external_predictions['ml_g1'],
                       'targets': None,
@@ -287,14 +282,7 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa
             g_hat1['targets'][z == 0] = np.nan
 
         if self._dml_data.binary_outcome:
-            binary_preds = (type_of_target(g_hat1['preds']) == 'binary')
-            zero_one_preds = np.all((np.power(g_hat1['preds'], 2) - g_hat1['preds']) == 0)
-            if binary_preds & zero_one_preds:
-                raise ValueError(f'For the binary outcome variable {self._dml_data.y_col}, '
-                                 f'predictions obtained with the ml_g learner {str(self._learner["ml_g"])} are also '
-                                 'observed to be binary with values 0 and 1. Make sure that for classifiers '
-                                 'probabilities and not labels are predicted.')
-
+            _check_binary_predictions(g_hat1['preds'], self._learner['ml_g'], 'ml_g', self._dml_data.y_col)
             _check_is_propensity(g_hat1['preds'], self._learner['ml_g'], 'ml_g', smpls, eps=1e-12)
 
         # nuisance m
diff --git a/doubleml/irm/irm.py b/doubleml/irm/irm.py
index 82d22d1d..1b1695c6 100644
--- a/doubleml/irm/irm.py
+++ b/doubleml/irm/irm.py
@@ -13,7 +13,7 @@
 
 from ..utils._estimation import _dml_cv_predict, _get_cond_smpls, _dml_tune, _trimm, _normalize_ipw, _cond_targets
 from ..utils._checks import _check_score, _check_trimming, _check_finite_predictions, _check_is_propensity, _check_integer, \
-    _check_weights
+    _check_weights, _check_binary_predictions
 
 
 class DoubleMLIRM(LinearScoreMixin, DoubleML):
@@ -143,7 +143,6 @@ def __init__(self,
         ml_g_is_classifier = self._check_learner(ml_g, 'ml_g', regressor=True, classifier=True)
         _ = self._check_learner(ml_m, 'ml_m', regressor=False, classifier=True)
         self._learner = {'ml_g': ml_g, 'ml_m': ml_m}
-        self._normalize_ipw = normalize_ipw
         if ml_g_is_classifier:
             if obj_dml_data.binary_outcome:
                 self._predict_method = {'ml_g': 'predict_proba', 'ml_m': 'predict_proba'}
@@ -154,6 +153,7 @@ def __init__(self,
             self._predict_method = {'ml_g': 'predict', 'ml_m': 'predict_proba'}
         self._initialize_ml_nuisance_params()
 
+        self._normalize_ipw = normalize_ipw
         if not isinstance(self.normalize_ipw, bool):
             raise TypeError('Normalization indicator has to be boolean. ' +
                             f'Object of type {str(type(self.normalize_ipw))} passed.')
@@ -275,13 +275,7 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa
             g_hat0['targets'] = _cond_targets(g_hat0['targets'], cond_sample=(d == 0))
 
             if self._dml_data.binary_outcome:
-                binary_preds = (type_of_target(g_hat0['preds']) == 'binary')
-                zero_one_preds = np.all((np.power(g_hat0['preds'], 2) - g_hat0['preds']) == 0)
-                if binary_preds & zero_one_preds:
-                    raise ValueError(f'For the binary outcome variable {self._dml_data.y_col}, '
-                                     f'predictions obtained with the ml_g learner {str(self._learner["ml_g"])} are also '
-                                     'observed to be binary with values 0 and 1. Make sure that for classifiers '
-                                     'probabilities and not labels are predicted.')
+                _check_binary_predictions(g_hat0['preds'], self._learner['ml_g'], 'ml_g', self._dml_data.y_col)
 
         if g1_external:
             # use external predictions
@@ -297,13 +291,7 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa
             g_hat1['targets'] = _cond_targets(g_hat1['targets'], cond_sample=(d == 1))
 
         if self._dml_data.binary_outcome & (self.score != 'ATTE'):
-            binary_preds = (type_of_target(g_hat1['preds']) == 'binary')
-            zero_one_preds = np.all((np.power(g_hat1['preds'], 2) - g_hat1['preds']) == 0)
-            if binary_preds & zero_one_preds:
-                raise ValueError(f'For the binary outcome variable {self._dml_data.y_col}, '
-                                 f'predictions obtained with the ml_g learner {str(self._learner["ml_g"])} are also '
-                                 'observed to be binary with values 0 and 1. Make sure that for classifiers '
-                                 'probabilities and not labels are predicted.')
+            _check_binary_predictions(g_hat1['preds'], self._learner['ml_g'], 'ml_g', self._dml_data.y_col)
 
         # nuisance m
         if m_external:
@@ -340,7 +328,6 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa
 
     def _score_elements(self, y, d, g_hat0, g_hat1, m_hat, smpls):
 
-        m_hat_adj = np.full_like(m_hat, np.nan, dtype='float64')
         if self.normalize_ipw:
             m_hat_adj = _normalize_ipw(m_hat, d)
         else:
@@ -376,11 +363,7 @@ def _sensitivity_element_est(self, preds):
 
         m_hat = preds['predictions']['ml_m']
         g_hat0 = preds['predictions']['ml_g0']
-        if self.score == 'ATE':
-            g_hat1 = preds['predictions']['ml_g1']
-        else:
-            assert self.score == 'ATTE'
-            g_hat1 = y
+        g_hat1 = preds['predictions']['ml_g1']
 
         # use weights make this extendable
         weights, weights_bar = self._get_weights(m_hat=m_hat)
@@ -424,7 +407,6 @@ def _nuisance_tuning(self, smpls, param_grids, scoring_methods, n_folds_tune, n_
         g0_tune_res = _dml_tune(y, x, train_inds_d0,
                                 self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'],
                                 n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search)
-        g1_tune_res = list()
         g1_tune_res = _dml_tune(y, x, train_inds_d1,
                                 self._learner['ml_g'], param_grids['ml_g'], scoring_methods['ml_g'],
                                 n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search)
diff --git a/doubleml/irm/qte.py b/doubleml/irm/qte.py
index 17894a60..2a212d77 100644
--- a/doubleml/irm/qte.py
+++ b/doubleml/irm/qte.py
@@ -15,6 +15,8 @@
 from ..utils.resampling import DoubleMLResampling
 from ..utils._checks import _check_score, _check_trimming, _check_zero_one_treatment, _check_sample_splitting
 
+from ..utils._descriptive import generate_summary
+
 
 class DoubleMLQTE:
     """Double machine learning for quantile treatment effects
@@ -138,7 +140,6 @@ def __init__(self,
         self._trimming_threshold = trimming_threshold
         _check_trimming(self._trimming_rule, self._trimming_threshold)
 
-        self._check_quantile()
         if not isinstance(self.normalize_ipw, bool):
             raise TypeError('Normalization indicator has to be boolean. ' +
                             f'Object of type {str(type(self.normalize_ipw))} passed.')
@@ -356,18 +357,13 @@ def summary(self):
         """
         A summary for the estimated causal effect after calling :meth:`fit`.
         """
-        col_names = ['coef', 'std err', 't', 'P>|t|']
         if self.framework is None:
+            col_names = ['coef', 'std err', 't', 'P>|t|']
             df_summary = pd.DataFrame(columns=col_names)
         else:
-            summary_stats = np.transpose(np.vstack(
-                [self.coef, self.se,
-                 self.t_stat, self.pval]))
-            df_summary = pd.DataFrame(summary_stats,
-                                      columns=col_names,
-                                      index=self.quantiles)
             ci = self.confint()
-            df_summary = df_summary.join(ci)
+            df_summary = generate_summary(self.coef, self.se, self.t_stat,
+                                          self.pval, ci, self.quantiles)
         return df_summary
 
     def fit(self, n_jobs_models=None, n_jobs_cv=None, store_predictions=True, store_models=False, external_predictions=None):
@@ -411,14 +407,13 @@ def fit(self, n_jobs_models=None, n_jobs_cv=None, store_predictions=True, store_
         framework_list = [None] * self.n_quantiles
 
         for i_quant in range(self.n_quantiles):
-            self._i_quant = i_quant
             # save the parallel fitted models in the right list
-            self._modellist_0[self._i_quant] = fitted_models[self._i_quant][0]
-            self._modellist_1[self._i_quant] = fitted_models[self._i_quant][1]
+            self._modellist_0[i_quant] = fitted_models[i_quant][0]
+            self._modellist_1[i_quant] = fitted_models[i_quant][1]
 
             # set up the framework
-            framework_list[self._i_quant] = self._modellist_1[self._i_quant].framework - \
-                self._modellist_0[self._i_quant].framework
+            framework_list[i_quant] = self._modellist_1[i_quant].framework - \
+                self._modellist_0[i_quant].framework
 
         # aggregate all frameworks
         self._framework = concat(framework_list)
@@ -626,7 +621,6 @@ def _initialize_models(self):
             'draw_sample_splitting': False
         }
         for i_quant in range(self.n_quantiles):
-            self._i_quant = i_quant
 
             # initialize models for both potential quantiles
             if self.score == 'PQ':
diff --git a/doubleml/plm/plr.py b/doubleml/plm/plr.py
index 3b7d90f5..d5810b97 100644
--- a/doubleml/plm/plr.py
+++ b/doubleml/plm/plr.py
@@ -1,7 +1,6 @@
 import numpy as np
 import pandas as pd
 from sklearn.utils import check_X_y
-from sklearn.utils.multiclass import type_of_target
 from sklearn.base import clone
 
 import warnings
@@ -12,7 +11,7 @@
 from ..utils.blp import DoubleMLBLP
 
 from ..utils._estimation import _dml_cv_predict, _dml_tune
-from ..utils._checks import _check_score, _check_finite_predictions, _check_is_propensity
+from ..utils._checks import _check_score, _check_finite_predictions, _check_is_propensity, _check_binary_predictions
 
 
 class DoubleMLPLR(LinearScoreMixin, DoubleML):
@@ -198,13 +197,7 @@ def _nuisance_est(self, smpls, n_jobs_cv, external_predictions, return_models=Fa
             _check_is_propensity(m_hat['preds'], self._learner['ml_m'], 'ml_m', smpls, eps=1e-12)
 
         if self._dml_data.binary_treats[self._dml_data.d_cols[self._i_treat]]:
-            binary_preds = (type_of_target(m_hat['preds']) == 'binary')
-            zero_one_preds = np.all((np.power(m_hat['preds'], 2) - m_hat['preds']) == 0)
-            if binary_preds & zero_one_preds:
-                raise ValueError(f'For the binary treatment variable {self._dml_data.d_cols[self._i_treat]}, '
-                                 f'predictions obtained with the ml_m learner {str(self._learner["ml_m"])} are also '
-                                 'observed to be binary with values 0 and 1. Make sure that for classifiers '
-                                 'probabilities and not labels are predicted.')
+            _check_binary_predictions(m_hat['preds'], self._learner['ml_m'], 'ml_m', self._dml_data.d_cols[self._i_treat])
 
         # an estimate of g is obtained for the IV-type score and callable scores
         g_hat = {'preds': None, 'targets': None, 'models': None}
diff --git a/doubleml/utils/_checks.py b/doubleml/utils/_checks.py
index e54c4041..d7d2881e 100644
--- a/doubleml/utils/_checks.py
+++ b/doubleml/utils/_checks.py
@@ -206,6 +206,16 @@ def _check_is_propensity(preds, learner, learner_name, smpls, eps=1e-12):
     return
 
 
+def _check_binary_predictions(pred, learner, learner_name, variable_name):
+    binary_preds = (type_of_target(pred) == 'binary')
+    zero_one_preds = np.all((np.power(pred, 2) - pred) == 0)
+    if binary_preds & zero_one_preds:
+        raise ValueError(f'For the binary variable {variable_name}, '
+                         f'predictions obtained with the {learner_name} learner {str(learner)} are also '
+                         'observed to be binary with values 0 and 1. Make sure that for classifiers '
+                         'probabilities and not labels are predicted.')
+
+
 def _check_benchmarks(benchmarks):
     if benchmarks is not None:
         if not isinstance(benchmarks, dict):
diff --git a/doubleml/utils/_descriptive.py b/doubleml/utils/_descriptive.py
new file mode 100644
index 00000000..54144bc8
--- /dev/null
+++ b/doubleml/utils/_descriptive.py
@@ -0,0 +1,13 @@
+import numpy as np
+import pandas as pd
+
+
+def generate_summary(coef, se, t_stat, pval, ci, index_names=None):
+    col_names = ['coef', 'std err', 't', 'P>|t|']
+    summary_stats = np.transpose(np.vstack(
+        [coef, se, t_stat, pval]))
+    df_summary = pd.DataFrame(summary_stats, columns=col_names)
+    if index_names is not None:
+        df_summary.index = index_names
+    df_summary = df_summary.join(ci)
+    return df_summary
diff --git a/doubleml/utils/gain_statistics.py b/doubleml/utils/gain_statistics.py
index 3c50d084..2fa233b3 100644
--- a/doubleml/utils/gain_statistics.py
+++ b/doubleml/utils/gain_statistics.py
@@ -19,44 +19,49 @@ def gain_statistics(dml_long, dml_short):
     benchmark_dict : dict
         Benchmarking dictionary (dict) with values for ``cf_d``, ``cf_y``, ``rho``, and ``delta_theta``.
     """
-    if not isinstance(dml_long.sensitivity_elements, dict):
+
+    # set input for readability
+    sensitivity_elements_long = dml_long.framework.sensitivity_elements
+    sensitivity_elements_short = dml_short.framework.sensitivity_elements
+
+    if not isinstance(sensitivity_elements_long, dict):
         raise TypeError("dml_long does not contain the necessary sensitivity elements. "
-                        "Expected dict for dml_long.sensitivity_elements.")
+                        "Expected dict for dml_long.framework.sensitivity_elements.")
     expected_keys = ['sigma2', 'nu2']
-    if not all(key in dml_long.sensitivity_elements.keys() for key in expected_keys):
+    if not all(key in sensitivity_elements_long.keys() for key in expected_keys):
         raise ValueError("dml_long does not contain the necessary sensitivity elements. "
                          "Required keys are: " + str(expected_keys))
-    if not isinstance(dml_short.sensitivity_elements, dict):
+    if not isinstance(sensitivity_elements_short, dict):
         raise TypeError("dml_short does not contain the necessary sensitivity elements. "
-                        "Expected dict for dml_short.sensitivity_elements.")
-    if not all(key in dml_short.sensitivity_elements.keys() for key in expected_keys):
+                        "Expected dict for dml_short.framework.sensitivity_elements.")
+    if not all(key in sensitivity_elements_short.keys() for key in expected_keys):
         raise ValueError("dml_short does not contain the necessary sensitivity elements. "
                          "Required keys are: " + str(expected_keys))
 
     for key in expected_keys:
-        if not isinstance(dml_long.sensitivity_elements[key], np.ndarray):
+        if not isinstance(sensitivity_elements_long[key], np.ndarray):
             raise TypeError("dml_long does not contain the necessary sensitivity elements. "
                             f"Expected numpy.ndarray for key {key}.")
-        if not isinstance(dml_short.sensitivity_elements[key], np.ndarray):
+        if not isinstance(sensitivity_elements_short[key], np.ndarray):
             raise TypeError("dml_short does not contain the necessary sensitivity elements. "
                             f"Expected numpy.ndarray for key {key}.")
-        if len(dml_long.sensitivity_elements[key].shape) != 3 or dml_long.sensitivity_elements[key].shape[0] != 1:
+        if len(sensitivity_elements_long[key].shape) != 3 or sensitivity_elements_long[key].shape[0] != 1:
             raise ValueError("dml_long does not contain the necessary sensitivity elements. "
                              f"Expected 3 dimensions of shape (1, n_coef, n_rep) for key {key}.")
-        if len(dml_short.sensitivity_elements[key].shape) != 3 or dml_short.sensitivity_elements[key].shape[0] != 1:
+        if len(sensitivity_elements_short[key].shape) != 3 or sensitivity_elements_short[key].shape[0] != 1:
             raise ValueError("dml_short does not contain the necessary sensitivity elements. "
                              f"Expected 3 dimensions of shape (1, n_coef, n_rep) for key {key}.")
-        if not np.array_equal(dml_long.sensitivity_elements[key].shape, dml_short.sensitivity_elements[key].shape):
+        if not np.array_equal(sensitivity_elements_long[key].shape, sensitivity_elements_short[key].shape):
             raise ValueError("dml_long and dml_short do not contain the same shape of sensitivity elements. "
-                             "Shapes of " + key + " are: " + str(dml_long.sensitivity_elements[key].shape) +
-                             " and " + str(dml_short.sensitivity_elements[key].shape))
+                             "Shapes of " + key + " are: " + str(sensitivity_elements_long[key].shape) +
+                             " and " + str(sensitivity_elements_short[key].shape))
 
     if not isinstance(dml_long.all_coef, np.ndarray):
         raise TypeError("dml_long.all_coef does not contain the necessary coefficients. Expected numpy.ndarray.")
     if not isinstance(dml_short.all_coef, np.ndarray):
         raise TypeError("dml_short.all_coef does not contain the necessary coefficients. Expected numpy.ndarray.")
 
-    expected_shape = (dml_long.sensitivity_elements['sigma2'].shape[2], dml_long.sensitivity_elements['sigma2'].shape[1])
+    expected_shape = (sensitivity_elements_long['sigma2'].shape[1], sensitivity_elements_long['sigma2'].shape[2])
     if dml_long.all_coef.shape != expected_shape:
         raise ValueError("dml_long.all_coef does not contain the necessary coefficients. Expected shape: " +
                          str(expected_shape))
@@ -66,10 +71,10 @@ def gain_statistics(dml_long, dml_short):
 
     # save elements for readability
     var_y = np.var(dml_long._dml_data.y)
-    var_y_residuals_long = np.squeeze(dml_long.sensitivity_elements['sigma2'], axis=0)
-    nu2_long = np.squeeze(dml_long.sensitivity_elements['nu2'], axis=0)
-    var_y_residuals_short = np.squeeze(dml_short.sensitivity_elements['sigma2'], axis=0)
-    nu2_short = np.squeeze(dml_short.sensitivity_elements['nu2'], axis=0)
+    var_y_residuals_long = np.squeeze(sensitivity_elements_long['sigma2'], axis=0)
+    nu2_long = np.squeeze(sensitivity_elements_long['nu2'], axis=0)
+    var_y_residuals_short = np.squeeze(sensitivity_elements_short['sigma2'], axis=0)
+    nu2_short = np.squeeze(sensitivity_elements_short['nu2'], axis=0)
 
     # compute nonparametric R2
     R2_y_long = 1.0 - np.divide(var_y_residuals_long, var_y)
@@ -79,12 +84,12 @@ def gain_statistics(dml_long, dml_short):
     # Gain statistics
     all_cf_y_benchmark = np.clip(np.divide((R2_y_long - R2_y_short), (1.0 - R2_y_long)), 0, 1)
     all_cf_d_benchmark = np.clip(np.divide((1.0 - R2_riesz), R2_riesz), 0, 1)
-    cf_y_benchmark = np.median(all_cf_y_benchmark, axis=0)
-    cf_d_benchmark = np.median(all_cf_d_benchmark, axis=0)
+    cf_y_benchmark = np.median(all_cf_y_benchmark, axis=1)
+    cf_d_benchmark = np.median(all_cf_d_benchmark, axis=1)
 
     # change in estimates (slightly different to paper)
-    all_delta_theta = np.transpose(dml_short.all_coef - dml_long.all_coef)
-    delta_theta = np.median(all_delta_theta, axis=0)
+    all_delta_theta = dml_short.all_coef - dml_long.all_coef
+    delta_theta = np.median(all_delta_theta, axis=1)
 
     # degree of adversity
     var_g = var_y_residuals_short - var_y_residuals_long
@@ -97,7 +102,7 @@ def gain_statistics(dml_long, dml_short):
                                    where=denom != 0),
                          0.0, 1.0)
     all_rho_benchmark = np.multiply(rho_values, rho_sign)
-    rho_benchmark = np.median(all_rho_benchmark, axis=0)
+    rho_benchmark = np.median(all_rho_benchmark, axis=1)
     benchmark_dict = {
         "cf_y": cf_y_benchmark,
         "cf_d": cf_d_benchmark,


Test patch:
diff --git a/doubleml/irm/tests/_utils_apo_manual.py b/doubleml/irm/tests/_utils_apo_manual.py
new file mode 100644
index 00000000..e22f80ff
--- /dev/null
+++ b/doubleml/irm/tests/_utils_apo_manual.py
@@ -0,0 +1,221 @@
+import numpy as np
+from sklearn.base import clone, is_classifier
+
+from ...tests._utils_boot import boot_manual, draw_weights
+from ...tests._utils import fit_predict, fit_predict_proba, tune_grid_search
+
+from ...utils._estimation import _normalize_ipw
+from ...utils._checks import _check_is_propensity
+
+
+def fit_apo(y, x, d,
+            learner_g, learner_m, treatment_level, all_smpls, score,
+            n_rep=1, g0_params=None, g1_params=None, m_params=None,
+            normalize_ipw=False, trimming_threshold=1e-2):
+    n_obs = len(y)
+    treated = (d == treatment_level)
+
+    thetas = np.zeros(n_rep)
+    ses = np.zeros(n_rep)
+    all_g_hat0 = list()
+    all_g_hat1 = list()
+    all_m_hat = list()
+
+    for i_rep in range(n_rep):
+        smpls = all_smpls[i_rep]
+        g_hat0, g_hat1, m_hat = fit_nuisance_apo(y, x, d, treated,
+                                                 learner_g, learner_m, smpls, score,
+                                                 g0_params=g0_params, g1_params=g1_params, m_params=m_params,
+                                                 trimming_threshold=trimming_threshold)
+
+        all_g_hat0.append(g_hat0)
+        all_g_hat1.append(g_hat1)
+        all_m_hat.append(m_hat)
+
+        thetas[i_rep], ses[i_rep] = apo_dml2(y, x, d, treated,
+                                             g_hat0, g_hat1, m_hat,
+                                             smpls, score, normalize_ipw)
+
+    theta = np.median(thetas)
+    se = np.sqrt(np.median(np.power(ses, 2) * n_obs + np.power(thetas - theta, 2)) / n_obs)
+
+    res = {'theta': theta, 'se': se,
+           'thetas': thetas, 'ses': ses,
+           'all_g_hat0': all_g_hat0, 'all_g_hat1': all_g_hat1, 'all_m_hat': all_m_hat}
+
+    return res
+
+
+def fit_nuisance_apo(y, x, d, treated,
+                     learner_g, learner_m, smpls, score,
+                     g0_params=None, g1_params=None, m_params=None,
+                     trimming_threshold=1e-12):
+    ml_g0 = clone(learner_g)
+    ml_g1 = clone(learner_g)
+
+    train_cond0 = np.where(treated == 0)[0]
+    if is_classifier(learner_g):
+        g_hat0_list = fit_predict_proba(y, x, ml_g0, g0_params, smpls,
+                                        train_cond=train_cond0)
+    else:
+        g_hat0_list = fit_predict(y, x, ml_g0, g0_params, smpls,
+                                  train_cond=train_cond0)
+
+    train_cond1 = np.where(treated == 1)[0]
+    if is_classifier(learner_g):
+        g_hat1_list = fit_predict_proba(y, x, ml_g1, g1_params, smpls,
+                                        train_cond=train_cond1)
+    else:
+        g_hat1_list = fit_predict(y, x, ml_g1, g1_params, smpls,
+                                  train_cond=train_cond1)
+
+    ml_m = clone(learner_m)
+    m_hat_list = fit_predict_proba(treated, x, ml_m, m_params, smpls,
+                                   trimming_threshold=trimming_threshold)
+
+    return g_hat0_list, g_hat1_list, m_hat_list
+
+
+def compute_residuals(y, g_hat0_list, g_hat1_list, m_hat_list, smpls):
+    u_hat0 = np.full_like(y, np.nan, dtype='float64')
+    u_hat1 = np.full_like(y, np.nan, dtype='float64')
+    g_hat0 = np.full_like(y, np.nan, dtype='float64')
+    g_hat1 = np.full_like(y, np.nan, dtype='float64')
+    m_hat = np.full_like(y, np.nan, dtype='float64')
+    for idx, (_, test_index) in enumerate(smpls):
+        u_hat0[test_index] = y[test_index] - g_hat0_list[idx]
+        u_hat1[test_index] = y[test_index] - g_hat1_list[idx]
+        g_hat0[test_index] = g_hat0_list[idx]
+        g_hat1[test_index] = g_hat1_list[idx]
+        m_hat[test_index] = m_hat_list[idx]
+
+    _check_is_propensity(m_hat, 'learner_m', 'ml_m', smpls, eps=1e-12)
+    return u_hat0, u_hat1, g_hat0, g_hat1, m_hat
+
+
+def apo_dml2(y, x, d, treated, g_hat0_list, g_hat1_list, m_hat_list, smpls, score, normalize_ipw):
+    n_obs = len(y)
+    u_hat0, u_hat1, g_hat0, g_hat1, m_hat = compute_residuals(
+        y, g_hat0_list, g_hat1_list, m_hat_list, smpls
+    )
+
+    if normalize_ipw:
+        m_hat_adj = _normalize_ipw(m_hat, treated)
+    else:
+        m_hat_adj = m_hat
+
+    theta_hat = apo_orth(g_hat0, g_hat1, m_hat_adj,
+                         u_hat0, u_hat1, treated, score)
+
+    se = np.sqrt(var_apo(theta_hat, g_hat0, g_hat1,
+                         m_hat_adj,
+                         u_hat0, u_hat1,
+                         treated, score, n_obs))
+
+    return theta_hat, se
+
+
+def apo_orth(g_hat0, g_hat1, m_hat, u_hat0, u_hat1, treated, score):
+    res = np.mean(g_hat1 + np.divide(np.multiply(treated, u_hat1), m_hat))
+    return res
+
+
+def var_apo(theta, g_hat0, g_hat1, m_hat, u_hat0, u_hat1, treated, score, n_obs):
+    var = 1/n_obs * np.mean(np.power(g_hat1 + np.divide(np.multiply(treated, u_hat1), m_hat) - theta, 2))
+    return var
+
+
+def boot_apo(y, d, treatment_level, thetas, ses, all_g_hat0, all_g_hat1, all_m_hat,
+             all_smpls, score, bootstrap, n_rep_boot,
+             n_rep=1, normalize_ipw=True):
+    treated = (d == treatment_level)
+    all_boot_t_stat = list()
+    for i_rep in range(n_rep):
+        smpls = all_smpls[i_rep]
+        n_obs = len(y)
+
+        weights = draw_weights(bootstrap, n_rep_boot, n_obs)
+        boot_t_stat = boot_apo_single_split(
+            thetas[i_rep], y, d, treated,
+            all_g_hat0[i_rep], all_g_hat1[i_rep], all_m_hat[i_rep], smpls,
+            score, ses[i_rep], weights, n_rep_boot, normalize_ipw)
+        all_boot_t_stat.append(boot_t_stat)
+
+    boot_t_stat = np.hstack(all_boot_t_stat)
+
+    return boot_t_stat
+
+
+def boot_apo_single_split(theta, y, d, treated, g_hat0_list, g_hat1_list, m_hat_list,
+                          smpls, score, se, weights, n_rep_boot, normalize_ipw):
+    _, u_hat1, _, g_hat1, m_hat = compute_residuals(
+        y, g_hat0_list, g_hat1_list, m_hat_list, smpls)
+
+    if normalize_ipw:
+        m_hat_adj = _normalize_ipw(m_hat, treated)
+    else:
+        m_hat_adj = m_hat
+
+    J = -1.0
+    psi = g_hat1 + np.divide(np.multiply(treated, u_hat1), m_hat_adj) - theta
+    boot_t_stat = boot_manual(psi, J, smpls, se, weights, n_rep_boot)
+
+    return boot_t_stat
+
+
+def fit_sensitivity_elements_apo(y, d, treatment_level, all_coef, predictions, score, n_rep):
+    n_treat = 1
+    n_obs = len(y)
+    treated = (d == treatment_level)
+
+    sigma2 = np.full(shape=(1, n_rep, n_treat), fill_value=np.nan)
+    nu2 = np.full(shape=(1, n_rep, n_treat), fill_value=np.nan)
+    psi_sigma2 = np.full(shape=(n_obs, n_rep, n_treat), fill_value=np.nan)
+    psi_nu2 = np.full(shape=(n_obs, n_rep, n_treat), fill_value=np.nan)
+
+    for i_rep in range(n_rep):
+
+        m_hat = predictions['ml_m'][:, i_rep, 0]
+        g_hat0 = predictions['ml_g0'][:, i_rep, 0]
+        g_hat1 = predictions['ml_g1'][:, i_rep, 0]
+
+        weights = np.ones_like(d)
+        weights_bar = np.ones_like(d)
+
+        sigma2_score_element = np.square(y - np.multiply(treated, g_hat1) - np.multiply(1.0-treated, g_hat0))
+        sigma2[0, i_rep, 0] = np.mean(sigma2_score_element)
+        psi_sigma2[:, i_rep, 0] = sigma2_score_element - sigma2[0, i_rep, 0]
+
+        # calc m(W,alpha) and Riesz representer
+        m_alpha = np.multiply(weights, np.multiply(weights_bar, np.divide(1.0, m_hat)))
+        rr = np.multiply(weights_bar, np.divide(treated, m_hat))
+
+        nu2_score_element = np.multiply(2.0, m_alpha) - np.square(rr)
+        nu2[0, i_rep, 0] = np.mean(nu2_score_element)
+        psi_nu2[:, i_rep, 0] = nu2_score_element - nu2[0, i_rep, 0]
+
+    element_dict = {'sigma2': sigma2,
+                    'nu2': nu2,
+                    'psi_sigma2': psi_sigma2,
+                    'psi_nu2': psi_nu2}
+    return element_dict
+
+
+def tune_nuisance_apo(y, x, d, treatment_level, ml_g, ml_m, smpls, score, n_folds_tune,
+                      param_grid_g, param_grid_m):
+    train_cond0 = np.where(d != treatment_level)[0]
+    g0_tune_res = tune_grid_search(y, x, ml_g, smpls, param_grid_g, n_folds_tune,
+                                   train_cond=train_cond0)
+
+    train_cond1 = np.where(d == treatment_level)[0]
+    g1_tune_res = tune_grid_search(y, x, ml_g, smpls, param_grid_g, n_folds_tune,
+                                   train_cond=train_cond1)
+
+    treated = (d == treatment_level)
+    m_tune_res = tune_grid_search(treated, x, ml_m, smpls, param_grid_m, n_folds_tune)
+
+    g0_best_params = [xx.best_params_ for xx in g0_tune_res]
+    g1_best_params = [xx.best_params_ for xx in g1_tune_res]
+    m_best_params = [xx.best_params_ for xx in m_tune_res]
+
+    return g0_best_params, g1_best_params, m_best_params
diff --git a/doubleml/irm/tests/_utils_apos_manual.py b/doubleml/irm/tests/_utils_apos_manual.py
new file mode 100644
index 00000000..cf47d645
--- /dev/null
+++ b/doubleml/irm/tests/_utils_apos_manual.py
@@ -0,0 +1,74 @@
+import numpy as np
+from sklearn.base import clone
+
+from ..apo import DoubleMLAPO
+from ...double_ml_data import DoubleMLData
+
+from ...tests._utils_boot import draw_weights
+
+
+def fit_apos(y, x, d,
+             learner_g, learner_m, treatment_levels, all_smpls, score,
+             n_rep=1, trimming_rule='truncate',
+             normalize_ipw=False, trimming_threshold=1e-2):
+    n_obs = len(y)
+    n_treatments = len(treatment_levels)
+    n_folds = len(all_smpls[0])
+
+    dml_data = DoubleMLData.from_arrays(x, y, d)
+
+    all_apos = np.zeros((n_treatments, n_rep))
+    all_se = np.zeros((n_treatments, n_rep))
+    apo_scaled_score = np.zeros((n_obs, n_treatments, n_rep))
+
+    for i_level in range(n_treatments):
+        model_APO = DoubleMLAPO(
+            dml_data,
+            clone(learner_g),
+            clone(learner_m),
+            treatment_level=treatment_levels[i_level],
+            n_folds=n_folds,
+            n_rep=n_rep,
+            score=score,
+            trimming_rule=trimming_rule,
+            trimming_threshold=trimming_threshold,
+            normalize_ipw=normalize_ipw,
+            draw_sample_splitting=False
+        )
+
+        # synchronize the sample splitting
+        model_APO.set_sample_splitting(all_smpls)
+        model_APO.fit()
+
+        all_apos[i_level, :] = model_APO.all_coef
+        all_se[i_level, :] = model_APO.all_se
+
+        for i_rep in range(n_rep):
+            J = model_APO.psi_deriv[:, i_rep, 0].mean()
+            apo_psi = model_APO.psi[:, i_rep, 0]
+
+            apo_scaled_score[:, i_level, i_rep] = apo_psi / J
+
+    apos = np.median(all_apos, axis=1)
+    se = np.zeros(n_treatments)
+    for i_level in range(n_treatments):
+        se[i_level] = np.sqrt(np.median(np.power(all_se[i_level, :], 2) * n_obs +
+                                        np.power(all_apos[i_level, :] - all_apos[i_level], 2)) / n_obs)
+
+    res = {'apos': apos, 'se': se,
+           'all_apos': all_apos, 'all_se': all_se,
+           'apo_scaled_score': apo_scaled_score}
+    return res
+
+
+def boot_apos(scaled_scores, ses, treatment_levels, all_smpls, n_rep, bootstrap, n_rep_boot):
+    n_treatment_levels = len(treatment_levels)
+    boot_t_stat = np.zeros((n_rep_boot, n_treatment_levels, n_rep))
+    for i_rep in range(n_rep):
+        n_obs = scaled_scores.shape[0]
+        weights = draw_weights(bootstrap, n_rep_boot, n_obs)
+        for i_treatment_levels in range(n_treatment_levels):
+            boot_t_stat[:, i_treatment_levels, i_rep] = np.matmul(weights, scaled_scores[:, i_treatment_levels, i_rep]) / \
+                (n_obs * ses[i_treatment_levels, i_rep])
+
+    return boot_t_stat
diff --git a/doubleml/irm/tests/_utils_irm_manual.py b/doubleml/irm/tests/_utils_irm_manual.py
index c249c9cf..5fbdd174 100644
--- a/doubleml/irm/tests/_utils_irm_manual.py
+++ b/doubleml/irm/tests/_utils_irm_manual.py
@@ -248,11 +248,7 @@ def fit_sensitivity_elements_irm(y, d, all_coef, predictions, score, n_rep):
 
         m_hat = predictions['ml_m'][:, i_rep, 0]
         g_hat0 = predictions['ml_g0'][:, i_rep, 0]
-        if score == 'ATE':
-            g_hat1 = predictions['ml_g1'][:, i_rep, 0]
-        else:
-            assert score == 'ATTE'
-            g_hat1 = y
+        g_hat1 = predictions['ml_g1'][:, i_rep, 0]
 
         if score == 'ATE':
             weights = np.ones_like(d)
diff --git a/doubleml/irm/tests/_utils_qte_manual.py b/doubleml/irm/tests/_utils_qte_manual.py
index ccaf96ee..5c177907 100644
--- a/doubleml/irm/tests/_utils_qte_manual.py
+++ b/doubleml/irm/tests/_utils_qte_manual.py
@@ -1,7 +1,5 @@
 import numpy as np
 from sklearn.base import clone
-import pandas as pd
-from scipy.stats import norm
 
 from ..pq import DoubleMLPQ
 from ...double_ml_data import DoubleMLData
@@ -99,21 +97,3 @@ def boot_qte(scaled_scores, ses, quantiles, all_smpls, n_rep, bootstrap, n_rep_b
                 (n_obs * ses[i_quant, i_rep])
 
     return boot_t_stat
-
-
-def confint_qte(coef, se, quantiles, boot_t_stat=None, joint=True, level=0.95):
-    a = (1 - level)
-    ab = np.array([a / 2, 1. - a / 2])
-    if joint:
-        assert boot_t_stat.shape[2] == 1
-        sim = np.amax(np.abs(boot_t_stat[:, :, 0]), 1)
-        hatc = np.quantile(sim, 1 - a)
-        ci = np.vstack((coef - se * hatc, coef + se * hatc)).T
-    else:
-        fac = norm.ppf(ab)
-        ci = np.vstack((coef + se * fac[0], coef + se * fac[1])).T
-
-    df_ci = pd.DataFrame(ci,
-                         columns=['{:.1f} %'.format(i * 100) for i in ab],
-                         index=quantiles)
-    return df_ci
diff --git a/doubleml/irm/tests/conftest.py b/doubleml/irm/tests/conftest.py
index 3f57b422..6fe207b0 100644
--- a/doubleml/irm/tests/conftest.py
+++ b/doubleml/irm/tests/conftest.py
@@ -14,8 +14,7 @@ def _g(x):
 
 @pytest.fixture(scope='session',
                 params=[(500, 10),
-                        (1000, 20),
-                        (1000, 100)])
+                        (1000, 20)])
 def generate_data_irm(request):
     n_p = request.param
     np.random.seed(1111)
diff --git a/doubleml/irm/tests/test_apo.py b/doubleml/irm/tests/test_apo.py
new file mode 100644
index 00000000..7082e399
--- /dev/null
+++ b/doubleml/irm/tests/test_apo.py
@@ -0,0 +1,245 @@
+import numpy as np
+import pandas as pd
+import pytest
+import math
+
+from sklearn.base import clone
+
+from sklearn.linear_model import LogisticRegression, LinearRegression
+from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
+
+import doubleml as dml
+from doubleml.datasets import make_irm_data_discrete_treatments, make_irm_data
+
+from ...tests._utils import draw_smpls
+from ._utils_apo_manual import fit_apo, boot_apo, fit_sensitivity_elements_apo
+
+
+@pytest.fixture(scope='module',
+                params=[[LinearRegression(),
+                         LogisticRegression(solver='lbfgs', max_iter=250, random_state=42)],
+                        [RandomForestRegressor(max_depth=5, n_estimators=10, random_state=42),
+                         RandomForestClassifier(max_depth=5, n_estimators=10, random_state=42)]])
+def learner(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[False, True])
+def normalize_ipw(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[0.2, 0.15])
+def trimming_threshold(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[0, 1])
+def treatment_level(request):
+    return request.param
+
+
+@pytest.fixture(scope='module')
+def dml_apo_fixture(generate_data_irm, learner, normalize_ipw, trimming_threshold, treatment_level):
+    boot_methods = ['normal']
+    n_folds = 2
+    n_rep_boot = 499
+
+    # Set machine learning methods for m & g
+    ml_g = clone(learner[0])
+    ml_m = clone(learner[1])
+
+    np.random.seed(3141)
+    n_obs = 500
+    data_apo = make_irm_data_discrete_treatments(n_obs=n_obs)
+    y = data_apo['y']
+    x = data_apo['x']
+    d = data_apo['d']
+    df_apo = pd.DataFrame(
+        np.column_stack((y, d, x)),
+        columns=['y', 'd'] + ['x' + str(i) for i in range(data_apo['x'].shape[1])]
+    )
+
+    dml_data = dml.DoubleMLData(df_apo, 'y', 'd')
+    all_smpls = draw_smpls(n_obs, n_folds, n_rep=1, groups=d)
+
+    np.random.seed(3141)
+    dml_obj = dml.DoubleMLAPO(dml_data,
+                              ml_g, ml_m,
+                              treatment_level=treatment_level,
+                              n_folds=n_folds,
+                              score='APO',
+                              normalize_ipw=normalize_ipw,
+                              draw_sample_splitting=False,
+                              trimming_threshold=trimming_threshold)
+
+    # synchronize the sample splitting
+    dml_obj.set_sample_splitting(all_smpls=all_smpls)
+    dml_obj.fit()
+
+    np.random.seed(3141)
+    res_manual = fit_apo(y, x, d,
+                         clone(learner[0]), clone(learner[1]),
+                         treatment_level=treatment_level,
+                         all_smpls=all_smpls,
+                         score='APO',
+                         normalize_ipw=normalize_ipw,
+                         trimming_threshold=trimming_threshold)
+
+    np.random.seed(3141)
+    # test with external nuisance predictions
+    dml_obj_ext = dml.DoubleMLAPO(dml_data,
+                                  ml_g, ml_m,
+                                  treatment_level=treatment_level,
+                                  n_folds=n_folds,
+                                  score='APO',
+                                  normalize_ipw=normalize_ipw,
+                                  draw_sample_splitting=False,
+                                  trimming_threshold=trimming_threshold)
+
+    # synchronize the sample splitting
+    dml_obj_ext.set_sample_splitting(all_smpls=all_smpls)
+
+    prediction_dict = {'d': {'ml_g0': dml_obj.predictions['ml_g0'].reshape(-1, 1),
+                             'ml_g1': dml_obj.predictions['ml_g1'].reshape(-1, 1),
+                             'ml_m': dml_obj.predictions['ml_m'].reshape(-1, 1)}}
+    dml_obj_ext.fit(external_predictions=prediction_dict)
+
+    res_dict = {'coef': dml_obj.coef,
+                'coef_manual': res_manual['theta'],
+                'coef_ext': dml_obj_ext.coef,
+                'se': dml_obj.se,
+                'se_manual': res_manual['se'],
+                'se_ext': dml_obj_ext.se,
+                'boot_methods': boot_methods}
+
+    for bootstrap in boot_methods:
+        np.random.seed(3141)
+        boot_t_stat = boot_apo(y, d, treatment_level, res_manual['thetas'], res_manual['ses'],
+                               res_manual['all_g_hat0'], res_manual['all_g_hat1'],
+                               res_manual['all_m_hat'],
+                               all_smpls,
+                               score='APO',
+                               bootstrap=bootstrap,
+                               n_rep_boot=n_rep_boot,
+                               normalize_ipw=normalize_ipw)
+
+        np.random.seed(3141)
+        dml_obj.bootstrap(method=bootstrap, n_rep_boot=n_rep_boot)
+        np.random.seed(3141)
+        dml_obj_ext.bootstrap(method=bootstrap, n_rep_boot=n_rep_boot)
+        res_dict['boot_t_stat' + bootstrap] = dml_obj.boot_t_stat
+        res_dict['boot_t_stat' + bootstrap + '_manual'] = boot_t_stat.reshape(-1, 1, 1)
+        res_dict['boot_t_stat' + bootstrap + '_ext'] = dml_obj_ext.boot_t_stat
+
+    # check if sensitivity score with rho=0 gives equal asymptotic standard deviation
+    dml_obj.sensitivity_analysis(rho=0.0)
+    res_dict['sensitivity_ses'] = dml_obj.sensitivity_params['se']
+
+    # sensitivity tests
+    res_dict['sensitivity_elements'] = dml_obj.sensitivity_elements
+    res_dict['sensitivity_elements_manual'] = fit_sensitivity_elements_apo(y, d,
+                                                                           treatment_level,
+                                                                           all_coef=dml_obj.all_coef,
+                                                                           predictions=dml_obj.predictions,
+                                                                           score='APO',
+                                                                           n_rep=1)
+    return res_dict
+
+
+@pytest.mark.ci
+def test_dml_apo_coef(dml_apo_fixture):
+    assert math.isclose(dml_apo_fixture['coef'][0],
+                        dml_apo_fixture['coef_manual'],
+                        rel_tol=1e-9, abs_tol=1e-4)
+    assert math.isclose(dml_apo_fixture['coef'][0],
+                        dml_apo_fixture['coef_ext'][0],
+                        rel_tol=1e-9, abs_tol=1e-4)
+
+
+@pytest.mark.ci
+def test_dml_apo_se(dml_apo_fixture):
+    assert math.isclose(dml_apo_fixture['se'][0],
+                        dml_apo_fixture['se_manual'],
+                        rel_tol=1e-9, abs_tol=1e-4)
+    assert math.isclose(dml_apo_fixture['se'][0],
+                        dml_apo_fixture['se_ext'][0],
+                        rel_tol=1e-9, abs_tol=1e-4)
+
+
+@pytest.mark.ci
+def test_dml_apo_boot(dml_apo_fixture):
+    for bootstrap in dml_apo_fixture['boot_methods']:
+        assert np.allclose(dml_apo_fixture['boot_t_stat' + bootstrap],
+                           dml_apo_fixture['boot_t_stat' + bootstrap + '_manual'],
+                           rtol=1e-9, atol=1e-4)
+        assert np.allclose(dml_apo_fixture['boot_t_stat' + bootstrap],
+                           dml_apo_fixture['boot_t_stat' + bootstrap + '_ext'],
+                           rtol=1e-9, atol=1e-4)
+
+
+@pytest.mark.ci
+def test_dml_apo_sensitivity_rho0(dml_apo_fixture):
+    assert np.allclose(dml_apo_fixture['se'],
+                       dml_apo_fixture['sensitivity_ses']['lower'],
+                       rtol=1e-9, atol=1e-4)
+    assert np.allclose(dml_apo_fixture['se'],
+                       dml_apo_fixture['sensitivity_ses']['upper'],
+                       rtol=1e-9, atol=1e-4)
+
+
+@pytest.mark.ci
+def test_dml_apo_sensitivity(dml_apo_fixture):
+    sensitivity_element_names = ['sigma2', 'nu2', 'psi_sigma2', 'psi_nu2']
+    for sensitivity_element in sensitivity_element_names:
+        assert np.allclose(dml_apo_fixture['sensitivity_elements'][sensitivity_element],
+                           dml_apo_fixture['sensitivity_elements_manual'][sensitivity_element],
+                           rtol=1e-9, atol=1e-4)
+
+
+@pytest.mark.ci
+def test_dml_apo_capo_gapo(treatment_level):
+    n = 20
+    # collect data
+    np.random.seed(42)
+    obj_dml_data = make_irm_data(n_obs=n, dim_x=2)
+
+    # First stage estimation
+    ml_g = RandomForestRegressor(n_estimators=10)
+    ml_m = RandomForestClassifier(n_estimators=10)
+
+    dml_obj = dml.DoubleMLAPO(obj_dml_data,
+                              ml_m=ml_m,
+                              ml_g=ml_g,
+                              treatment_level=treatment_level,
+                              trimming_threshold=0.05,
+                              n_folds=5)
+
+    dml_obj.fit()
+    # create a random basis
+    random_basis = pd.DataFrame(np.random.normal(0, 1, size=(n, 5)))
+    capo = dml_obj.capo(random_basis)
+    assert isinstance(capo, dml.utils.blp.DoubleMLBLP)
+    assert isinstance(capo.confint(), pd.DataFrame)
+
+    groups_1 = pd.DataFrame(np.column_stack([obj_dml_data.data['X1'] <= -1.0,
+                                             obj_dml_data.data['X1'] > 0.2]),
+                            columns=['Group 1', 'Group 2'])
+    msg = ('At least one group effect is estimated with less than 6 observations.')
+    with pytest.warns(UserWarning, match=msg):
+        gapo_1 = dml_obj.gapo(groups_1)
+    assert isinstance(gapo_1, dml.utils.blp.DoubleMLBLP)
+    assert isinstance(gapo_1.confint(), pd.DataFrame)
+    assert all(gapo_1.confint().index == groups_1.columns.to_list())
+
+    np.random.seed(42)
+    groups_2 = pd.DataFrame(np.random.choice(["1", "2"], n, p=[0.1, 0.9]))
+    msg = ('At least one group effect is estimated with less than 6 observations.')
+    with pytest.warns(UserWarning, match=msg):
+        gapo_2 = dml_obj.gapo(groups_2)
+    assert isinstance(gapo_2, dml.utils.blp.DoubleMLBLP)
+    assert isinstance(gapo_2.confint(), pd.DataFrame)
+    assert all(gapo_2.confint().index == ["Group_1", "Group_2"])
diff --git a/doubleml/irm/tests/test_apo_classifier.py b/doubleml/irm/tests/test_apo_classifier.py
new file mode 100644
index 00000000..860a61ef
--- /dev/null
+++ b/doubleml/irm/tests/test_apo_classifier.py
@@ -0,0 +1,117 @@
+import numpy as np
+import pytest
+import math
+
+from sklearn.base import clone
+
+from sklearn.linear_model import LogisticRegression
+from sklearn.ensemble import RandomForestClassifier
+
+import doubleml as dml
+
+from ...tests._utils import draw_smpls
+from ._utils_apo_manual import fit_apo, boot_apo
+
+
+@pytest.fixture(scope='module',
+                params=[[LogisticRegression(solver='lbfgs', max_iter=250),
+                         LogisticRegression(solver='lbfgs', max_iter=250)],
+                        [RandomForestClassifier(max_depth=2, n_estimators=10, random_state=42),
+                         RandomForestClassifier(max_depth=2, n_estimators=10, random_state=42)]])
+def learner(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[True, False])
+def normalize_ipw(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[0.01, 0.05])
+def trimming_threshold(request):
+    return request.param
+
+
+@pytest.fixture(scope='module')
+def dml_apo_classifier_fixture(generate_data_irm_binary, learner, normalize_ipw, trimming_threshold):
+    boot_methods = ['normal']
+    n_folds = 2
+    n_rep_boot = 499
+
+    treatment_level = 0
+    score = "APO"
+
+    # collect data
+    (x, y, d) = generate_data_irm_binary
+    n_obs = len(y)
+    all_smpls = draw_smpls(n_obs, n_folds, n_rep=1, groups=d)
+
+    # Set machine learning methods for m & g
+    ml_g = clone(learner[0])
+    ml_m = clone(learner[1])
+
+    np.random.seed(3141)
+    obj_dml_data = dml.DoubleMLData.from_arrays(x, y, d)
+    dml_obj = dml.DoubleMLAPO(obj_dml_data,
+                              ml_g, ml_m,
+                              treatment_level=treatment_level,
+                              n_folds=n_folds,
+                              score=score,
+                              normalize_ipw=normalize_ipw,
+                              trimming_threshold=trimming_threshold,
+                              draw_sample_splitting=False)
+    # synchronize the sample splitting
+    dml_obj.set_sample_splitting(all_smpls=all_smpls)
+    dml_obj.fit()
+
+    np.random.seed(3141)
+    res_manual = fit_apo(y, x, d,
+                         clone(learner[0]), clone(learner[1]),
+                         treatment_level,
+                         all_smpls, score,
+                         normalize_ipw=normalize_ipw, trimming_threshold=trimming_threshold)
+
+    res_dict = {'coef': dml_obj.coef,
+                'coef_manual': res_manual['theta'],
+                'se': dml_obj.se,
+                'se_manual': res_manual['se'],
+                'boot_methods': boot_methods}
+
+    for bootstrap in boot_methods:
+        np.random.seed(3141)
+        boot_t_stat = boot_apo(y, d, treatment_level, res_manual['thetas'], res_manual['ses'],
+                               res_manual['all_g_hat0'], res_manual['all_g_hat1'],
+                               res_manual['all_m_hat'],
+                               all_smpls, score, bootstrap, n_rep_boot,
+                               normalize_ipw=normalize_ipw)
+
+        np.random.seed(3141)
+        dml_obj.bootstrap(method=bootstrap, n_rep_boot=n_rep_boot)
+        res_dict['boot_t_stat' + bootstrap] = dml_obj.boot_t_stat
+        res_dict['boot_t_stat' + bootstrap + '_manual'] = boot_t_stat.reshape(-1, 1, 1)
+
+    return res_dict
+
+
+@pytest.mark.ci
+def test_dml_apo_coef(dml_apo_classifier_fixture):
+    assert math.isclose(dml_apo_classifier_fixture['coef'],
+                        dml_apo_classifier_fixture['coef_manual'],
+                        rel_tol=1e-9, abs_tol=1e-4)
+
+
+@pytest.mark.ci
+def test_dml_apo_se(dml_apo_classifier_fixture):
+    assert math.isclose(dml_apo_classifier_fixture['se'],
+                        dml_apo_classifier_fixture['se_manual'],
+                        rel_tol=1e-9, abs_tol=1e-4)
+
+
+@pytest.mark.ci
+def test_dml_apo_boot(dml_apo_classifier_fixture):
+    for bootstrap in dml_apo_classifier_fixture['boot_methods']:
+        assert np.allclose(dml_apo_classifier_fixture['boot_t_stat' + bootstrap],
+                           dml_apo_classifier_fixture['boot_t_stat' + bootstrap + '_manual'],
+                           rtol=1e-9, atol=1e-4)
diff --git a/doubleml/irm/tests/test_apo_exceptions.py b/doubleml/irm/tests/test_apo_exceptions.py
new file mode 100644
index 00000000..31fa6b44
--- /dev/null
+++ b/doubleml/irm/tests/test_apo_exceptions.py
@@ -0,0 +1,198 @@
+import pytest
+import pandas as pd
+import numpy as np
+
+from doubleml import DoubleMLAPO, DoubleMLData
+from doubleml.datasets import make_irm_data_discrete_treatments, make_iivm_data, make_irm_data
+
+from sklearn.linear_model import Lasso, LogisticRegression
+from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
+
+n = 100
+data_apo = make_irm_data_discrete_treatments(n_obs=n)
+df_apo = pd.DataFrame(np.column_stack((data_apo['y'], data_apo['d'], data_apo['x'])),
+                      columns=['y', 'd'] + ['x' + str(i) for i in range(data_apo['x'].shape[1])])
+
+dml_data = DoubleMLData(df_apo, 'y', 'd')
+
+ml_g = Lasso()
+ml_m = LogisticRegression()
+
+
+@pytest.mark.ci
+def test_apo_exception_data():
+    msg = 'The data must be of DoubleMLData or DoubleMLClusterData type.'
+    with pytest.raises(TypeError, match=msg):
+        _ = DoubleMLAPO(pd.DataFrame(), ml_g, ml_m, treatment_level=0)
+
+    msg = 'Only one treatment variable is allowed. Got 2 treatment variables.'
+    with pytest.raises(ValueError, match=msg):
+        dml_data_multiple = DoubleMLData(df_apo, 'y', ['d', 'x1'])
+        _ = DoubleMLAPO(dml_data_multiple, ml_g, ml_m, treatment_level=0)
+
+    dml_data_z = make_iivm_data()
+    msg = r'Incompatible data. z have been set as instrumental variable\(s\).'
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data_z, ml_g, ml_m, treatment_level=0)
+
+    msg = 'The number of treated observations is less than 5. Number of treated observations: 0 for treatment level 1.1.'
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=1.1)
+
+    msg = r'The proportion of observations with treatment level 42 is less than 5\%. Got 0.70\%.'
+    # test warning
+    with pytest.warns(UserWarning, match=msg):
+        data_apo_warn = make_irm_data_discrete_treatments(n_obs=1000)
+        data_apo_warn['d'][0:7] = 42
+        df_apo_warn = pd.DataFrame(
+            np.column_stack((data_apo_warn['y'], data_apo_warn['d'], data_apo_warn['x'])),
+            columns=['y', 'd'] + ['x' + str(i) for i in range(data_apo_warn['x'].shape[1])]
+        )
+        dml_data_warn = DoubleMLData(df_apo_warn, 'y', 'd')
+
+        _ = DoubleMLAPO(dml_data_warn, ml_g, ml_m, treatment_level=42)
+
+
+@pytest.mark.ci
+def test_apo_exception_learner():
+    msg = (r'The ml_g learner LogisticRegression\(\) was identified as classifier but the outcome variable is not'
+           ' binary with values 0 and 1.')
+    with pytest.raises(ValueError, match=msg):
+        ml_g_classifier = LogisticRegression()
+        _ = DoubleMLAPO(dml_data, ml_g_classifier, ml_m, treatment_level=0)
+
+
+@pytest.mark.ci
+def test_apo_exception_scores():
+    msg = 'Invalid score MAR. Valid score APO.'
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0, score='MAR')
+
+
+@pytest.mark.ci
+def test_apo_exception_trimming_rule():
+    msg = 'Invalid trimming_rule discard. Valid trimming_rule truncate.'
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0, trimming_rule='discard')
+
+    # check the trimming_threshold exceptions
+    msg = "trimming_threshold has to be a float. Object of type <class 'str'> passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0,
+                        trimming_rule='truncate', trimming_threshold="0.1")
+
+    msg = 'Invalid trimming_threshold 0.6. trimming_threshold has to be between 0 and 0.5.'
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0,
+                        trimming_rule='truncate', trimming_threshold=0.6)
+
+
+@pytest.mark.ci
+def test_apo_exception_ipw_normalization():
+    msg = "Normalization indicator has to be boolean. Object of type <class 'int'> passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0, normalize_ipw=1)
+
+
+@pytest.mark.ci
+def test_apo_exception_weights():
+    msg = "weights must be a numpy array or dictionary. weights of type <class 'int'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0, weights=1)
+    msg = r"weights must have keys \['weights', 'weights_bar'\]. keys dict_keys\(\['d'\]\) were passed."
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0, weights={'d': [1, 2, 3]})
+
+    # shape checks
+    msg = rf"weights must have shape \({n},\). weights of shape \(1,\) was passed."
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0, weights=np.ones(1))
+    msg = rf"weights must have shape \({n},\). weights of shape \({n}, 2\) was passed."
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0, weights=np.ones((n, 2)))
+
+    msg = rf"weights must have shape \({n},\). weights of shape \(1,\) was passed."
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0,
+                        weights={'weights': np.ones(1), 'weights_bar': np.ones(1)})
+    msg = rf"weights must have shape \({n},\). weights of shape \({n}, 2\) was passed."
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0,
+                        weights={'weights': np.ones((n, 2)), 'weights_bar': np.ones((n, 2))})
+    msg = rf"weights_bar must have shape \({n}, 1\). weights_bar of shape \({n}, 2\) was passed."
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0,
+                        weights={'weights': np.ones(n), 'weights_bar': np.ones((n, 2))})
+
+    # value checks
+    msg = "All weights values must be greater or equal 0."
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0,
+                        weights=-1*np.ones(n,))
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0,
+                        weights={'weights': -1*np.ones(n,), 'weights_bar': np.ones((n, 1))})
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0,
+                        weights={'weights': np.ones(n,), 'weights_bar': -1*np.ones((n, 1))})
+
+    msg = "At least one weight must be non-zero."
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0,
+                        weights=np.zeros((dml_data.d.shape[0], )))
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0,
+                        weights={'weights': np.zeros((dml_data.d.shape[0], )),
+                                 'weights_bar': np.ones((dml_data.d.shape[0], 1))})
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPO(dml_data, ml_g, ml_m, treatment_level=0,
+                        weights={'weights': np.ones((dml_data.d.shape[0], )),
+                                 'weights_bar': np.zeros((dml_data.d.shape[0], 1))})
+
+
+@pytest.mark.ci
+def test_apo_exception_capo_gapo():
+    n = 20
+    # collect data
+    np.random.seed(42)
+    obj_dml_data = make_irm_data(n_obs=n, dim_x=2)
+
+    # First stage estimation
+    ml_g = RandomForestRegressor(n_estimators=10)
+    ml_m = RandomForestClassifier(n_estimators=10)
+
+    dml_obj = DoubleMLAPO(obj_dml_data,
+                          ml_m=ml_m,
+                          ml_g=ml_g,
+                          treatment_level=0)
+
+    dml_obj.fit()
+    # create a random basis
+    random_basis = pd.DataFrame(np.random.normal(0, 1, size=(n, 5)))
+
+    msg = "Invalid score APO_2. Valid score APO."
+    with pytest.raises(ValueError, match=msg):
+        dml_obj._score = 'APO_2'
+        _ = dml_obj.capo(random_basis)
+    # reset the score
+    dml_obj._score = 'APO'
+
+    msg = "Only implemented for one repetition. Number of repetitions is 2."
+    with pytest.raises(NotImplementedError, match=msg):
+        dml_obj._n_rep = 2
+        dml_obj.capo(random_basis)
+    # reset the number of repetitions
+    dml_obj._n_rep = 1
+
+    msg = "Groups must be of DataFrame type. Groups of type <class 'int'> was passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_obj.gapo(1)
+
+    groups_1 = pd.DataFrame(
+        np.column_stack([obj_dml_data.data['X1'] > 0.2, np.ones_like(obj_dml_data.data['X1'])]),
+        columns=['Group 1', 'Group 2']
+    )
+    msg = (r'Columns of groups must be of bool type or int type \(dummy coded\). Alternatively,'
+           ' groups should only contain one column.')
+    with pytest.raises(TypeError, match=msg):
+        _ = dml_obj.gapo(groups_1)
diff --git a/doubleml/irm/tests/test_apo_external_predictions.py b/doubleml/irm/tests/test_apo_external_predictions.py
new file mode 100644
index 00000000..a3f77dea
--- /dev/null
+++ b/doubleml/irm/tests/test_apo_external_predictions.py
@@ -0,0 +1,96 @@
+import pytest
+import numpy as np
+import pandas as pd
+import math
+
+from sklearn.linear_model import LinearRegression, LogisticRegression
+from doubleml import DoubleMLAPO, DoubleMLData
+from doubleml.datasets import make_irm_data_discrete_treatments
+from doubleml.utils import DMLDummyRegressor, DMLDummyClassifier
+
+from ...tests._utils import draw_smpls
+
+
+@pytest.fixture(scope="module", params=[1, 3])
+def n_rep(request):
+    return request.param
+
+
+@pytest.fixture(scope="module", params=[True, False])
+def set_ml_m_ext(request):
+    return request.param
+
+
+@pytest.fixture(scope="module", params=[True, False])
+def set_ml_g_ext(request):
+    return request.param
+
+
+@pytest.fixture(scope="module")
+def doubleml_apo_ext_fixture(n_rep, set_ml_m_ext, set_ml_g_ext):
+
+    score = "APO"
+    treatment_level = 0
+    ext_predictions = {"d": {}}
+
+    np.random.seed(3141)
+    n_obs = 500
+    data_apo = make_irm_data_discrete_treatments(n_obs=n_obs)
+    df_apo = pd.DataFrame(
+        np.column_stack((data_apo['y'], data_apo['d'], data_apo['x'])),
+        columns=['y', 'd'] + ['x' + str(i) for i in range(data_apo['x'].shape[1])]
+    )
+
+    dml_data = DoubleMLData(df_apo, 'y', 'd')
+    d = data_apo['d']
+    all_smpls = draw_smpls(n_obs, n_folds=5, n_rep=n_rep, groups=d)
+
+    kwargs = {
+        "obj_dml_data": dml_data,
+        "score": score,
+        "treatment_level": treatment_level,
+        "n_rep": n_rep,
+        "draw_sample_splitting": False
+    }
+
+    dml_obj = DoubleMLAPO(ml_g=LinearRegression(), ml_m=LogisticRegression(), **kwargs)
+    dml_obj.set_sample_splitting(all_smpls=all_smpls)
+
+    np.random.seed(3141)
+    dml_obj.fit(store_predictions=True)
+
+    if set_ml_m_ext:
+        ext_predictions["d"]["ml_m"] = dml_obj.predictions["ml_m"][:, :, 0]
+        ml_m = DMLDummyClassifier()
+    else:
+        ml_m = LogisticRegression(random_state=42)
+
+    if set_ml_g_ext:
+        ext_predictions["d"]["ml_g0"] = dml_obj.predictions["ml_g0"][:, :, 0]
+        ext_predictions["d"]["ml_g1"] = dml_obj.predictions["ml_g1"][:, :, 0]
+        ml_g = DMLDummyRegressor()
+    else:
+        ml_g = LinearRegression()
+
+    dml_obj_ext = DoubleMLAPO(ml_g=ml_g, ml_m=ml_m, **kwargs)
+    dml_obj_ext.set_sample_splitting(all_smpls=all_smpls)
+
+    np.random.seed(3141)
+    dml_obj_ext.fit(external_predictions=ext_predictions)
+
+    res_dict = {
+        "coef_normal": dml_obj.coef[0],
+        "coef_ext": dml_obj_ext.coef[0]
+    }
+
+    return res_dict
+
+
+@pytest.mark.ci
+def test_doubleml_apo_ext_coef(doubleml_apo_ext_fixture):
+    assert math.isclose(
+        doubleml_apo_ext_fixture["coef_normal"],
+        doubleml_apo_ext_fixture["coef_ext"],
+        rel_tol=1e-9,
+        abs_tol=1e-4
+    )
diff --git a/doubleml/irm/tests/test_apo_tune.py b/doubleml/irm/tests/test_apo_tune.py
new file mode 100644
index 00000000..3a818fca
--- /dev/null
+++ b/doubleml/irm/tests/test_apo_tune.py
@@ -0,0 +1,159 @@
+import numpy as np
+import pytest
+import math
+
+from sklearn.base import clone
+
+from sklearn.linear_model import LogisticRegression
+from sklearn.ensemble import RandomForestRegressor
+
+import doubleml as dml
+
+from ...tests._utils import draw_smpls
+from ._utils_apo_manual import fit_apo, boot_apo, tune_nuisance_apo
+
+
+@pytest.fixture(scope='module',
+                params=[RandomForestRegressor(random_state=42)])
+def learner_g(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[LogisticRegression(random_state=42)])
+def learner_m(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=['APO'])
+def score(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[True, False])
+def normalize_ipw(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[True, False])
+def tune_on_folds(request):
+    return request.param
+
+
+def get_par_grid(learner):
+    if learner.__class__ in [RandomForestRegressor]:
+        par_grid = {'n_estimators': [5, 10, 20]}
+    else:
+        assert learner.__class__ in [LogisticRegression]
+        par_grid = {'C': np.logspace(-4, 2, 10)}
+    return par_grid
+
+
+@pytest.fixture(scope='module')
+def dml_apo_tune_fixture(generate_data_irm, learner_g, learner_m, score, normalize_ipw, tune_on_folds):
+    par_grid = {'ml_g': get_par_grid(learner_g),
+                'ml_m': get_par_grid(learner_m)}
+    n_folds_tune = 4
+
+    boot_methods = ['normal']
+    n_folds = 2
+    n_rep_boot = 499
+    treatment_level = 0
+
+    # collect data
+    (x, y, d) = generate_data_irm
+    n_obs = len(y)
+    all_smpls = draw_smpls(n_obs, n_folds, n_rep=1, groups=d)
+
+    # Set machine learning methods for m & g
+    ml_g = clone(learner_g)
+    ml_m = clone(learner_m)
+
+    np.random.seed(3141)
+    obj_dml_data = dml.DoubleMLData.from_arrays(x, y, d)
+    dml_obj = dml.DoubleMLAPO(obj_dml_data,
+                              ml_g, ml_m,
+                              treatment_level=treatment_level,
+                              n_folds=n_folds,
+                              score=score,
+                              normalize_ipw=normalize_ipw,
+                              draw_sample_splitting=False)
+    # synchronize the sample splitting
+    dml_obj.set_sample_splitting(all_smpls=all_smpls)
+    np.random.seed(3141)
+    # tune hyperparameters
+    tune_res = dml_obj.tune(par_grid, tune_on_folds=tune_on_folds, n_folds_tune=n_folds_tune,
+                            return_tune_res=False)
+    assert isinstance(tune_res, dml.DoubleMLAPO)
+
+    dml_obj.fit()
+
+    np.random.seed(3141)
+    smpls = all_smpls[0]
+
+    if tune_on_folds:
+        g0_params, g1_params, m_params = tune_nuisance_apo(y, x, d, treatment_level,
+                                                           clone(learner_g), clone(learner_m), smpls, score,
+                                                           n_folds_tune,
+                                                           par_grid['ml_g'], par_grid['ml_m'])
+    else:
+        xx = [(np.arange(len(y)), np.array([]))]
+        g0_params, g1_params, m_params = tune_nuisance_apo(y, x, d, treatment_level,
+                                                           clone(learner_g), clone(learner_m), xx, score,
+                                                           n_folds_tune,
+                                                           par_grid['ml_g'], par_grid['ml_m'])
+        g0_params = g0_params * n_folds
+        m_params = m_params * n_folds
+        g1_params = g1_params * n_folds
+
+    res_manual = fit_apo(y, x, d, clone(learner_g), clone(learner_m),
+                         treatment_level,
+                         all_smpls, score,
+                         normalize_ipw=normalize_ipw,
+                         g0_params=g0_params, g1_params=g1_params, m_params=m_params)
+
+    res_dict = {'coef': dml_obj.coef,
+                'coef_manual': res_manual['theta'],
+                'se': dml_obj.se,
+                'se_manual': res_manual['se'],
+                'boot_methods': boot_methods}
+
+    for bootstrap in boot_methods:
+        np.random.seed(3141)
+        boot_t_stat = boot_apo(y, d, treatment_level, res_manual['thetas'], res_manual['ses'],
+                               res_manual['all_g_hat0'], res_manual['all_g_hat1'],
+                               res_manual['all_m_hat'],
+                               all_smpls, score, bootstrap, n_rep_boot,
+                               normalize_ipw=normalize_ipw)
+
+        np.random.seed(3141)
+        dml_obj.bootstrap(method=bootstrap, n_rep_boot=n_rep_boot)
+        res_dict['boot_t_stat' + bootstrap] = dml_obj.boot_t_stat
+        res_dict['boot_t_stat' + bootstrap + '_manual'] = boot_t_stat.reshape(-1, 1, 1)
+
+    return res_dict
+
+
+@pytest.mark.ci
+def test_dml_apo_tune_coef(dml_apo_tune_fixture):
+    assert math.isclose(dml_apo_tune_fixture['coef'],
+                        dml_apo_tune_fixture['coef_manual'],
+                        rel_tol=1e-9, abs_tol=1e-4)
+
+
+@pytest.mark.ci
+def test_dml_apo_tune_se(dml_apo_tune_fixture):
+    assert math.isclose(dml_apo_tune_fixture['se'],
+                        dml_apo_tune_fixture['se_manual'],
+                        rel_tol=1e-9, abs_tol=1e-4)
+
+
+@pytest.mark.ci
+def test_dml_apo_tune_boot(dml_apo_tune_fixture):
+    for bootstrap in dml_apo_tune_fixture['boot_methods']:
+        assert np.allclose(dml_apo_tune_fixture['boot_t_stat' + bootstrap],
+                           dml_apo_tune_fixture['boot_t_stat' + bootstrap + '_manual'],
+                           rtol=1e-9, atol=1e-4)
diff --git a/doubleml/irm/tests/test_apo_weighted_scores.py b/doubleml/irm/tests/test_apo_weighted_scores.py
new file mode 100644
index 00000000..5551e5dd
--- /dev/null
+++ b/doubleml/irm/tests/test_apo_weighted_scores.py
@@ -0,0 +1,120 @@
+import pytest
+import numpy as np
+
+from sklearn.base import clone
+from sklearn.linear_model import LogisticRegression, LinearRegression
+from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
+
+from ...tests._utils import draw_smpls
+import doubleml as dml
+
+
+@pytest.fixture(scope='module',
+                params=[[LinearRegression(),
+                         LogisticRegression(solver='lbfgs', max_iter=250)],
+                        [RandomForestRegressor(max_depth=5, n_estimators=10, random_state=42),
+                         RandomForestClassifier(max_depth=5, n_estimators=10, random_state=42)]])
+def learner(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=['APO'])
+def score(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[1, 3])
+def n_rep(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[False, True])
+def normalize_ipw(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[0.2, 0.15])
+def trimming_threshold(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[0, 1])
+def treatment_level(request):
+    return request.param
+
+
+@pytest.fixture(scope='module')
+def weighted_apo_score_fixture(generate_data_irm, learner, score, n_rep, normalize_ipw, trimming_threshold,
+                               treatment_level):
+    n_folds = 2
+
+    # collect data
+    (x, y, d) = generate_data_irm
+    n_obs = len(y)
+    all_smpls = draw_smpls(n_obs, n_folds, n_rep=1, groups=d)
+    obj_dml_data = dml.DoubleMLData.from_arrays(x, y, d)
+
+    input_args = {
+        "obj_dml_data": obj_dml_data,
+        "ml_g": clone(learner[0]),
+        "ml_m": clone(learner[1]),
+        "treatment_level": treatment_level,
+        "n_folds": n_folds,
+        "n_rep": n_rep,
+        "score": score,
+        "normalize_ipw": normalize_ipw,
+        "trimming_threshold": trimming_threshold,
+        "draw_sample_splitting": False,
+    }
+
+    np.random.seed(42)
+    dml_obj = dml.DoubleMLAPO(**input_args)
+    dml_obj.set_sample_splitting(all_smpls=all_smpls)
+    dml_obj.fit()
+
+    np.random.seed(42)
+    weights = 0.5 * np.ones_like(obj_dml_data.y)
+    dml_obj_weighted = dml.DoubleMLAPO(weights=weights, **input_args)
+    dml_obj_weighted.set_sample_splitting(all_smpls=all_smpls)
+    dml_obj_weighted.fit()
+
+    np.random.seed(42)
+    weights_dict = {
+        'weights': weights,
+        'weights_bar': np.tile(weights[:, np.newaxis], (1, n_rep)),
+    }
+    dml_obj_weighted_dict = dml.DoubleMLAPO(weights=weights_dict, **input_args)
+    dml_obj_weighted_dict.set_sample_splitting(all_smpls=all_smpls)
+    dml_obj_weighted_dict.fit()
+
+    result_dict = {
+        'coef': dml_obj.coef,
+        'weighted_coef': dml_obj_weighted.coef,
+        'weighted_coef_dict': dml_obj_weighted_dict.coef,
+        'default_weights': dml_obj.weights,
+    }
+    return result_dict
+
+
+@pytest.mark.ci
+def test_apo_weighted_coef(weighted_apo_score_fixture):
+    assert np.allclose(0.5 * weighted_apo_score_fixture['coef'],
+                       weighted_apo_score_fixture['weighted_coef'])
+    assert np.allclose(0.5 * weighted_apo_score_fixture['coef'],
+                       weighted_apo_score_fixture['weighted_coef_dict'])
+
+
+@pytest.mark.ci
+def test_apo_default_weights(weighted_apo_score_fixture):
+    assert isinstance(weighted_apo_score_fixture['default_weights'], dict)
+
+    expected_keys = {'weights'}
+    assert set(weighted_apo_score_fixture['default_weights'].keys()) == expected_keys
+
+    assert np.allclose(weighted_apo_score_fixture['default_weights']['weights'],
+                       np.ones_like(weighted_apo_score_fixture['default_weights']['weights']))
diff --git a/doubleml/irm/tests/test_apos.py b/doubleml/irm/tests/test_apos.py
new file mode 100644
index 00000000..92a372ff
--- /dev/null
+++ b/doubleml/irm/tests/test_apos.py
@@ -0,0 +1,290 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from sklearn.base import clone
+
+from sklearn.linear_model import LogisticRegression, LinearRegression
+from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
+
+import doubleml as dml
+from doubleml.datasets import make_irm_data_discrete_treatments, make_irm_data
+
+from ._utils_apos_manual import fit_apos, boot_apos
+from ...tests._utils import confint_manual
+
+
+@pytest.mark.ci
+def test_apo_properties():
+    n = 20
+    # collect data
+    np.random.seed(42)
+    obj_dml_data = make_irm_data(n_obs=n, dim_x=2)
+
+    dml_obj = dml.DoubleMLAPOS(obj_dml_data,
+                               ml_g=RandomForestRegressor(n_estimators=10),
+                               ml_m=RandomForestClassifier(n_estimators=10),
+                               treatment_levels=0)
+
+    # check properties before fit
+    assert dml_obj.n_rep_boot is None
+    assert dml_obj.coef is None
+    assert dml_obj.all_coef is None
+    assert dml_obj.se is None
+    assert dml_obj.all_se is None
+    assert dml_obj.t_stat is None
+    assert dml_obj.pval is None
+    assert dml_obj.n_rep_boot is None
+    assert dml_obj.boot_t_stat is None
+    assert dml_obj.boot_method is None
+    assert dml_obj.sensitivity_elements is None
+    assert dml_obj.sensitivity_params is None
+
+    # check properties after fit
+    dml_obj.fit()
+    assert dml_obj.coef is not None
+    assert dml_obj.all_coef is not None
+    assert dml_obj.se is not None
+    assert dml_obj.all_se is not None
+    assert dml_obj.t_stat is not None
+    assert dml_obj.pval is not None
+    assert dml_obj.n_rep_boot is None
+    assert dml_obj.boot_t_stat is None
+    assert dml_obj.boot_method is None
+    assert dml_obj.sensitivity_elements is not None
+    assert dml_obj.sensitivity_params is None
+
+    # check properties after bootstrap
+    dml_obj.bootstrap()
+    assert dml_obj.n_rep_boot is not None
+    assert dml_obj.boot_t_stat is not None
+    assert dml_obj.boot_method is not None
+
+    # check properties after sensitivity analysis
+    dml_obj.sensitivity_analysis()
+    assert dml_obj.sensitivity_params is not None
+
+
+@pytest.fixture(scope='module',
+                params=[[LinearRegression(),
+                         LogisticRegression(solver='lbfgs', max_iter=250, random_state=42)],
+                        [RandomForestRegressor(max_depth=5, n_estimators=10, random_state=42),
+                         RandomForestClassifier(max_depth=5, n_estimators=10, random_state=42)]])
+def learner(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[1, 5])
+def n_rep(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[False, True])
+def normalize_ipw(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[0.2, 0.15])
+def trimming_threshold(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[[0, 1, 2], [0]])
+def treatment_levels(request):
+    return request.param
+
+
+@pytest.fixture(scope='module')
+def dml_apos_fixture(learner, n_rep, normalize_ipw, trimming_threshold, treatment_levels):
+    boot_methods = ['normal']
+    n_folds = 2
+    n_rep_boot = 499
+
+    np.random.seed(3141)
+    n_obs = 500
+    data = make_irm_data_discrete_treatments(n_obs=n_obs)
+    y = data['y']
+    x = data['x']
+    d = data['d']
+    df = pd.DataFrame(
+        np.column_stack((y, d, x)),
+        columns=['y', 'd'] + ['x' + str(i) for i in range(data['x'].shape[1])]
+    )
+
+    dml_data = dml.DoubleMLData(df, 'y', 'd')
+
+    input_args = {
+        'obj_dml_data': dml_data,
+        'ml_g': clone(learner[0]),
+        'ml_m': clone(learner[1]),
+        "treatment_levels": treatment_levels,
+        "n_folds": n_folds,
+        "n_rep": n_rep,
+        "score": 'APO',
+        "normalize_ipw": normalize_ipw,
+        "trimming_rule": 'truncate',
+        "trimming_threshold": trimming_threshold,
+        }
+
+    unfitted_apos_model = dml.DoubleMLAPOS(**input_args)
+    np.random.seed(42)
+    dml_obj = dml.DoubleMLAPOS(**input_args)
+    dml_obj.fit()
+    # get the sample splitting
+    all_smpls = dml_obj.smpls
+
+    np.random.seed(42)
+    dml_obj_ext_smpls = dml.DoubleMLAPOS(**input_args, draw_sample_splitting=False)
+    dml_obj_ext_smpls.set_sample_splitting(dml_obj.smpls)
+    dml_obj_ext_smpls.fit()
+
+    np.random.seed(42)
+    res_manual = fit_apos(
+        y, x, d,
+        clone(learner[0]), clone(learner[1]),
+        treatment_levels=treatment_levels,
+        all_smpls=all_smpls,
+        n_rep=n_rep,
+        score='APO',
+        trimming_rule='truncate',
+        normalize_ipw=normalize_ipw,
+        trimming_threshold=trimming_threshold)
+
+    ci = dml_obj.confint(joint=False, level=0.95)
+    ci_ext_smpls = dml_obj_ext_smpls.confint(joint=False, level=0.95)
+    ci_manual = confint_manual(
+        res_manual['apos'], res_manual['se'], treatment_levels,
+        boot_t_stat=None, joint=False, level=0.95
+        )
+
+    res_dict = {
+        'coef': dml_obj.coef,
+        'coef_ext_smpls': dml_obj_ext_smpls.coef,
+        'coef_manual': res_manual['apos'],
+        'se': dml_obj.se,
+        'se_ext_smpls': dml_obj_ext_smpls.se,
+        'se_manual': res_manual['se'],
+        'boot_methods': boot_methods,
+        'n_treatment_levels': len(treatment_levels),
+        'n_rep': n_rep,
+        'ci': ci.to_numpy(),
+        'ci_ext_smpls': ci_ext_smpls.to_numpy(),
+        'ci_manual': ci_manual.to_numpy(),
+        'apos_model': dml_obj,
+        'unfitted_apos_model': unfitted_apos_model
+    }
+    if n_rep == 1:
+        for bootstrap in boot_methods:
+            np.random.seed(42)
+            boot_t_stat = boot_apos(res_manual['apo_scaled_score'], res_manual['all_se'], treatment_levels,
+                                    all_smpls, n_rep, bootstrap, n_rep_boot)
+
+            np.random.seed(42)
+            dml_obj.bootstrap(method=bootstrap, n_rep_boot=n_rep_boot)
+
+            res_dict['boot_t_stat_' + bootstrap] = dml_obj.boot_t_stat
+            res_dict['boot_t_stat_' + bootstrap + '_manual'] = boot_t_stat
+
+            ci = dml_obj.confint(joint=True, level=0.95)
+            ci_manual = confint_manual(
+                res_manual['apos'], res_manual['se'], treatment_levels,
+                boot_t_stat=boot_t_stat, joint=True, level=0.95)
+            res_dict['boot_ci_' + bootstrap] = ci.to_numpy()
+            res_dict['boot_ci_' + bootstrap + '_manual'] = ci_manual.to_numpy()
+
+    # causal contrasts
+    if len(treatment_levels) > 1:
+        acc_single = dml_obj.causal_contrast(reference_levels=[treatment_levels[0]])
+        res_dict['causal_contrast_single'] = acc_single
+        acc_multiple = dml_obj.causal_contrast(reference_levels=treatment_levels)
+        res_dict['causal_contrast_multiple'] = acc_multiple
+
+    return res_dict
+
+
+@pytest.mark.ci
+def test_dml_apos_coef(dml_apos_fixture):
+    assert np.allclose(dml_apos_fixture['coef'],
+                       dml_apos_fixture['coef_manual'],
+                       rtol=1e-9, atol=1e-9)
+    assert np.allclose(dml_apos_fixture['coef'],
+                       dml_apos_fixture['coef_ext_smpls'],
+                       rtol=1e-9, atol=1e-9)
+
+
+@pytest.mark.ci
+def test_dml_apos_se(dml_apos_fixture):
+    if dml_apos_fixture['n_rep'] != 1:
+        pytest.skip("Skipping test as n_rep is not 1")
+    assert np.allclose(dml_apos_fixture['se'],
+                       dml_apos_fixture['se_manual'],
+                       rtol=1e-9, atol=1e-9)
+    assert np.allclose(dml_apos_fixture['se'],
+                       dml_apos_fixture['se_ext_smpls'],
+                       rtol=1e-9, atol=1e-9)
+
+
+@pytest.mark.ci
+def test_dml_apos_boot(dml_apos_fixture):
+    if dml_apos_fixture['n_rep'] != 1:
+        pytest.skip("Skipping test as n_rep is not 1")
+    for bootstrap in dml_apos_fixture['boot_methods']:
+        assert np.allclose(dml_apos_fixture['boot_t_stat_' + bootstrap],
+                           dml_apos_fixture['boot_t_stat_' + bootstrap + '_manual'],
+                           rtol=1e-9, atol=1e-4)
+
+
+@pytest.mark.ci
+def test_dml_apos_ci(dml_apos_fixture):
+    if dml_apos_fixture['n_rep'] != 1:
+        pytest.skip("Skipping test as n_rep is not 1")
+    for bootstrap in dml_apos_fixture['boot_methods']:
+        assert np.allclose(dml_apos_fixture['ci'],
+                           dml_apos_fixture['ci_manual'],
+                           rtol=1e-9, atol=1e-4)
+        assert np.allclose(dml_apos_fixture['ci'],
+                           dml_apos_fixture['ci_ext_smpls'],
+                           rtol=1e-9, atol=1e-4)
+        assert np.allclose(dml_apos_fixture['boot_ci_' + bootstrap],
+                           dml_apos_fixture['boot_ci_' + bootstrap + '_manual'],
+                           rtol=1e-9, atol=1e-4)
+
+
+@pytest.mark.ci
+def test_doubleml_apos_return_types(dml_apos_fixture):
+    assert isinstance(dml_apos_fixture['apos_model'].__str__(), str)
+    assert isinstance(dml_apos_fixture['apos_model'].summary, pd.DataFrame)
+
+    assert dml_apos_fixture['apos_model'].all_coef.shape == (
+        dml_apos_fixture['n_treatment_levels'],
+        dml_apos_fixture['n_rep']
+    )
+    assert isinstance(dml_apos_fixture['unfitted_apos_model'].summary, pd.DataFrame)
+    if dml_apos_fixture['n_treatment_levels'] > 1:
+        assert isinstance(dml_apos_fixture['causal_contrast_single'], dml.DoubleMLFramework)
+        assert isinstance(dml_apos_fixture['causal_contrast_multiple'], dml.DoubleMLFramework)
+
+    benchmark = dml_apos_fixture['apos_model'].sensitivity_benchmark(benchmarking_set=['x1'])
+    assert isinstance(benchmark, pd.DataFrame)
+
+
+@pytest.mark.ci
+def test_doubleml_apos_causal_contrast(dml_apos_fixture):
+    if dml_apos_fixture['n_treatment_levels'] == 1:
+        pytest.skip("Skipping test as n_treatment_levels is 1")
+
+    acc_single = dml_apos_fixture['apos_model'].all_coef[1:, ] - dml_apos_fixture['apos_model'].all_coef[0, ]
+    assert np.allclose(dml_apos_fixture['causal_contrast_single'].all_thetas,
+                       acc_single,
+                       rtol=1e-9, atol=1e-9)
+
+    acc_multiple = np.append(acc_single,
+                             dml_apos_fixture['apos_model'].all_coef[2:3, ] - dml_apos_fixture['apos_model'].all_coef[1:2, ],
+                             axis=0)
+    assert np.allclose(dml_apos_fixture['causal_contrast_multiple'].all_thetas,
+                       acc_multiple,
+                       rtol=1e-9, atol=1e-9)
diff --git a/doubleml/irm/tests/test_apos_classfier.py b/doubleml/irm/tests/test_apos_classfier.py
new file mode 100644
index 00000000..9c3e7d35
--- /dev/null
+++ b/doubleml/irm/tests/test_apos_classfier.py
@@ -0,0 +1,200 @@
+import numpy as np
+import pandas as pd
+import pytest
+
+from sklearn.base import clone
+
+from sklearn.linear_model import LogisticRegression
+from sklearn.ensemble import RandomForestClassifier
+
+import doubleml as dml
+from doubleml.datasets import make_irm_data_discrete_treatments
+
+from ._utils_apos_manual import fit_apos, boot_apos
+from ...tests._utils import confint_manual
+
+
+@pytest.fixture(scope='module',
+                params=[[LogisticRegression(solver='lbfgs', max_iter=250),
+                         LogisticRegression(solver='lbfgs', max_iter=250)],
+                        [RandomForestClassifier(max_depth=2, n_estimators=10, random_state=42),
+                         RandomForestClassifier(max_depth=2, n_estimators=10, random_state=42)]])
+def learner(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[1])
+def n_rep(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[False, True])
+def normalize_ipw(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[0.2, 0.15])
+def trimming_threshold(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[[0, 1, 2], [0]])
+def treatment_levels(request):
+    return request.param
+
+
+@pytest.fixture(scope='module')
+def dml_apos_classifier_fixture(learner, n_rep, normalize_ipw, trimming_threshold, treatment_levels):
+    boot_methods = ['normal']
+    n_folds = 2
+    n_rep_boot = 499
+
+    np.random.seed(3141)
+    n_obs = 500
+    data = make_irm_data_discrete_treatments(n_obs=n_obs)
+    y = np.random.binomial(1, 0.5, n_obs)
+    x = data['x']
+    d = data['d']
+    df = pd.DataFrame(
+        np.column_stack((y, d, x)),
+        columns=['y', 'd'] + ['x' + str(i) for i in range(data['x'].shape[1])]
+    )
+
+    dml_data = dml.DoubleMLData(df, 'y', 'd')
+
+    input_args = {
+        'obj_dml_data': dml_data,
+        'ml_g': clone(learner[0]),
+        'ml_m': clone(learner[1]),
+        "treatment_levels": treatment_levels,
+        "n_folds": n_folds,
+        "n_rep": n_rep,
+        "score": 'APO',
+        "normalize_ipw": normalize_ipw,
+        "trimming_rule": 'truncate',
+        "trimming_threshold": trimming_threshold,
+        }
+
+    unfitted_apos_model = dml.DoubleMLAPOS(**input_args)
+    np.random.seed(42)
+    dml_obj = dml.DoubleMLAPOS(**input_args)
+    dml_obj.fit()
+    # get the sample splitting
+    all_smpls = dml_obj.smpls
+
+    np.random.seed(42)
+    dml_obj_ext_smpls = dml.DoubleMLAPOS(**input_args, draw_sample_splitting=False)
+    dml_obj_ext_smpls.set_sample_splitting(dml_obj.smpls)
+    dml_obj_ext_smpls.fit()
+
+    np.random.seed(42)
+    res_manual = fit_apos(
+        y, x, d,
+        clone(learner[0]), clone(learner[1]),
+        treatment_levels=treatment_levels,
+        all_smpls=all_smpls,
+        score='APO',
+        trimming_rule='truncate',
+        normalize_ipw=normalize_ipw,
+        trimming_threshold=trimming_threshold)
+
+    ci = dml_obj.confint(joint=False, level=0.95)
+    ci_ext_smpls = dml_obj_ext_smpls.confint(joint=False, level=0.95)
+    ci_manual = confint_manual(
+        res_manual['apos'], res_manual['se'], treatment_levels,
+        boot_t_stat=None, joint=False, level=0.95
+        )
+
+    res_dict = {
+        'coef': dml_obj.coef,
+        'coef_ext_smpls': dml_obj_ext_smpls.coef,
+        'coef_manual': res_manual['apos'],
+        'se': dml_obj.se,
+        'se_ext_smpls': dml_obj_ext_smpls.se,
+        'se_manual': res_manual['se'],
+        'boot_methods': boot_methods,
+        'n_treatment_levels': len(treatment_levels),
+        'n_rep': n_rep,
+        'ci': ci.to_numpy(),
+        'ci_ext_smpls': ci_ext_smpls.to_numpy(),
+        'ci_manual': ci_manual.to_numpy(),
+        'apos_model': dml_obj,
+        'unfitted_apos_model': unfitted_apos_model
+    }
+
+    for bootstrap in boot_methods:
+        np.random.seed(42)
+        boot_t_stat = boot_apos(res_manual['apo_scaled_score'], res_manual['all_se'], treatment_levels,
+                                all_smpls, n_rep, bootstrap, n_rep_boot)
+
+        np.random.seed(42)
+        dml_obj.bootstrap(method=bootstrap, n_rep_boot=n_rep_boot)
+
+        res_dict['boot_t_stat_' + bootstrap] = dml_obj.boot_t_stat
+        res_dict['boot_t_stat_' + bootstrap + '_manual'] = boot_t_stat
+
+        ci = dml_obj.confint(joint=True, level=0.95)
+        ci_manual = confint_manual(res_manual['apos'], res_manual['se'], treatment_levels,
+                                   boot_t_stat=boot_t_stat, joint=True, level=0.95)
+        res_dict['boot_ci_' + bootstrap] = ci.to_numpy()
+        res_dict['boot_ci_' + bootstrap + '_manual'] = ci_manual.to_numpy()
+
+    return res_dict
+
+
+@pytest.mark.ci
+def test_dml_apos_coef(dml_apos_classifier_fixture):
+    assert np.allclose(dml_apos_classifier_fixture['coef'],
+                       dml_apos_classifier_fixture['coef_manual'],
+                       rtol=1e-9, atol=1e-9)
+    assert np.allclose(dml_apos_classifier_fixture['coef'],
+                       dml_apos_classifier_fixture['coef_ext_smpls'],
+                       rtol=1e-9, atol=1e-9)
+
+
+@pytest.mark.ci
+def test_dml_apos_se(dml_apos_classifier_fixture):
+    assert np.allclose(dml_apos_classifier_fixture['se'],
+                       dml_apos_classifier_fixture['se_manual'],
+                       rtol=1e-9, atol=1e-9)
+    assert np.allclose(dml_apos_classifier_fixture['se'],
+                       dml_apos_classifier_fixture['se_ext_smpls'],
+                       rtol=1e-9, atol=1e-9)
+
+
+@pytest.mark.ci
+def test_dml_apos_boot(dml_apos_classifier_fixture):
+    for bootstrap in dml_apos_classifier_fixture['boot_methods']:
+        assert np.allclose(dml_apos_classifier_fixture['boot_t_stat_' + bootstrap],
+                           dml_apos_classifier_fixture['boot_t_stat_' + bootstrap + '_manual'],
+                           rtol=1e-9, atol=1e-4)
+
+
+@pytest.mark.ci
+def test_dml_apos_ci(dml_apos_classifier_fixture):
+    for bootstrap in dml_apos_classifier_fixture['boot_methods']:
+        assert np.allclose(dml_apos_classifier_fixture['ci'],
+                           dml_apos_classifier_fixture['ci_manual'],
+                           rtol=1e-9, atol=1e-4)
+        assert np.allclose(dml_apos_classifier_fixture['ci'],
+                           dml_apos_classifier_fixture['ci_ext_smpls'],
+                           rtol=1e-9, atol=1e-4)
+        assert np.allclose(dml_apos_classifier_fixture['boot_ci_' + bootstrap],
+                           dml_apos_classifier_fixture['boot_ci_' + bootstrap + '_manual'],
+                           rtol=1e-9, atol=1e-4)
+
+
+@pytest.mark.ci
+def test_doubleml_apos_return_types(dml_apos_classifier_fixture):
+    assert isinstance(dml_apos_classifier_fixture['apos_model'].__str__(), str)
+    assert isinstance(dml_apos_classifier_fixture['apos_model'].summary, pd.DataFrame)
+
+    assert dml_apos_classifier_fixture['apos_model'].all_coef.shape == (
+        dml_apos_classifier_fixture['n_treatment_levels'],
+        dml_apos_classifier_fixture['n_rep']
+    )
+    assert isinstance(dml_apos_classifier_fixture['unfitted_apos_model'].summary, pd.DataFrame)
diff --git a/doubleml/irm/tests/test_apos_exceptions.py b/doubleml/irm/tests/test_apos_exceptions.py
new file mode 100644
index 00000000..0c20efe5
--- /dev/null
+++ b/doubleml/irm/tests/test_apos_exceptions.py
@@ -0,0 +1,185 @@
+import pytest
+import pandas as pd
+import numpy as np
+
+from doubleml import DoubleMLAPOS, DoubleMLData
+from doubleml.datasets import make_irm_data_discrete_treatments, make_iivm_data
+
+from sklearn.linear_model import Lasso, LogisticRegression
+
+n = 100
+data = make_irm_data_discrete_treatments(n_obs=n)
+df = pd.DataFrame(
+    np.column_stack((data['y'], data['d'], data['x'])),
+    columns=['y', 'd'] + ['x' + str(i) for i in range(data['x'].shape[1])]
+)
+
+dml_data = DoubleMLData(df, 'y', 'd')
+
+ml_g = Lasso()
+ml_m = LogisticRegression()
+
+
+@pytest.mark.ci
+def test_apos_exception_data():
+    msg = 'The data must be of DoubleMLData or DoubleMLClusterData type.'
+    with pytest.raises(TypeError, match=msg):
+        _ = DoubleMLAPOS(pd.DataFrame(), ml_g, ml_m, treatment_levels=0)
+
+    msg = 'The data must not contain instrumental variables.'
+    with pytest.raises(ValueError, match=msg):
+        dml_data_z = make_iivm_data()
+        _ = DoubleMLAPOS(dml_data_z, ml_g, ml_m, treatment_levels=0)
+
+    msg = ('Invalid reference_levels. reference_levels has to be an iterable subset or '
+           'a single element of the unique treatment levels in the data.')
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=[1.1])
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=1.1)
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=[1, 2.2])
+
+
+@pytest.mark.ci
+def test_apos_exception_learner():
+    msg = (r'The ml_g learner LogisticRegression\(\) was identified as classifier but the outcome variable is not'
+           ' binary with values 0 and 1.')
+    with pytest.raises(ValueError, match=msg):
+        ml_g_classifier = LogisticRegression()
+        _ = DoubleMLAPOS(dml_data, ml_g_classifier, ml_m, treatment_levels=0)
+
+
+@pytest.mark.ci
+def test_apos_exception_scores():
+    msg = 'Invalid score MAR. Valid score APO.'
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=0, score='MAR')
+
+
+@pytest.mark.ci
+def test_apos_exception_trimming_rule():
+    msg = 'Invalid trimming_rule discard. Valid trimming_rule truncate.'
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=0, trimming_rule='discard')
+
+    # check the trimming_threshold exceptions
+    msg = "trimming_threshold has to be a float. Object of type <class 'str'> passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=0,
+                         trimming_rule='truncate', trimming_threshold="0.1")
+
+    msg = 'Invalid trimming_threshold 0.6. trimming_threshold has to be between 0 and 0.5.'
+    with pytest.raises(ValueError, match=msg):
+        _ = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=0,
+                         trimming_rule='truncate', trimming_threshold=0.6)
+
+
+@pytest.mark.ci
+def test_apos_exception_ipw_normalization():
+    msg = "Normalization indicator has to be boolean. Object of type <class 'int'> passed."
+    with pytest.raises(TypeError, match=msg):
+        _ = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=0, normalize_ipw=1)
+
+
+@pytest.mark.ci
+def test_apos_exception_properties_and_methods():
+    # properties
+    dml_obj = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=0, draw_sample_splitting=False)
+    msg = r'Sample splitting not specified. Draw samples via .draw_sample splitting\(\). External samples not implemented yet.'
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_obj.smpls
+
+    # methods
+    dml_obj = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=0)
+    msg = r'Apply fit\(\) before confint\(\).'
+    with pytest.raises(ValueError, match=msg):
+        dml_obj.confint()
+    msg = r'Apply fit\(\) before bootstrap\(\).'
+    with pytest.raises(ValueError, match=msg):
+        dml_obj.bootstrap()
+    msg = r'Apply fit\(\) before sensitivity_analysis\(\).'
+    with pytest.raises(ValueError, match=msg):
+        dml_obj.sensitivity_analysis()
+    msg = r'Apply fit\(\) before sensitivity_plot\(\).'
+    with pytest.raises(ValueError, match=msg):
+        dml_obj.sensitivity_plot()
+    msg = r'Apply sensitivity_analysis\(\) before sensitivity_summary.'
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_obj.sensitivity_summary
+
+
+@pytest.mark.ci
+def test_apos_exception_ext_pred():
+    dml_obj = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=0)
+    external_predictions = [0, 1]
+    msg = r'external_predictions must be a dictionary. Object of type <class \'list\'> passed.'
+    with pytest.raises(TypeError, match=msg):
+        dml_obj.fit(external_predictions=external_predictions)
+
+    # test with a level subset
+    external_predictions = {
+        0: "dummy",
+        1: "dummy"
+    }
+    msg = (
+        r"external_predictions must be a subset of all treatment levels\. "
+        r"Expected keys: \{0\}\. "
+        r"Passed keys: \{0, 1\}\."
+    )
+    with pytest.raises(ValueError, match=msg):
+        dml_obj.fit(external_predictions=external_predictions)
+
+    external_predictions = {
+        0: "dummy",
+    }
+    msg = r'external_predictions\[0\] must be a dictionary. Object of type <class \'str\'> passed.'
+    with pytest.raises(TypeError, match=msg):
+        dml_obj.fit(external_predictions=external_predictions)
+
+    external_predictions = {
+        0: {"ml_g": "dummy"}
+    }
+    msg = r"external_predictions\[0\] must be a subset of \{.*\}. Passed keys: \{'ml_g'\}\."
+    with pytest.raises(ValueError, match=msg):
+        dml_obj.fit(external_predictions=external_predictions)
+
+    # test with all levels
+    dml_obj = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=[0, 1, 2, 3])
+    external_predictions = {
+        0: "dummy",
+        1: "dummy",
+        4: "dummy"
+    }
+    msg = (
+        r"external_predictions must be a subset of all treatment levels\. "
+        r"Expected keys: \{0, 1, 2, 3\}\. "
+        r"Passed keys: \{0, 1, 4\}\."
+    )
+    with pytest.raises(ValueError, match=msg):
+        dml_obj.fit(external_predictions=external_predictions)
+
+
+@pytest.mark.ci
+def test_causal_contrast_exceptions():
+    msg = r"Apply fit\(\) before causal_contrast\(\)."
+    with pytest.raises(ValueError, match=msg):
+        dml_obj = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=[0, 1])
+        dml_obj.causal_contrast(reference_levels=0)
+
+    msg = 'Only one treatment level. No causal contrast can be computed.'
+    with pytest.raises(ValueError, match=msg):
+        dml_obj = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=[0])
+        dml_obj.fit()
+        dml_obj.causal_contrast(reference_levels=0)
+
+    dml_obj = DoubleMLAPOS(dml_data, ml_g, ml_m, treatment_levels=[0, 1])
+    dml_obj.fit()
+    msg = ('Invalid reference_levels. reference_levels has to be an iterable subset of treatment_levels or '
+           'a single treatment level.')
+    with pytest.raises(ValueError, match=msg):
+        dml_obj.causal_contrast(reference_levels=2)
+    with pytest.raises(ValueError, match=msg):
+        dml_obj.causal_contrast(reference_levels=[2])
+    with pytest.raises(ValueError, match=msg):
+        dml_obj.causal_contrast(reference_levels=[0, 2])
diff --git a/doubleml/irm/tests/test_apos_external_predictions.py b/doubleml/irm/tests/test_apos_external_predictions.py
new file mode 100644
index 00000000..b6a2c8ee
--- /dev/null
+++ b/doubleml/irm/tests/test_apos_external_predictions.py
@@ -0,0 +1,118 @@
+import pytest
+import numpy as np
+import pandas as pd
+import math
+
+from sklearn.linear_model import LinearRegression, LogisticRegression
+from doubleml import DoubleMLAPOS, DoubleMLData
+from doubleml.datasets import make_irm_data_discrete_treatments
+from doubleml.utils import DMLDummyRegressor, DMLDummyClassifier
+
+from ...tests._utils import draw_smpls
+
+
+@pytest.fixture(scope="module", params=[1, 3])
+def n_rep(request):
+    return request.param
+
+
+@pytest.fixture(scope="module", params=[[0, 1, 2, 3], [0, 1]])
+def treatment_levels(request):
+    return request.param
+
+
+@pytest.fixture(scope="module", params=[True, False])
+def set_ml_m_ext(request):
+    return request.param
+
+
+@pytest.fixture(scope="module", params=[True, False])
+def set_ml_g_ext(request):
+    return request.param
+
+
+@pytest.fixture(scope="module")
+def doubleml_apos_ext_fixture(n_rep, treatment_levels, set_ml_m_ext, set_ml_g_ext):
+    score = "APO"
+    ext_predictions = {
+        treatment_level: {} for treatment_level in treatment_levels
+    }
+
+    np.random.seed(3141)
+    n_obs = 500
+    data_apo = make_irm_data_discrete_treatments(n_obs=n_obs)
+    df_apo = pd.DataFrame(
+        np.column_stack((data_apo['y'], data_apo['d'], data_apo['x'])),
+        columns=['y', 'd'] + ['x' + str(i) for i in range(data_apo['x'].shape[1])]
+    )
+
+    dml_data = DoubleMLData(df_apo, 'y', 'd')
+    d = data_apo['d']
+    all_smpls = draw_smpls(n_obs, n_folds=5, n_rep=n_rep, groups=d)
+
+    kwargs = {
+        "obj_dml_data": dml_data,
+        "score": score,
+        "treatment_levels": treatment_levels,
+        "n_rep": n_rep,
+        "draw_sample_splitting": False
+    }
+
+    dml_obj = DoubleMLAPOS(ml_g=LinearRegression(), ml_m=LogisticRegression(), **kwargs)
+    dml_obj.set_sample_splitting(all_smpls=all_smpls)
+
+    np.random.seed(3141)
+    dml_obj.fit(store_predictions=True)
+
+    if set_ml_m_ext:
+        for i_treatment_level, treatment_level in enumerate(treatment_levels):
+            ext_predictions[treatment_level]["ml_m"] = dml_obj.modellist[i_treatment_level].predictions["ml_m"][:, :, 0]
+        ml_m = DMLDummyClassifier()
+    else:
+        ml_m = LogisticRegression(random_state=42)
+
+    if set_ml_g_ext:
+        for i_treatment_level, treatment_level in enumerate(treatment_levels):
+            ext_predictions[treatment_level]["ml_g0"] = dml_obj.modellist[i_treatment_level].predictions["ml_g0"][:, :, 0]
+            ext_predictions[treatment_level]["ml_g1"] = dml_obj.modellist[i_treatment_level].predictions["ml_g1"][:, :, 0]
+        ml_g = DMLDummyRegressor()
+    else:
+        ml_g = LinearRegression()
+
+    dml_obj_ext = DoubleMLAPOS(ml_g=ml_g, ml_m=ml_m, **kwargs)
+    dml_obj_ext.set_sample_splitting(all_smpls=all_smpls)
+
+    np.random.seed(3141)
+    dml_obj_ext.fit(external_predictions=ext_predictions)
+
+    res_dict = {
+        "coef_normal": dml_obj.coef[0],
+        "coef_ext": dml_obj_ext.coef[0],
+        "dml_obj": dml_obj,
+        "dml_obj_ext": dml_obj_ext,
+        "treatment_levels": treatment_levels
+    }
+
+    return res_dict
+
+
+@pytest.mark.ci
+def test_doubleml_apos_ext_coef(doubleml_apos_ext_fixture):
+    assert math.isclose(
+        doubleml_apos_ext_fixture["coef_normal"],
+        doubleml_apos_ext_fixture["coef_ext"],
+        rel_tol=1e-9,
+        abs_tol=1e-4
+    )
+
+
+@pytest.mark.ci
+def test_doubleml_apos_ext_pred_nuisance(doubleml_apos_ext_fixture):
+    for i_level, _ in enumerate(doubleml_apos_ext_fixture["treatment_levels"]):
+        for nuisance_key in ["ml_g0", "ml_g1", "ml_m"]:
+            assert np.allclose(
+                doubleml_apos_ext_fixture["dml_obj"].modellist[i_level].nuisance_loss[nuisance_key],
+                doubleml_apos_ext_fixture["dml_obj_ext"].modellist[i_level].nuisance_loss[nuisance_key],
+                rtol=1e-9,
+                atol=1e-4
+            )
diff --git a/doubleml/irm/tests/test_apos_weighted_scores.py b/doubleml/irm/tests/test_apos_weighted_scores.py
new file mode 100644
index 00000000..3ab8db6a
--- /dev/null
+++ b/doubleml/irm/tests/test_apos_weighted_scores.py
@@ -0,0 +1,128 @@
+import pytest
+import numpy as np
+import pandas as pd
+
+from sklearn.base import clone
+from sklearn.linear_model import LogisticRegression, LinearRegression
+from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
+
+import doubleml as dml
+from doubleml.datasets import make_irm_data_discrete_treatments
+
+
+@pytest.fixture(scope='module',
+                params=[[LinearRegression(),
+                         LogisticRegression(solver='lbfgs', max_iter=250)],
+                        [RandomForestRegressor(max_depth=5, n_estimators=10, random_state=42),
+                         RandomForestClassifier(max_depth=5, n_estimators=10, random_state=42)]])
+def learner(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=['APO'])
+def score(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[1, 3])
+def n_rep(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[False, True])
+def normalize_ipw(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[0.2, 0.15])
+def trimming_threshold(request):
+    return request.param
+
+
+@pytest.fixture(scope='module',
+                params=[[0, 1, 2], [0]])
+def treatment_levels(request):
+    return request.param
+
+
+@pytest.fixture(scope='module')
+def weighted_apos_score_fixture(learner, score, n_rep, normalize_ipw, trimming_threshold,
+                                treatment_levels):
+    n_obs = 500
+    n_folds = 2
+
+    # collect data
+    data = make_irm_data_discrete_treatments(n_obs=n_obs)
+    y = data['y']
+    x = data['x']
+    d = data['d']
+    df = pd.DataFrame(
+        np.column_stack((y, d, x)),
+        columns=['y', 'd'] + ['x' + str(i) for i in range(data['x'].shape[1])]
+    )
+
+    obj_dml_data = dml.DoubleMLData(df, 'y', 'd')
+
+    input_args = {
+        'obj_dml_data': obj_dml_data,
+        'ml_g': clone(learner[0]),
+        'ml_m': clone(learner[1]),
+        'treatment_levels': treatment_levels,
+        'n_folds': n_folds,
+        'n_rep': n_rep,
+        'score': score,
+        'normalize_ipw': normalize_ipw,
+        'trimming_threshold': trimming_threshold,
+        'trimming_rule': 'truncate'
+    }
+
+    np.random.seed(42)
+    dml_obj = dml.DoubleMLAPOS(**input_args)
+    dml_obj.fit()
+
+    np.random.seed(42)
+    weights = 0.5 * np.ones_like(obj_dml_data.y)
+    dml_obj_weighted = dml.DoubleMLAPOS(draw_sample_splitting=False,
+                                        weights=weights,
+                                        **input_args)
+    dml_obj_weighted.set_sample_splitting(all_smpls=dml_obj.smpls)
+    dml_obj_weighted.fit()
+
+    np.random.seed(42)
+    weights_dict = {
+        'weights': weights,
+        'weights_bar': np.tile(weights[:, np.newaxis], (1, n_rep)),
+    }
+    dml_obj_weighted_dict = dml.DoubleMLAPOS(draw_sample_splitting=False,
+                                             weights=weights_dict,
+                                             **input_args)
+    dml_obj_weighted_dict.set_sample_splitting(all_smpls=dml_obj.smpls)
+    dml_obj_weighted_dict.fit()
+
+    result_dict = {
+        'coef': dml_obj.coef,
+        'weighted_coef': dml_obj_weighted.coef,
+        'weighted_coef_dict': dml_obj_weighted_dict.coef,
+        'default_weights': dml_obj.weights,
+    }
+    return result_dict
+
+
+@pytest.mark.ci
+def test_apos_weighted_coef(weighted_apos_score_fixture):
+    assert np.allclose(0.5 * weighted_apos_score_fixture['coef'],
+                       weighted_apos_score_fixture['weighted_coef'])
+    assert np.allclose(0.5 * weighted_apos_score_fixture['coef'],
+                       weighted_apos_score_fixture['weighted_coef_dict'])
+
+
+@pytest.mark.ci
+def test_apos_default_weights(weighted_apos_score_fixture):
+    assert isinstance(weighted_apos_score_fixture['default_weights'], np.ndarray)
+
+    assert np.allclose(weighted_apos_score_fixture['default_weights'],
+                       np.ones_like(weighted_apos_score_fixture['default_weights']))
diff --git a/doubleml/irm/tests/test_irm_classifier.py b/doubleml/irm/tests/test_irm_classifier.py
index 46cdfb77..cfea434d 100644
--- a/doubleml/irm/tests/test_irm_classifier.py
+++ b/doubleml/irm/tests/test_irm_classifier.py
@@ -1,4 +1,3 @@
-
 import numpy as np
 import pytest
 import math
diff --git a/doubleml/irm/tests/test_qte.py b/doubleml/irm/tests/test_qte.py
index bdcd695d..7c7b8c1d 100644
--- a/doubleml/irm/tests/test_qte.py
+++ b/doubleml/irm/tests/test_qte.py
@@ -9,8 +9,8 @@
 from sklearn.linear_model import LogisticRegression
 from sklearn.ensemble import RandomForestClassifier
 
-from ...tests._utils import draw_smpls
-from ._utils_qte_manual import fit_qte, boot_qte, confint_qte
+from ...tests._utils import draw_smpls, confint_manual
+from ._utils_qte_manual import fit_qte, boot_qte
 
 from doubleml.datasets import make_irm_data
 from ...utils._estimation import _default_kde
@@ -94,8 +94,8 @@ def dml_qte_fixture(generate_data_quantiles, learner, normalize_ipw, kde):
                          draw_sample_splitting=True)
 
     ci = dml_qte_obj.confint(joint=False, level=0.95)
-    ci_manual = confint_qte(res_manual['qte'], res_manual['se'], quantiles,
-                            boot_t_stat=None, joint=False, level=0.95)
+    ci_manual = confint_manual(res_manual['qte'], res_manual['se'], quantiles,
+                               boot_t_stat=None, joint=False, level=0.95)
     res_dict = {'coef': dml_qte_obj.coef,
                 'coef_manual': res_manual['qte'],
                 'coef_ext_smpls': dml_qte_obj_ext_smpls.coef,
@@ -120,8 +120,8 @@ def dml_qte_fixture(generate_data_quantiles, learner, normalize_ipw, kde):
         res_dict['boot_t_stat_' + bootstrap + '_manual'] = boot_t_stat
 
         ci = dml_qte_obj.confint(joint=True, level=0.95)
-        ci_manual = confint_qte(res_manual['qte'], res_manual['se'], quantiles,
-                                boot_t_stat=boot_t_stat, joint=True, level=0.95)
+        ci_manual = confint_manual(res_manual['qte'], res_manual['se'], quantiles,
+                                   boot_t_stat=boot_t_stat, joint=True, level=0.95)
         res_dict['boot_ci_' + bootstrap] = ci.to_numpy()
         res_dict['boot_ci_' + bootstrap + '_manual'] = ci_manual.to_numpy()
     return res_dict
@@ -181,6 +181,7 @@ def test_doubleml_qte_exceptions():
         _ = dml_obj.smpls
 
 
+@pytest.mark.ci
 def test_doubleml_qte_return_types(dml_qte_fixture):
     assert isinstance(dml_qte_fixture['qte_model'].__str__(), str)
     assert isinstance(dml_qte_fixture['qte_model'].summary, pd.DataFrame)
diff --git a/doubleml/tests/_utils.py b/doubleml/tests/_utils.py
index b6c8fbc2..18ceef88 100644
--- a/doubleml/tests/_utils.py
+++ b/doubleml/tests/_utils.py
@@ -1,8 +1,21 @@
 import numpy as np
 from sklearn.model_selection import KFold, GridSearchCV, StratifiedKFold
 from sklearn.base import clone
+import pandas as pd
+from scipy.stats import norm
 
 from ..utils._estimation import _var_est, _aggregate_coefs_and_ses
+from ..double_ml_data import DoubleMLBaseData
+
+
+class DummyDataClass(DoubleMLBaseData):
+    def __init__(self,
+                 data):
+        DoubleMLBaseData.__init__(self, data)
+
+    @property
+    def n_coefs(self):
+        return 1
 
 
 def draw_smpls(n_obs, n_folds, n_rep=1, groups=None):
@@ -111,3 +124,21 @@ def generate_dml_dict(psi_a, psi_b):
     }
 
     return doubleml_dict
+
+
+def confint_manual(coef, se, index_names, boot_t_stat=None, joint=True, level=0.95):
+    a = (1 - level)
+    ab = np.array([a / 2, 1. - a / 2])
+    if joint:
+        assert boot_t_stat.shape[2] == 1
+        sim = np.amax(np.abs(boot_t_stat[:, :, 0]), 1)
+        hatc = np.quantile(sim, 1 - a)
+        ci = np.vstack((coef - se * hatc, coef + se * hatc)).T
+    else:
+        fac = norm.ppf(ab)
+        ci = np.vstack((coef + se * fac[0], coef + se * fac[1])).T
+
+    df_ci = pd.DataFrame(ci,
+                         columns=['{:.1f} %'.format(i * 100) for i in ab],
+                         index=index_names)
+    return df_ci
diff --git a/doubleml/tests/test_datasets.py b/doubleml/tests/test_datasets.py
index d1227869..80f71d0e 100644
--- a/doubleml/tests/test_datasets.py
+++ b/doubleml/tests/test_datasets.py
@@ -5,7 +5,8 @@
 from doubleml import DoubleMLData, DoubleMLClusterData
 from doubleml.datasets import fetch_401K, fetch_bonus, make_plr_CCDDHNR2018, make_plr_turrell2018, \
     make_irm_data, make_iivm_data, _make_pliv_data, make_pliv_CHS2015, make_pliv_multiway_cluster_CKMS2021, \
-    make_did_SZ2020, make_confounded_irm_data, make_confounded_plr_data, make_heterogeneous_data, make_ssm_data
+    make_did_SZ2020, make_confounded_irm_data, make_confounded_plr_data, make_heterogeneous_data, make_ssm_data, \
+    make_irm_data_discrete_treatments
 
 msg_inv_return_type = 'Invalid return_type.'
 
@@ -288,3 +289,34 @@ def test_make_ssm_data_return_types():
     assert isinstance(s, np.ndarray)
     with pytest.raises(ValueError, match=msg_inv_return_type):
         _ = make_ssm_data(n_obs=100, return_type='matrix')
+
+
+@pytest.fixture(scope='function',
+                params=[3, 5])
+def n_levels(request):
+    return request.param
+
+
+def test_make_data_discrete_treatments(n_levels):
+    np.random.seed(3141)
+    n = 100
+    data_apo = make_irm_data_discrete_treatments(n_obs=n, n_levels=3)
+    assert isinstance(data_apo, dict)
+    assert isinstance(data_apo['y'], np.ndarray)
+    assert isinstance(data_apo['d'], np.ndarray)
+    assert isinstance(data_apo['x'], np.ndarray)
+    assert isinstance(data_apo['oracle_values'], dict)
+
+    assert isinstance(data_apo['oracle_values']['cont_d'], np.ndarray)
+    assert isinstance(data_apo['oracle_values']['level_bounds'], np.ndarray)
+    assert isinstance(data_apo['oracle_values']['potential_level'], np.ndarray)
+    assert isinstance(data_apo['oracle_values']['ite'], np.ndarray)
+    assert isinstance(data_apo['oracle_values']['y0'], np.ndarray)
+
+    msg = 'n_levels must be at least 2.'
+    with pytest.raises(ValueError, match=msg):
+        _ = make_irm_data_discrete_treatments(n_obs=n, n_levels=1)
+
+    msg = 'n_levels must be an integer.'
+    with pytest.raises(ValueError, match=msg):
+        _ = make_irm_data_discrete_treatments(n_obs=n, n_levels=1.1)
diff --git a/doubleml/tests/test_exceptions.py b/doubleml/tests/test_exceptions.py
index 215c0a08..cacd1edf 100644
--- a/doubleml/tests/test_exceptions.py
+++ b/doubleml/tests/test_exceptions.py
@@ -8,7 +8,8 @@
     DoubleMLDIDCS, DoubleMLBLP
 from doubleml.datasets import make_plr_CCDDHNR2018, make_irm_data, make_pliv_CHS2015, make_iivm_data, \
     make_pliv_multiway_cluster_CKMS2021, make_did_SZ2020
-from doubleml.double_ml_data import DoubleMLBaseData
+
+from ._utils import DummyDataClass
 
 from sklearn.linear_model import Lasso, LogisticRegression
 from sklearn.base import BaseEstimator
@@ -38,16 +39,6 @@
 dml_data_iivm_binary_outcome = DoubleMLData.from_arrays(x, y, d, z)
 
 
-class DummyDataClass(DoubleMLBaseData):
-    def __init__(self,
-                 data):
-        DoubleMLBaseData.__init__(self, data)
-
-    @property
-    def n_coefs(self):
-        return 1
-
-
 @pytest.mark.ci
 def test_doubleml_exception_data():
     msg = 'The data must be of DoubleMLData or DoubleMLClusterData type.'
@@ -269,6 +260,14 @@ def test_doubleml_exception_data():
                           Lasso(), LogisticRegression())
 
 
+@pytest.mark.ci
+def test_doubleml_exception_framework():
+    msg = r'Apply fit\(\) before sensitivity_analysis\(\).'
+    with pytest.raises(ValueError, match=msg):
+        dml_obj = DoubleMLPLR(dml_data, ml_l, ml_m)
+        dml_obj.sensitivity_analysis()
+
+
 @pytest.mark.ci
 def test_doubleml_exception_scores():
     # PLR
@@ -950,7 +949,7 @@ def test_doubleml_exception_learner():
            'nor a classifier. Method predict is used for prediction.')
     with pytest.warns(UserWarning, match=msg):
         dml_plr_hidden_classifier = DoubleMLPLR(dml_data_irm, Lasso(), log_reg)
-    msg = (r'For the binary treatment variable d, predictions obtained with the ml_m learner LogisticRegression\(\) '
+    msg = (r'For the binary variable d, predictions obtained with the ml_m learner LogisticRegression\(\) '
            'are also observed to be binary with values 0 and 1. Make sure that for classifiers probabilities and not '
            'labels are predicted.')
     with pytest.raises(ValueError, match=msg):
@@ -966,7 +965,7 @@ def test_doubleml_exception_learner():
     with pytest.warns(UserWarning, match=msg):
         dml_irm_hidden_classifier = DoubleMLIRM(dml_data_irm_binary_outcome,
                                                 log_reg, LogisticRegression())
-    msg = (r'For the binary outcome variable y, predictions obtained with the ml_g learner '
+    msg = (r'For the binary variable y, predictions obtained with the ml_g learner '
            r'LogisticRegressionManipulatedPredict\(\) are also observed to be binary with values 0 and 1. Make sure '
            'that for classifiers probabilities and not labels are predicted.')
     with pytest.raises(ValueError, match=msg):
@@ -980,7 +979,7 @@ def test_doubleml_exception_learner():
     with pytest.warns(UserWarning, match=msg):
         dml_iivm_hidden_classifier = DoubleMLIIVM(dml_data_iivm_binary_outcome,
                                                   log_reg, LogisticRegression(), LogisticRegression())
-    msg = (r'For the binary outcome variable y, predictions obtained with the ml_g learner '
+    msg = (r'For the binary variable y, predictions obtained with the ml_g learner '
            r'LogisticRegressionManipulatedPredict\(\) are also observed to be binary with values 0 and 1. Make sure '
            'that for classifiers probabilities and not labels are predicted.')
     with pytest.raises(ValueError, match=msg):
@@ -1108,6 +1107,13 @@ def test_doubleml_sensitivity_inputs():
         dml_irm.sensitivity_analysis()
 
 
+def test_doubleml_sensitivity_summary():
+    dml_irm = DoubleMLIRM(dml_data_irm, Lasso(), LogisticRegression(), trimming_threshold=0.1)
+    msg = r'Apply sensitivity_analysis\(\) before sensitivity_summary.'
+    with pytest.raises(ValueError, match=msg):
+        _ = dml_irm.sensitivity_summary
+
+
 @pytest.mark.ci
 def test_doubleml_sensitivity_benchmark():
     dml_irm = DoubleMLIRM(dml_data_irm, Lasso(), LogisticRegression(), trimming_threshold=0.1)
diff --git a/doubleml/tests/test_framework_exceptions.py b/doubleml/tests/test_framework_exceptions.py
index 7dc8849b..b80cfac2 100644
--- a/doubleml/tests/test_framework_exceptions.py
+++ b/doubleml/tests/test_framework_exceptions.py
@@ -142,6 +142,33 @@ def test_input_exceptions():
         test_dict['cluster_dict'] = {'cluster_ids': np.ones(shape=(n_obs, n_rep))}
         DoubleMLFramework(test_dict)
 
+    test_dict = copy.deepcopy(doubleml_dict)
+    framework_names = DoubleMLFramework(test_dict)
+
+    msg = "treatment_names must be a list. Got 1 of type <class 'int'>."
+    with pytest.raises(TypeError, match=msg):
+        test_dict = copy.deepcopy(doubleml_dict)
+        test_dict['treatment_names'] = 1
+        DoubleMLFramework(test_dict)
+    with pytest.raises(TypeError, match=msg):
+        framework_names.treatment_names = 1
+
+    msg = r"treatment_names must be a list of strings. At least one element is not a string: \['test', 1\]."
+    with pytest.raises(TypeError, match=msg):
+        test_dict = copy.deepcopy(doubleml_dict)
+        test_dict['treatment_names'] = ['test', 1]
+        DoubleMLFramework(test_dict)
+    with pytest.raises(TypeError, match=msg):
+        framework_names.treatment_names = ['test', 1]
+
+    msg = "The length of treatment_names does not match the number of treatments. Got 2 treatments and 3 treatment names."
+    with pytest.raises(ValueError, match=msg):
+        test_dict = copy.deepcopy(doubleml_dict)
+        test_dict['treatment_names'] = ['test', 'test2', 'test3']
+        DoubleMLFramework(test_dict)
+    with pytest.raises(ValueError, match=msg):
+        framework_names.treatment_names = ['test', 'test2', 'test3']
+
 
 def test_operation_exceptions():
     # addition
diff --git a/doubleml/tests/test_framework_sensitivity.py b/doubleml/tests/test_framework_sensitivity.py
index 5b1be8bd..044d89d2 100644
--- a/doubleml/tests/test_framework_sensitivity.py
+++ b/doubleml/tests/test_framework_sensitivity.py
@@ -43,6 +43,7 @@ def dml_framework_sensitivity_fixture(n_rep, generate_data_simple):
         'dml_obj': dml_irm_obj,
         'dml_obj_2': dml_irm_obj_2,
         'dml_framework_obj': dml_framework_obj,
+        'dml_framework_obj_2': dml_framework_obj_2,
         'dml_framework_obj_add_obj': dml_framework_obj_add_obj,
         'dml_framework_obj_sub_obj': dml_framework_obj_sub_obj,
         'dml_framework_obj_mul_obj': dml_framework_obj_mul_obj,
@@ -59,6 +60,7 @@ def test_dml_framework_sensitivity_shapes(dml_framework_sensitivity_fixture):
     n_obs = dml_framework_sensitivity_fixture['dml_framework_obj'].n_obs
 
     object_list = ['dml_framework_obj',
+                   'dml_framework_obj_2',
                    'dml_framework_obj_add_obj',
                    'dml_framework_obj_sub_obj',
                    'dml_framework_obj_mul_obj']
@@ -81,3 +83,24 @@ def test_dml_framework_sensitivity_shapes(dml_framework_sensitivity_fixture):
     for key in score_keys:
         assert dml_framework_sensitivity_fixture['dml_framework_obj_concat']._sensitivity_elements[key].shape == \
             (n_obs, 2, n_rep)
+
+
+@pytest.mark.ci
+def test_dml_framework_sensitivity_summary(dml_framework_sensitivity_fixture):
+    # summary without sensitivity analysis
+    sensitivity_summary = dml_framework_sensitivity_fixture['dml_framework_obj_2'].sensitivity_summary
+    substring = 'Apply sensitivity_analysis() to generate sensitivity_summary.'
+    assert substring in sensitivity_summary
+
+    # summary with sensitivity analysis
+    sensitivity_summary = dml_framework_sensitivity_fixture['dml_framework_obj'].sensitivity_summary
+    assert isinstance(sensitivity_summary, str)
+    substrings = [
+        '\n------------------ Scenario          ------------------\n',
+        '\n------------------ Bounds with CI    ------------------\n',
+        '\n------------------ Robustness Values ------------------\n',
+        'Significance Level: level=',
+        'Sensitivity parameters: cf_y='
+    ]
+    for substring in substrings:
+        assert substring in sensitivity_summary
diff --git a/doubleml/tests/test_model_defaults.py b/doubleml/tests/test_model_defaults.py
index 5da2fbca..8d7234d6 100644
--- a/doubleml/tests/test_model_defaults.py
+++ b/doubleml/tests/test_model_defaults.py
@@ -1,8 +1,7 @@
 import pytest
 import numpy as np
 
-from doubleml import DoubleMLPLR, DoubleMLIRM, DoubleMLIIVM, DoubleMLPLIV, DoubleMLCVAR, DoubleMLPQ, \
-    DoubleMLLPQ, DoubleMLQTE, DoubleMLDID, DoubleMLDIDCS, DoubleMLSSM
+import doubleml as dml
 from doubleml.datasets import make_plr_CCDDHNR2018, make_irm_data, make_pliv_CHS2015, make_iivm_data, make_did_SZ2020, \
     make_ssm_data
 
@@ -19,47 +18,34 @@
 dml_data_ssm = make_ssm_data(n_obs=2000, mar=True)
 
 # linear models
-dml_plr = DoubleMLPLR(dml_data_plr, Lasso(), Lasso())
-dml_pliv = DoubleMLPLIV(dml_data_pliv, Lasso(), Lasso(), Lasso())
-dml_irm = DoubleMLIRM(dml_data_irm, Lasso(), LogisticRegression())
-dml_iivm = DoubleMLIIVM(dml_data_iivm, Lasso(), LogisticRegression(), LogisticRegression())
-dml_cvar = DoubleMLCVAR(dml_data_irm, ml_g=RandomForestRegressor(), ml_m=RandomForestClassifier())
-dml_did = DoubleMLDID(dml_data_did, Lasso(), LogisticRegression())
-dml_did_cs = DoubleMLDIDCS(dml_data_did_cs, Lasso(), LogisticRegression())
-dml_ssm = DoubleMLSSM(dml_data_ssm, Lasso(), LogisticRegression(), LogisticRegression())
-
-dml_plr.fit()
-dml_pliv.fit()
-dml_irm.fit()
-dml_iivm.fit()
-dml_cvar.fit()
-dml_did.fit()
-dml_did_cs.fit()
-dml_ssm.fit()
-
-dml_plr.bootstrap()
-dml_pliv.bootstrap()
-dml_irm.bootstrap()
-dml_iivm.bootstrap()
-dml_cvar.bootstrap()
-dml_did.bootstrap()
-dml_did_cs.bootstrap()
-dml_ssm.bootstrap()
+dml_plr = dml.DoubleMLPLR(dml_data_plr, Lasso(), Lasso())
+dml_pliv = dml.DoubleMLPLIV(dml_data_pliv, Lasso(), Lasso(), Lasso())
+dml_irm = dml.DoubleMLIRM(dml_data_irm, Lasso(), LogisticRegression())
+dml_iivm = dml.DoubleMLIIVM(dml_data_iivm, Lasso(), LogisticRegression(), LogisticRegression())
+dml_cvar = dml.DoubleMLCVAR(dml_data_irm, ml_g=RandomForestRegressor(), ml_m=RandomForestClassifier())
+dml_did = dml.DoubleMLDID(dml_data_did, Lasso(), LogisticRegression())
+dml_did_cs = dml.DoubleMLDIDCS(dml_data_did_cs, Lasso(), LogisticRegression())
+dml_ssm = dml.DoubleMLSSM(dml_data_ssm, Lasso(), LogisticRegression(), LogisticRegression())
+dml_apo = dml.DoubleMLAPO(dml_data_irm, Lasso(), LogisticRegression(), treatment_level=0)
+dml_apos = dml.DoubleMLAPOS(dml_data_irm, Lasso(), LogisticRegression(), treatment_levels=[0, 1])
 
 # nonlinear models
-dml_pq = DoubleMLPQ(dml_data_irm, ml_g=LogisticRegression(), ml_m=LogisticRegression())
-dml_lpq = DoubleMLLPQ(dml_data_iivm, ml_g=RandomForestClassifier(), ml_m=RandomForestClassifier())
-dml_qte = DoubleMLQTE(dml_data_irm, ml_g=RandomForestClassifier(), ml_m=RandomForestClassifier())
+dml_pq = dml.DoubleMLPQ(dml_data_irm, ml_g=LogisticRegression(), ml_m=LogisticRegression())
+dml_lpq = dml.DoubleMLLPQ(dml_data_iivm, ml_g=RandomForestClassifier(), ml_m=RandomForestClassifier())
+dml_qte = dml.DoubleMLQTE(dml_data_irm, ml_g=RandomForestClassifier(), ml_m=RandomForestClassifier())
 
-dml_pq.fit()
-dml_lpq.fit()
-dml_qte.fit()
 
-dml_pq.bootstrap()
-dml_lpq.bootstrap()
-dml_qte.bootstrap()
+def _assert_is_none(dml_obj):
+    assert dml_obj.n_rep_boot is None
+    assert dml_obj.boot_method is None
+    assert dml_obj.framework is None
+    assert dml_obj.sensitivity_params is None
+    assert dml_obj.boot_t_stat is None
 
-policy_tree = dml_irm.policy_tree(features=dml_data_irm.data.drop(columns=["y", "d"]))
+
+def _fit_bootstrap(dml_obj):
+    dml_obj.fit()
+    dml_obj.bootstrap()
 
 
 def _assert_resampling_default_settings(dml_obj):
@@ -84,12 +70,16 @@ def _assert_resampling_default_settings(dml_obj):
 
 @pytest.mark.ci
 def test_plr_defaults():
+    _assert_is_none(dml_plr)
+    _fit_bootstrap(dml_plr)
     _assert_resampling_default_settings(dml_plr)
     assert dml_plr.score == 'partialling out'
 
 
 @pytest.mark.ci
 def test_pliv_defaults():
+    _assert_is_none(dml_pliv)
+    _fit_bootstrap(dml_pliv)
     _assert_resampling_default_settings(dml_pliv)
     assert dml_pliv.score == 'partialling out'
     assert dml_pliv.partialX
@@ -98,6 +88,8 @@ def test_pliv_defaults():
 
 @pytest.mark.ci
 def test_irm_defaults():
+    _assert_is_none(dml_irm)
+    _fit_bootstrap(dml_irm)
     _assert_resampling_default_settings(dml_irm)
     assert dml_irm.score == 'ATE'
     assert dml_irm.trimming_rule == 'truncate'
@@ -109,6 +101,8 @@ def test_irm_defaults():
 
 @pytest.mark.ci
 def test_iivm_defaults():
+    _assert_is_none(dml_iivm)
+    _fit_bootstrap(dml_iivm)
     _assert_resampling_default_settings(dml_iivm)
     assert dml_iivm.score == 'LATE'
     assert dml_iivm.subgroups == {'always_takers': True, 'never_takers': True}
@@ -119,6 +113,8 @@ def test_iivm_defaults():
 
 @pytest.mark.ci
 def test_cvar_defaults():
+    _assert_is_none(dml_cvar)
+    _fit_bootstrap(dml_cvar)
     _assert_resampling_default_settings(dml_cvar)
     assert dml_cvar.quantile == 0.5
     assert dml_cvar.treatment == 1
@@ -129,6 +125,8 @@ def test_cvar_defaults():
 
 @pytest.mark.ci
 def test_pq_defaults():
+    _assert_is_none(dml_pq)
+    _fit_bootstrap(dml_pq)
     _assert_resampling_default_settings(dml_pq)
     assert dml_pq.quantile == 0.5
     assert dml_pq.treatment == 1
@@ -140,6 +138,8 @@ def test_pq_defaults():
 
 @pytest.mark.ci
 def test_lpq_defaults():
+    _assert_is_none(dml_lpq)
+    _fit_bootstrap(dml_lpq)
     _assert_resampling_default_settings(dml_lpq)
     assert dml_lpq.quantile == 0.5
     assert dml_lpq.treatment == 1
@@ -151,6 +151,11 @@ def test_lpq_defaults():
 
 @pytest.mark.ci
 def test_qte_defaults():
+    assert dml_qte.n_rep_boot is None
+    assert dml_qte.boot_method is None
+    assert dml_qte.framework is None
+    assert dml_qte.boot_t_stat is None
+    _fit_bootstrap(dml_qte)
     # not fix since its a differen object added in future versions _assert_resampling_default_settings(dml_qte)
     assert dml_qte.quantiles == 0.5
     assert dml_qte.score == 'PQ'
@@ -161,6 +166,8 @@ def test_qte_defaults():
 
 @pytest.mark.ci
 def test_did_defaults():
+    _assert_is_none(dml_did)
+    _fit_bootstrap(dml_did)
     _assert_resampling_default_settings(dml_did)
     assert dml_did.score == 'observational'
     assert dml_did.in_sample_normalization
@@ -170,6 +177,8 @@ def test_did_defaults():
 
 @pytest.mark.ci
 def test_did_cs_defaults():
+    _assert_is_none(dml_did_cs)
+    _fit_bootstrap(dml_did_cs)
     _assert_resampling_default_settings(dml_did_cs)
     assert dml_did.score == 'observational'
     assert dml_did_cs.in_sample_normalization
@@ -179,6 +188,8 @@ def test_did_cs_defaults():
 
 @pytest.mark.ci
 def test_ssm_defaults():
+    _assert_is_none(dml_ssm)
+    _fit_bootstrap(dml_ssm)
     _assert_resampling_default_settings(dml_ssm)
     assert dml_ssm.score == 'missing-at-random'
     assert dml_ssm.trimming_rule == 'truncate'
@@ -186,6 +197,33 @@ def test_ssm_defaults():
     assert not dml_ssm.normalize_ipw
 
 
+@pytest.mark.ci
+def test_apo_defaults():
+    _assert_is_none(dml_apo)
+    _fit_bootstrap(dml_apo)
+    _assert_resampling_default_settings(dml_apo)
+    assert dml_apo.score == 'APO'
+    assert dml_apo.trimming_rule == 'truncate'
+    assert dml_apo.trimming_threshold == 1e-2
+    assert not dml_apo.normalize_ipw
+    assert set(dml_apo.weights.keys()) == set(['weights'])
+    assert np.array_equal(dml_apo.weights['weights'], np.ones((dml_apo._dml_data.n_obs,)))
+
+
+@pytest.mark.ci
+def test_apos_defaults():
+    assert dml_apos.n_rep_boot is None
+    assert dml_apos.boot_method is None
+    assert dml_apos.framework is None
+    assert dml_apos.boot_t_stat is None
+    _fit_bootstrap(dml_qte)
+    assert dml_apos.score == 'APO'
+    assert dml_apos.trimming_rule == 'truncate'
+    assert dml_apos.trimming_threshold == 1e-2
+    assert not dml_apos.normalize_ipw
+    assert np.array_equal(dml_apos.weights, np.ones((dml_apos._dml_data.n_obs,)))
+
+
 @pytest.mark.ci
 def test_sensitivity_defaults():
     input_dict = {'cf_y': 0.03,
@@ -200,6 +238,9 @@ def test_sensitivity_defaults():
 
 @pytest.mark.ci
 def test_policytree_defaults():
+    dml_irm = dml.DoubleMLIRM(dml_data_irm, Lasso(), LogisticRegression())
+    dml_irm.fit()
+    policy_tree = dml_irm.policy_tree(features=dml_data_irm.data.drop(columns=["y", "d"]))
     assert policy_tree.policy_tree.max_depth == 2
     assert policy_tree.policy_tree.min_samples_leaf == 8
     assert policy_tree.policy_tree.ccp_alpha == 0.01
diff --git a/doubleml/tests/test_return_types.py b/doubleml/tests/test_return_types.py
index d76f2d14..a9014d08 100644
--- a/doubleml/tests/test_return_types.py
+++ b/doubleml/tests/test_return_types.py
@@ -18,6 +18,7 @@
     DoubleMLPolicyTree,
     DoubleMLFramework,
     DoubleMLSSM,
+    DoubleMLAPO
 )
 from doubleml.datasets import (
     make_plr_CCDDHNR2018,
@@ -61,6 +62,7 @@
 dml_did_cs = DoubleMLDIDCS(dml_data_did_cs, Lasso(), LogisticRegression())
 dml_did_cs_binary_outcome = DoubleMLDIDCS(dml_data_did_cs_binary_outcome, LogisticRegression(), LogisticRegression())
 dml_ssm = DoubleMLSSM(dml_data_ssm, ml_g=Lasso(), ml_m=LogisticRegression(), ml_pi=LogisticRegression())
+dml_apo = DoubleMLAPO(dml_data_irm, Lasso(), LogisticRegression(), treatment_level=0)
 
 
 @pytest.mark.ci
@@ -77,7 +79,8 @@
                           (dml_did_binary_outcome, DoubleMLDID),
                           (dml_did_cs, DoubleMLDIDCS),
                           (dml_did_cs_binary_outcome, DoubleMLDIDCS),
-                          (dml_ssm, DoubleMLSSM)])
+                          (dml_ssm, DoubleMLSSM),
+                          (dml_apo, DoubleMLAPO)])
 def test_return_types(dml_obj, cls):
     # ToDo: A second test case with multiple treatment variables would be helpful
     assert isinstance(dml_obj.__str__(), str)
@@ -168,11 +171,16 @@ def test_return_types(dml_obj, cls):
 ssm_obj.fit()
 ssm_obj.bootstrap(n_rep_boot=n_rep_boot)
 
+apo_obj = DoubleMLAPO(dml_data_irm, Lasso(), LogisticRegression(), treatment_level=0,
+                      n_rep=n_rep, n_folds=n_folds)
+apo_obj.fit()
+apo_obj.bootstrap(n_rep_boot=n_rep_boot)
+
 
 @pytest.mark.ci
 @pytest.mark.parametrize('dml_obj',
                          [plr_obj, pliv_obj,  irm_obj,  iivm_obj, cvar_obj, pq_obj, lpq_obj,
-                          did_obj, did_cs_obj])
+                          did_obj, did_cs_obj, ssm_obj, apo_obj])
 def test_property_types_and_shapes(dml_obj):
     # not checked: learner, learner_names, params, params_names, score
     # already checked: summary
@@ -300,6 +308,10 @@ def test_stored_predictions():
     assert ssm_obj.predictions['ml_m'].shape == (n_obs, n_rep, n_treat)
     assert ssm_obj.predictions['ml_pi'].shape == (n_obs, n_rep, n_treat)
 
+    assert apo_obj.predictions['ml_g0'].shape == (n_obs, n_rep, n_treat)
+    assert apo_obj.predictions['ml_g1'].shape == (n_obs, n_rep, n_treat)
+    assert apo_obj.predictions['ml_m'].shape == (n_obs, n_rep, n_treat)
+
 
 @pytest.mark.ci
 def test_stored_nuisance_targets():
@@ -347,6 +359,10 @@ def test_stored_nuisance_targets():
     assert ssm_obj.nuisance_targets['ml_m'].shape == (n_obs, n_rep, n_treat)
     assert ssm_obj.nuisance_targets['ml_pi'].shape == (n_obs, n_rep, n_treat)
 
+    assert apo_obj.nuisance_targets['ml_g0'].shape == (n_obs, n_rep, n_treat)
+    assert apo_obj.nuisance_targets['ml_g1'].shape == (n_obs, n_rep, n_treat)
+    assert apo_obj.nuisance_targets['ml_m'].shape == (n_obs, n_rep, n_treat)
+
 
 @pytest.mark.ci
 def test_nuisance_loss():
@@ -394,100 +410,55 @@ def test_nuisance_loss():
     assert ssm_obj.nuisance_loss['ml_m'].shape == (n_rep, n_treat)
     assert ssm_obj.nuisance_loss['ml_pi'].shape == (n_rep, n_treat)
 
+    assert apo_obj.nuisance_loss['ml_g0'].shape == (n_rep, n_treat)
+    assert apo_obj.nuisance_loss['ml_g1'].shape == (n_rep, n_treat)
+    assert apo_obj.nuisance_loss['ml_m'].shape == (n_rep, n_treat)
 
-@pytest.mark.ci
-def test_sensitivity():
 
-    var_keys = ['sigma2', 'nu2']
-    score_keys = ['psi_sigma2', 'psi_nu2', 'riesz_rep']
-    benchmarks = {'cf_y': [0.1, 0.2], 'cf_d': [0.15, 0.2], 'name': ["test1", "test2"]}
+def _test_sensitivity_return_types(dml_obj, n_rep, n_treat, benchmarking_set):
+    assert isinstance(dml_obj.sensitivity_elements, dict)
+    for key in ['sigma2', 'nu2']:
+        assert isinstance(dml_obj.sensitivity_elements[key], np.ndarray)
+        assert dml_obj.sensitivity_elements[key].shape == (1, n_rep, n_treat)
+    for key in ['psi_sigma2', 'psi_nu2', 'riesz_rep']:
+        assert isinstance(dml_obj.sensitivity_elements[key], np.ndarray)
+        assert dml_obj.sensitivity_elements[key].shape == (n_obs, n_rep, n_treat)
 
-    # PLR
-    assert isinstance(plr_obj.sensitivity_elements, dict)
-    for key in var_keys:
-        assert isinstance(plr_obj.sensitivity_elements[key], np.ndarray)
-        assert plr_obj.sensitivity_elements[key].shape == (1, n_rep, n_treat)
-    for key in score_keys:
-        assert isinstance(plr_obj.sensitivity_elements[key], np.ndarray)
-        assert plr_obj.sensitivity_elements[key].shape == (n_obs, n_rep, n_treat)
-
-    assert isinstance(plr_obj.sensitivity_summary, str)
-    plr_obj.sensitivity_analysis()
-    assert isinstance(plr_obj.sensitivity_summary, str)
-    assert isinstance(plr_obj.sensitivity_plot(), plotly.graph_objs._figure.Figure)
-    assert isinstance(plr_obj.sensitivity_plot(value='ci', benchmarks=benchmarks), plotly.graph_objs._figure.Figure)
-    assert isinstance(plr_obj.framework._calc_sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95), dict)
-    assert isinstance(
-        plr_obj.framework._calc_robustness_value(null_hypothesis=0.0, level=0.95, rho=1.0, idx_treatment=0),
-        tuple)
-    plr_benchmark = plr_obj.sensitivity_benchmark(benchmarking_set=["X1"])
-    assert isinstance(plr_benchmark, pd.DataFrame)
+    assert isinstance(dml_obj.sensitivity_summary, str)
+    dml_obj.sensitivity_analysis()
+    assert isinstance(dml_obj.sensitivity_summary, str)
+    assert isinstance(dml_obj.sensitivity_plot(), plotly.graph_objs._figure.Figure)
+    benchmarks = {'cf_y': [0.1, 0.2], 'cf_d': [0.15, 0.2], 'name': ["test1", "test2"]}
+    assert isinstance(dml_obj.sensitivity_plot(value='ci', benchmarks=benchmarks), plotly.graph_objs._figure.Figure)
 
-    # DID
-    assert isinstance(irm_obj.sensitivity_elements, dict)
-    for key in var_keys:
-        assert isinstance(irm_obj.sensitivity_elements[key], np.ndarray)
-        assert irm_obj.sensitivity_elements[key].shape == (1, n_rep, n_treat)
-    for key in score_keys:
-        assert isinstance(irm_obj.sensitivity_elements[key], np.ndarray)
-        assert irm_obj.sensitivity_elements[key].shape == (n_obs, n_rep, n_treat)
-
-    assert isinstance(irm_obj.sensitivity_summary, str)
-    irm_obj.sensitivity_analysis()
-    assert isinstance(irm_obj.sensitivity_summary, str)
-    assert isinstance(irm_obj.sensitivity_plot(), plotly.graph_objs._figure.Figure)
-    assert isinstance(irm_obj.sensitivity_plot(value='ci', benchmarks=benchmarks), plotly.graph_objs._figure.Figure)
-    assert isinstance(irm_obj.framework._calc_sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95), dict)
+    assert isinstance(dml_obj.framework._calc_sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95), dict)
     assert isinstance(
-        irm_obj.framework._calc_robustness_value(null_hypothesis=0.0, level=0.95, rho=1.0, idx_treatment=0),
+        dml_obj.framework._calc_robustness_value(null_hypothesis=0.0, level=0.95, rho=1.0, idx_treatment=0),
         tuple
     )
-    irm_benchmark = irm_obj.sensitivity_benchmark(benchmarking_set=["X1"])
-    assert isinstance(irm_benchmark, pd.DataFrame)
+    benchmark = dml_obj.sensitivity_benchmark(benchmarking_set=benchmarking_set)
+    assert isinstance(benchmark, pd.DataFrame)
+
+    return
+
+
+@pytest.mark.ci
+def test_sensitivity():
+
+    # PLR
+    _test_sensitivity_return_types(plr_obj, n_rep, n_treat, benchmarking_set=["X1"])
+
+    # IRM
+    _test_sensitivity_return_types(irm_obj, n_rep, n_treat, benchmarking_set=["X1"])
 
     # DID
-    assert isinstance(did_obj.sensitivity_elements, dict)
-    for key in var_keys:
-        assert isinstance(did_obj.sensitivity_elements[key], np.ndarray)
-        assert did_obj.sensitivity_elements[key].shape == (1, n_rep, n_treat)
-    for key in score_keys:
-        assert isinstance(did_obj.sensitivity_elements[key], np.ndarray)
-        assert did_obj.sensitivity_elements[key].shape == (n_obs, n_rep, n_treat)
-
-    assert isinstance(did_obj.sensitivity_summary, str)
-    did_obj.sensitivity_analysis()
-    assert isinstance(did_obj.sensitivity_summary, str)
-    assert isinstance(did_obj.sensitivity_plot(), plotly.graph_objs._figure.Figure)
-    assert isinstance(did_obj.sensitivity_plot(value='ci', benchmarks=benchmarks), plotly.graph_objs._figure.Figure)
-    assert isinstance(did_obj.framework._calc_sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95), dict)
-    assert isinstance(
-        did_obj.framework._calc_robustness_value(null_hypothesis=0.0, level=0.95, rho=1.0, idx_treatment=0),
-        tuple
-    )
-    did_benchmark = did_obj.sensitivity_benchmark(benchmarking_set=['Z1'])
-    assert isinstance(did_benchmark, pd.DataFrame)
+    _test_sensitivity_return_types(did_obj, n_rep, n_treat, benchmarking_set=["Z1"])
 
     # DIDCS
-    assert isinstance(did_cs_obj.sensitivity_elements, dict)
-    for key in var_keys:
-        assert isinstance(did_cs_obj.sensitivity_elements[key], np.ndarray)
-        assert did_cs_obj.sensitivity_elements[key].shape == (1, n_rep, n_treat)
-    for key in score_keys:
-        assert isinstance(did_cs_obj.sensitivity_elements[key], np.ndarray)
-        assert did_cs_obj.sensitivity_elements[key].shape == (n_obs, n_rep, n_treat)
-
-    assert isinstance(did_cs_obj.sensitivity_summary, str)
-    did_cs_obj.sensitivity_analysis()
-    assert isinstance(did_cs_obj.sensitivity_summary, str)
-    assert isinstance(did_cs_obj.sensitivity_plot(), plotly.graph_objs._figure.Figure)
-    assert isinstance(did_cs_obj.sensitivity_plot(value='ci', benchmarks=benchmarks), plotly.graph_objs._figure.Figure)
-    assert isinstance(did_cs_obj.framework._calc_sensitivity_analysis(cf_y=0.03, cf_d=0.03, rho=1.0, level=0.95), dict)
-    assert isinstance(
-        did_cs_obj.framework._calc_robustness_value(null_hypothesis=0.0, level=0.95, rho=1.0, idx_treatment=0),
-        tuple
-    )
-    did_cs_benchmark = did_cs_obj.sensitivity_benchmark(benchmarking_set=['Z1'])
-    assert isinstance(did_cs_benchmark, pd.DataFrame)
+    _test_sensitivity_return_types(did_cs_obj, n_rep, n_treat, benchmarking_set=["Z1"])
+
+    # APO
+    _test_sensitivity_return_types(apo_obj, n_rep, n_treat, benchmarking_set=["X1"])
 
 
 @pytest.mark.ci
diff --git a/doubleml/utils/tests/test_exceptions_gain_statistics.py b/doubleml/utils/tests/test_exceptions_gain_statistics.py
index 805a84ed..734185eb 100644
--- a/doubleml/utils/tests/test_exceptions_gain_statistics.py
+++ b/doubleml/utils/tests/test_exceptions_gain_statistics.py
@@ -4,9 +4,14 @@
 from doubleml.utils.gain_statistics import gain_statistics
 
 
+class test_framework():
+    def __init__(self, sensitivity_elements):
+        self.sensitivity_elements = sensitivity_elements
+
+
 class test_dml_class():
     def __init__(self, sensitivity_elements, all_coef):
-        self.sensitivity_elements = sensitivity_elements
+        self.framework = test_framework(sensitivity_elements)
         self.all_coef = all_coef
 
 
@@ -22,19 +27,20 @@ def test_doubleml_exception_data():
             'sigma2': np.random.normal(size=(n_obs, n_rep, n_coef)),
             'nu2': np.random.normal(size=(n_obs, n_rep, n_coef))
         },
-        all_coef=np.random.normal(size=(n_coef, n_rep))
+        all_coef=np.random.normal(size=(n_rep, n_coef))
     )
 
     # incorrect types
     dml_incorrect = test_dml_class(
             sensitivity_elements=np.random.normal(size=(n_obs, n_rep, n_coef)),
-            all_coef=np.random.normal(size=(n_coef, n_rep))
+            all_coef=np.random.normal(size=(n_rep, n_coef))
         )
-    msg = r"dml_long does not contain the necessary sensitivity elements\. Expected dict for dml_long\.sensitivity_elements\."
+    msg = r"dml_long does not contain the necessary sensitivity elements\. "
+    msg += r"Expected dict for dml_long\.framework\.sensitivity_elements\."
     with pytest.raises(TypeError, match=msg):
         _ = gain_statistics(dml_incorrect, dml_correct)
     msg = r"dml_short does not contain the necessary sensitivity elements\. "
-    msg += r"Expected dict for dml_short\.sensitivity_elements\."
+    msg += r"Expected dict for dml_short\.framework\.sensitivity_elements\."
     with pytest.raises(TypeError, match=msg):
         _ = gain_statistics(dml_correct, dml_incorrect)
 
@@ -43,7 +49,7 @@ def test_doubleml_exception_data():
             sensitivity_elements={
                 'sigma2': np.random.normal(size=(n_obs, n_rep, n_coef)),
             },
-            all_coef=np.random.normal(size=(n_coef, n_rep))
+            all_coef=np.random.normal(size=(n_rep, n_coef))
         )
     msg = r"dml_long does not contain the necessary sensitivity elements\. Required keys are: \['sigma2', 'nu2'\]"
     with pytest.raises(ValueError, match=msg):
@@ -58,7 +64,7 @@ def test_doubleml_exception_data():
                 'sigma2': {},
                 'nu2': np.random.normal(size=(n_obs, n_rep, n_coef))
             },
-            all_coef=np.random.normal(size=(n_coef, n_rep))
+            all_coef=np.random.normal(size=(n_rep, n_coef))
         )
     msg = r"dml_long does not contain the necessary sensitivity elements\. Expected numpy\.ndarray for key sigma2\."
     with pytest.raises(TypeError, match=msg):
@@ -72,7 +78,7 @@ def test_doubleml_exception_data():
             'sigma2': np.random.normal(size=(n_obs, n_rep, n_coef)),
             'nu2': {}
         },
-        all_coef=np.random.normal(size=(n_coef, n_rep))
+        all_coef=np.random.normal(size=(n_rep, n_coef))
     )
     msg = r"dml_long does not contain the necessary sensitivity elements\. Expected numpy\.ndarray for key nu2\."
     with pytest.raises(TypeError, match=msg):
@@ -87,7 +93,7 @@ def test_doubleml_exception_data():
                 'sigma2': np.random.normal(size=(n_obs + 1, n_rep, n_coef)),
                 'nu2': np.random.normal(size=(n_obs, n_rep, n_coef))
             },
-            all_coef=np.random.normal(size=(n_coef, n_rep))
+            all_coef=np.random.normal(size=(n_rep, n_coef))
         )
     msg = (r"dml_long does not contain the necessary sensitivity elements\. "
            r"Expected 3 dimensions of shape \(1, n_coef, n_rep\) for key sigma2\.")
@@ -103,7 +109,7 @@ def test_doubleml_exception_data():
                 'sigma2': np.random.normal(size=(n_obs, n_rep, n_coef)),
                 'nu2': np.random.normal(size=(n_obs + 1, n_rep, n_coef))
             },
-            all_coef=np.random.normal(size=(n_coef, n_rep))
+            all_coef=np.random.normal(size=(n_rep, n_coef))
         )
     msg = (r"dml_long does not contain the necessary sensitivity elements\. "
            r"Expected 3 dimensions of shape \(1, n_coef, n_rep\) for key nu2\.")
@@ -120,7 +126,7 @@ def test_doubleml_exception_data():
                 'sigma2': np.random.normal(size=(n_obs, n_rep + 1, n_coef)),
                 'nu2': np.random.normal(size=(n_obs, n_rep, n_coef))
             },
-            all_coef=np.random.normal(size=(n_coef, n_rep))
+            all_coef=np.random.normal(size=(n_rep, n_coef))
         )
     msg = r"dml_long and dml_short do not contain the same shape of sensitivity elements\. "
     msg += r"Shapes of sigma2 are: \(1, 4, 5\) and \(1, 3, 5\)"
@@ -136,7 +142,7 @@ def test_doubleml_exception_data():
                 'sigma2': np.random.normal(size=(n_obs, n_rep, n_coef)),
                 'nu2': np.random.normal(size=(n_obs, n_rep + 1, n_coef))
             },
-            all_coef=np.random.normal(size=(n_coef, n_rep))
+            all_coef=np.random.normal(size=(n_rep, n_coef))
         )
     msg = r"dml_long and dml_short do not contain the same shape of sensitivity elements\. "
     msg += r"Shapes of nu2 are: \(1, 4, 5\) and \(1, 3, 5\)"
@@ -168,11 +174,11 @@ def test_doubleml_exception_data():
                 'sigma2': np.random.normal(size=(n_obs, n_rep, n_coef)),
                 'nu2': np.random.normal(size=(n_obs, n_rep, n_coef))
             },
-            all_coef=np.random.normal(size=(n_coef + 1, n_rep))
+            all_coef=np.random.normal(size=(n_rep, n_coef + 1))
         )
-    msg = r"dml_long\.all_coef does not contain the necessary coefficients\. Expected shape: \(5, 3\)"
+    msg = r"dml_long\.all_coef does not contain the necessary coefficients\. Expected shape: \(3, 5\)"
     with pytest.raises(ValueError, match=msg):
         _ = gain_statistics(dml_incorrect, dml_correct)
-    msg = r"dml_short\.all_coef does not contain the necessary coefficients\. Expected shape: \(5, 3\)"
+    msg = r"dml_short\.all_coef does not contain the necessary coefficients\. Expected shape: \(3, 5\)"
     with pytest.raises(ValueError, match=msg):
         _ = gain_statistics(dml_correct, dml_incorrect)