Repository: sbi-dev/sbi

Issue:
Confusing name `force_first_round_loss`
We recently introduced `force_first_round_loss`. Without setting `force_first_round_loss=True`, one can not train the `SNPE` object multiple times without proposal.

However, I just realized that, if `force_first_round_loss=True` but a `proposal` is passed, we still use the atomic loss (or whichever multi-round loss is used...). I think we should rename `force_first_round_loss` to `allow_first_round_loss` to avoid confusion.



Discussion:

Resolving patch:
diff --git a/sbi/inference/abc/mcabc.py b/sbi/inference/abc/mcabc.py
--- a/sbi/inference/abc/mcabc.py
+++ b/sbi/inference/abc/mcabc.py
@@ -177,8 +177,9 @@ def __call__(
             kde_dist = get_kde(final_theta, **kde_kwargs)
 
             if return_summary:
-                return kde_dist, dict(
-                    theta=final_theta, distances=distances_accepted, x=x_accepted
+                return (
+                    kde_dist,
+                    dict(theta=final_theta, distances=distances_accepted, x=x_accepted),
                 )
             else:
                 return kde_dist
diff --git a/sbi/inference/posteriors/mcmc_posterior.py b/sbi/inference/posteriors/mcmc_posterior.py
--- a/sbi/inference/posteriors/mcmc_posterior.py
+++ b/sbi/inference/posteriors/mcmc_posterior.py
@@ -285,7 +285,8 @@ def _build_mcmc_init_fn(
             init_strategy: Specifies the initialization method. Either of
                 [`proposal`|`sir`|`latest_sample`].
             kwargs: Passed on to init function. This way, init specific keywords can
-                be set through `mcmc_parameters`. Unused arguments should be absorbed.
+                be set through `mcmc_parameters`. Unused arguments will be absorbed by
+                the intitialization method.
 
         Returns: Initialization function.
         """
@@ -311,6 +312,7 @@ def _get_initial_params(
         num_chains: int,
         num_workers: int,
         show_progress_bars: bool,
+        **kwargs,
     ) -> Tensor:
         """Return initial parameters for MCMC obtained with given init strategy.
 
@@ -322,6 +324,7 @@ def _get_initial_params(
             num_chains: number of MCMC chains, generates initial params for each
             num_workers: number of CPU cores for parallization
             show_progress_bars: whether to show progress bars for SIR init
+            kwargs: Passed on to `_build_mcmc_init_fn`.
 
         Returns:
             Tensor: initial parameters, one for each chain
@@ -332,6 +335,7 @@ def _get_initial_params(
             self.potential_fn,
             transform=self.theta_transform,
             init_strategy=init_strategy,  # type: ignore
+            **kwargs,
         )
 
         # Parallelize inits for SIR only.
diff --git a/sbi/inference/snpe/snpe_base.py b/sbi/inference/snpe/snpe_base.py
--- a/sbi/inference/snpe/snpe_base.py
+++ b/sbi/inference/snpe/snpe_base.py
@@ -302,7 +302,12 @@ def train(
                 )
 
                 train_losses = self._loss(
-                    theta_batch, x_batch, masks_batch, proposal, calibration_kernel
+                    theta_batch,
+                    x_batch,
+                    masks_batch,
+                    proposal,
+                    calibration_kernel,
+                    force_first_round_loss=force_first_round_loss,
                 )
                 train_loss = torch.mean(train_losses)
                 train_log_probs_sum -= train_losses.sum().item()
@@ -339,6 +344,7 @@ def train(
                         masks_batch,
                         proposal,
                         calibration_kernel,
+                        force_first_round_loss=force_first_round_loss,
                     )
                     val_log_prob_sum -= val_losses.sum().item()
 
@@ -499,6 +505,7 @@ def _loss(
         masks: Tensor,
         proposal: Optional[Any],
         calibration_kernel: Callable,
+        force_first_round_loss: bool = False,
     ) -> Tensor:
         """Return loss with proposal correction (`round_>0`) or without it (`round_=0`).
 
@@ -507,8 +514,11 @@ def _loss(
 
         Returns:
             Calibration kernel-weighted negative log prob.
+            force_first_round_loss: If `True`, train with maximum likelihood,
+                i.e., potentially ignoring the correction for using a proposal
+                distribution different from the prior.
         """
-        if self._round == 0:
+        if self._round == 0 or force_first_round_loss:
             # Use posterior log prob (without proposal correction) for first round.
             log_prob = self._neural_net.log_prob(theta, x)
         else:


Test patch:
diff --git a/tests/linearGaussian_snpe_test.py b/tests/linearGaussian_snpe_test.py
--- a/tests/linearGaussian_snpe_test.py
+++ b/tests/linearGaussian_snpe_test.py
@@ -374,6 +374,48 @@ def test_api_snpe_c_posterior_correction(sample_with, mcmc_method, prior_str):
     _ = posterior.log_prob(samples)
 
 
+# Testing rejection and mcmc sampling methods.
+@pytest.mark.parametrize(
+    "force_first_round_loss, pass_proposal_to_append",
+    (
+        (True, True),
+        (True, False),
+        (False, True),
+        pytest.param(False, False, marks=pytest.mark.xfail),
+    ),
+)
+def test_api_force_first_round_loss(
+    force_first_round_loss: bool, pass_proposal_to_append: bool
+):
+    """Test that leakage correction applied to sampling works, with both MCMC and
+    rejection.
+
+    """
+
+    num_dim = 2
+    x_o = zeros(1, num_dim)
+
+    # likelihood_mean will be likelihood_shift+theta
+    likelihood_shift = -1.0 * ones(num_dim)
+    likelihood_cov = 0.3 * eye(num_dim)
+    prior = utils.BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
+
+    simulator, prior = prepare_for_sbi(
+        lambda theta: linear_gaussian(theta, likelihood_shift, likelihood_cov), prior
+    )
+    inference = SNPE_C(prior, show_progress_bars=False)
+
+    proposal = prior
+    for _ in range(2):
+        train_proposal = proposal if pass_proposal_to_append else None
+        theta, x = simulate_for_sbi(simulator, proposal, 1000)
+        _ = inference.append_simulations(theta, x, proposal=train_proposal).train(
+            force_first_round_loss=force_first_round_loss, max_num_epochs=2
+        )
+        posterior = inference.build_posterior().set_default_x(x_o)
+        proposal = posterior
+
+
 @pytest.mark.slow
 def test_sample_conditional():
     """