fix: update init_action handling in CEMSolver and LagrangianSolver tests to ensure correct tensor shapes and values

Author: quentinll

tests/solver/test_continuous.py

@@ -1,7 +1,6 @@
 """Tests for continuous solvers (CEM, MPPI, GD, Nevergrad)."""
 
 import numpy as np
-import pytest
 import torch
 from gymnasium import spaces as gym_spaces
 
@@ -15,7 +14,9 @@
 class DummyCostModel:
     """Simple Costable implementation for tests."""
 
-    def get_cost(self, info_dict: dict, action_candidates: torch.Tensor) -> torch.Tensor:
+    def get_cost(
+        self, info_dict: dict, action_candidates: torch.Tensor
+    ) -> torch.Tensor:
         # Quadratic cost: sum over horizon and action dims
         cost = action_candidates.pow(2).sum(dim=(-1, -2))
         return cost
@@ -39,7 +40,9 @@ def test_cem_solver_configure():
     """Test CEMSolver configuration."""
     model = DummyCostModel()
     solver = CEMSolver(model=model, n_steps=10)
-    action_space = gym_spaces.Box(low=-1, high=1, shape=(4, 2), dtype=np.float32)
+    action_space = gym_spaces.Box(
+        low=-1, high=1, shape=(4, 2), dtype=np.float32
+    )
     config = PlanConfig(horizon=5, receding_horizon=3, action_block=1)
 
     solver.configure(action_space=action_space, n_envs=2, config=config)
@@ -58,48 +61,60 @@ def test_cem_solver_configure_discrete_warning(caplog):
     config = PlanConfig(horizon=5, receding_horizon=3)
 
     solver.configure(action_space=action_space, n_envs=1, config=config)
-    assert "discrete" in caplog.text.lower() or solver._configured  # warning logged
+    assert (
+        'discrete' in caplog.text.lower() or solver._configured
+    )  # warning logged
 
 
 def test_cem_solver_init_action_distrib():
     """Test CEMSolver action distribution initialization."""
     model = DummyCostModel()
     solver = CEMSolver(model=model, n_steps=10)
-    action_space = gym_spaces.Box(low=-1, high=1, shape=(2, 3), dtype=np.float32)
+    action_space = gym_spaces.Box(
+        low=-1, high=1, shape=(2, 3), dtype=np.float32
+    )
     config = PlanConfig(horizon=5, receding_horizon=3)
     solver.configure(action_space=action_space, n_envs=2, config=config)
 
-    mean, var = solver.init_action_distrib()
+    actions = torch.zeros(2, 5, 3)
+    mean, var = solver.init_action_distrib(actions)
     assert mean.shape == (2, 5, 3)
     assert var.shape == (2, 5, 3)
 
 
 def test_cem_solver_init_action_distrib_with_init():
-    """Test CEMSolver action distribution with initial actions."""
+    """Test CEMSolver action distribution returns provided actions as mean."""
     model = DummyCostModel()
     solver = CEMSolver(model=model, n_steps=10)
-    action_space = gym_spaces.Box(low=-1, high=1, shape=(2, 3), dtype=np.float32)
+    action_space = gym_spaces.Box(
+        low=-1, high=1, shape=(2, 3), dtype=np.float32
+    )
     config = PlanConfig(horizon=5, receding_horizon=3)
     solver.configure(action_space=action_space, n_envs=2, config=config)
 
-    init_actions = torch.randn(2, 2, 3)
+    init_actions = torch.randn(2, 5, 3)
     mean, var = solver.init_action_distrib(init_actions)
-    assert mean.shape == (2, 5, 3)  # Padded to horizon
+    assert mean.shape == (2, 5, 3)
+    assert mean is init_actions
 
 
 def test_cem_solver_call():
     """Test CEMSolver __call__ method."""
     model = DummyCostModel()
-    solver = CEMSolver(model=model, n_steps=2, num_samples=50, batch_size=2, topk=10)
-    action_space = gym_spaces.Box(low=-1, high=1, shape=(2, 2), dtype=np.float32)
+    solver = CEMSolver(
+        model=model, n_steps=2, num_samples=50, batch_size=2, topk=10
+    )
+    action_space = gym_spaces.Box(
+        low=-1, high=1, shape=(2, 2), dtype=np.float32
+    )
     config = PlanConfig(horizon=3, receding_horizon=2)
     solver.configure(action_space=action_space, n_envs=2, config=config)
 
-    info_dict = {"pixels": torch.randn(2, 1, 3, 64, 64)}
+    info_dict = {'pixels': torch.randn(2, 1, 3, 64, 64)}
     outputs = solver(info_dict)
 
-    assert "actions" in outputs
-    assert outputs["actions"].shape == (2, 3, 2)
+    assert 'actions' in outputs
+    assert outputs['actions'].shape == (2, 3, 2)
 
 
 ###########################
@@ -110,7 +125,9 @@ def test_cem_solver_call():
 def test_icem_solver_init():
     """Test ICEMSolver initialization."""
     model = DummyCostModel()
-    solver = ICEMSolver(model=model, n_steps=10, num_samples=100, noise_beta=2.0)
+    solver = ICEMSolver(
+        model=model, n_steps=10, num_samples=100, noise_beta=2.0
+    )
     assert solver.model is model
     assert solver.n_steps == 10
     assert solver.num_samples == 100
@@ -123,7 +140,9 @@ def test_icem_solver_configure():
     """Test ICEMSolver configuration."""
     model = DummyCostModel()
     solver = ICEMSolver(model=model, n_steps=10)
-    action_space = gym_spaces.Box(low=-1, high=1, shape=(4, 2), dtype=np.float32)
+    action_space = gym_spaces.Box(
+        low=-1, high=1, shape=(4, 2), dtype=np.float32
+    )
     config = PlanConfig(horizon=5, receding_horizon=3, action_block=1)
 
     solver.configure(action_space=action_space, n_envs=2, config=config)
@@ -140,43 +159,59 @@ def test_icem_solver_init_action_distrib():
     """Test ICEMSolver action distribution initialization."""
     model = DummyCostModel()
     solver = ICEMSolver(model=model, n_steps=10)
-    action_space = gym_spaces.Box(low=-1, high=1, shape=(2, 3), dtype=np.float32)
+    action_space = gym_spaces.Box(
+        low=-1, high=1, shape=(2, 3), dtype=np.float32
+    )
     config = PlanConfig(horizon=5, receding_horizon=3)
     solver.configure(action_space=action_space, n_envs=2, config=config)
 
-    mean, var = solver.init_action_distrib()
+    actions = torch.zeros(2, 5, 3)
+    mean, var = solver.init_action_distrib(actions)
     assert mean.shape == (2, 5, 3)
     assert var.shape == (2, 5, 3)
 
 
 def test_icem_solver_call():
     """Test ICEMSolver __call__ method."""
     model = DummyCostModel()
-    solver = ICEMSolver(model=model, n_steps=2, num_samples=50, batch_size=2, topk=10)
-    action_space = gym_spaces.Box(low=-1, high=1, shape=(2, 2), dtype=np.float32)
+    solver = ICEMSolver(
+        model=model, n_steps=2, num_samples=50, batch_size=2, topk=10
+    )
+    action_space = gym_spaces.Box(
+        low=-1, high=1, shape=(2, 2), dtype=np.float32
+    )
     config = PlanConfig(horizon=3, receding_horizon=2)
     solver.configure(action_space=action_space, n_envs=2, config=config)
 
-    info_dict = {"pixels": torch.randn(2, 1, 3, 64, 64)}
+    info_dict = {'pixels': torch.randn(2, 1, 3, 64, 64)}
     outputs = solver(info_dict)
 
-    assert "actions" in outputs
-    assert outputs["actions"].shape == (2, 3, 2)
+    assert 'actions' in outputs
+    assert outputs['actions'].shape == (2, 3, 2)
 
 
 def test_icem_solver_white_noise_fallback():
     """Test ICEMSolver with beta=0 (white noise, equivalent to standard CEM)."""
     model = DummyCostModel()
-    solver = ICEMSolver(model=model, n_steps=2, num_samples=50, batch_size=2, topk=10, noise_beta=0.0)
-    action_space = gym_spaces.Box(low=-1, high=1, shape=(2, 2), dtype=np.float32)
+    solver = ICEMSolver(
+        model=model,
+        n_steps=2,
+        num_samples=50,
+        batch_size=2,
+        topk=10,
+        noise_beta=0.0,
+    )
+    action_space = gym_spaces.Box(
+        low=-1, high=1, shape=(2, 2), dtype=np.float32
+    )
     config = PlanConfig(horizon=3, receding_horizon=2)
     solver.configure(action_space=action_space, n_envs=2, config=config)
 
-    info_dict = {"pixels": torch.randn(2, 1, 3, 64, 64)}
+    info_dict = {'pixels': torch.randn(2, 1, 3, 64, 64)}
     outputs = solver(info_dict)
 
-    assert "actions" in outputs
-    assert outputs["actions"].shape == (2, 3, 2)
+    assert 'actions' in outputs
+    assert outputs['actions'].shape == (2, 3, 2)
 
 
 ###########################
@@ -197,7 +232,9 @@ def test_mppi_solver_configure():
     """Test MPPISolver configuration."""
     model = DummyCostModel()
     solver = MPPISolver(model=model, n_steps=10)
-    action_space = gym_spaces.Box(low=-1, high=1, shape=(4, 2), dtype=np.float32)
+    action_space = gym_spaces.Box(
+        low=-1, high=1, shape=(4, 2), dtype=np.float32
+    )
     config = PlanConfig(horizon=5, receding_horizon=3)
 
     solver.configure(action_space=action_space, n_envs=2, config=config)
@@ -212,11 +249,14 @@ def test_mppi_solver_init_action_distrib():
     """Test MPPISolver action distribution initialization."""
     model = DummyCostModel()
     solver = MPPISolver(model=model, n_steps=10)
-    action_space = gym_spaces.Box(low=-1, high=1, shape=(2, 3), dtype=np.float32)
+    action_space = gym_spaces.Box(
+        low=-1, high=1, shape=(2, 3), dtype=np.float32
+    )
     config = PlanConfig(horizon=5, receding_horizon=3)
     solver.configure(action_space=action_space, n_envs=2, config=config)
 
-    mean, var = solver.init_action_distrib()
+    actions = torch.zeros(2, 5, 3)
+    mean, var = solver.init_action_distrib(actions)
     assert mean.shape == (2, 5, 3)
     assert var.shape == (2, 5, 3)
 
@@ -225,15 +265,17 @@ def test_mppi_solver_call():
     """Test MPPISolver __call__ method."""
     model = DummyCostModel()
     solver = MPPISolver(model=model, n_steps=2, num_samples=10, batch_size=2)
-    action_space = gym_spaces.Box(low=-1, high=1, shape=(2, 2), dtype=np.float32)
+    action_space = gym_spaces.Box(
+        low=-1, high=1, shape=(2, 2), dtype=np.float32
+    )
     config = PlanConfig(horizon=3, receding_horizon=2)
     solver.configure(action_space=action_space, n_envs=2, config=config)
 
-    info_dict = {"pixels": torch.randn(2, 1, 3, 64, 64)}
+    info_dict = {'pixels': torch.randn(2, 1, 3, 64, 64)}
     outputs = solver(info_dict)
 
-    assert "actions" in outputs
-    assert outputs["actions"].shape == (2, 3, 2)
+    assert 'actions' in outputs
+    assert outputs['actions'].shape == (2, 3, 2)
 
 
 ###########################
@@ -254,7 +296,9 @@ def test_gradient_solver_configure():
     """Test GradientSolver configuration."""
     model = DummyCostModel()
     solver = GradientSolver(model=model, n_steps=10)
-    action_space = gym_spaces.Box(low=-1, high=1, shape=(4, 2), dtype=np.float32)
+    action_space = gym_spaces.Box(
+        low=-1, high=1, shape=(4, 2), dtype=np.float32
+    )
     config = PlanConfig(horizon=5, receding_horizon=3)
 
     solver.configure(action_space=action_space, n_envs=2, config=config)
@@ -269,25 +313,36 @@ def test_gradient_solver_init_action():
     """Test GradientSolver action initialization."""
     model = DummyCostModel()
     solver = GradientSolver(model=model, n_steps=10, num_samples=3)
-    action_space = gym_spaces.Box(low=-1, high=1, shape=(2, 3), dtype=np.float32)
+    action_space = gym_spaces.Box(
+        low=-1, high=1, shape=(2, 3), dtype=np.float32
+    )
     config = PlanConfig(horizon=5, receding_horizon=3)
     solver.configure(action_space=action_space, n_envs=2, config=config)
 
-    solver.init_action()
-    assert hasattr(solver, "init")
-    assert solver.init.shape == (2, 3, 5, 3)  # (n_envs, num_samples, horizon, action_dim)
+    solver.init_action(torch.zeros(2, 5, 3))
+    assert hasattr(solver, 'init')
+    assert solver.init.shape == (
+        2,
+        3,
+        5,
+        3,
+    )  # (n_envs, num_samples, horizon, action_dim)
 
 
 def test_gradient_solver_call():
     """Test GradientSolver __call__ method."""
     model = DummyCostModel()
-    solver = GradientSolver(model=model, n_steps=2, num_samples=2, batch_size=2)
-    action_space = gym_spaces.Box(low=-1, high=1, shape=(2, 2), dtype=np.float32)
+    solver = GradientSolver(
+        model=model, n_steps=2, num_samples=2, batch_size=2
+    )
+    action_space = gym_spaces.Box(
+        low=-1, high=1, shape=(2, 2), dtype=np.float32
+    )
     config = PlanConfig(horizon=3, receding_horizon=2)
     solver.configure(action_space=action_space, n_envs=2, config=config)
 
-    info_dict = {"pixels": torch.randn(2, 1, 3, 64, 64)}
+    info_dict = {'pixels': torch.randn(2, 1, 3, 64, 64)}
     outputs = solver(info_dict)
 
-    assert "actions" in outputs
-    assert outputs["actions"].shape == (2, 3, 2)
+    assert 'actions' in outputs
+    assert outputs['actions'].shape == (2, 3, 2)

tests/solver/test_lagrangian.py

@@ -1,11 +1,8 @@
 """Tests for LagrangianSolver class."""
 
-import dataclasses
-
 import numpy as np
 import pytest
 import torch
-import torch.nn.functional as F
 from gymnasium import spaces as gym_spaces
 
 from stable_worldmodel.policy import PlanConfig
@@ -170,31 +167,30 @@ def test_n_envs_before_configure_is_none():
 ###########################
 
 
-def test_init_action_none_creates_zeros():
-    """init_action(None) creates zero tensor of correct shape."""
+def test_init_action_zeros_creates_zero_init():
+    """init_action with a zeros tensor creates a zero init of correct shape."""
     solver = make_solver(num_samples=1, var_scale=0.0)
     configure(solver, action_dim=4, n_envs=2, horizon=5, action_block=1)
 
     with torch.no_grad():
-        solver.init_action(None)
+        solver.init_action(torch.zeros(2, 5, 4))
 
     # shape: (n_envs, num_samples, horizon, action_dim)
     assert solver.init.shape == (2, 1, 5, 4)
     assert torch.allclose(solver.init, torch.zeros_like(solver.init))
 
 
-def test_init_action_partial_fills_remaining_with_zeros():
-    """Providing fewer steps pads the rest with zeros (first sample only)."""
+def test_init_action_full_tensor_preserved_as_first_sample():
+    """Providing a full-horizon tensor preserves it as the first sample."""
     solver = make_solver(num_samples=1, var_scale=0.0)
     configure(solver, action_dim=4, n_envs=2, horizon=6, action_block=1)
 
-    init = torch.ones(2, 3, 4)  # 3 of 6 steps
+    init = torch.ones(2, 6, 4)
     with torch.no_grad():
         solver.init_action(init)
 
     first_sample = solver.init[:, 0]  # (n_envs, horizon, action_dim)
-    assert torch.allclose(first_sample[:, :3], init)
-    assert torch.allclose(first_sample[:, 3:], torch.zeros(2, 3, 4))
+    assert torch.allclose(first_sample, init)
 
 
 def test_init_action_full_horizon_preserved():
@@ -449,7 +445,7 @@ def test_persist_multipliers_warm_starts():
     configure(solver, action_dim=4, n_envs=2, horizon=4, action_block=1)
 
     out1 = solver.solve({})
-    lambdas_after_first = out1['lambdas'].clone()
+    _lambdas_after_first = out1['lambdas'].clone()
 
     out2 = solver.solve({})
     lambdas_after_second = out2['lambdas']

tests/test_policy.py

@@ -723,7 +723,7 @@ def test_worldmodel_policy_warmstart_extends_partial_plan(actionable_setup):
 
 
 def test_worldmodel_policy_no_warmstart_without_actionable():
-    """Solver receives None init_action when model does not implement Actionable."""
+    """Solver receives a full zero init_action when model does not implement Actionable."""
     non_actionable_model = MagicMock(spec=['get_cost'])
     solver = MockSolverWithWarmStart(model=non_actionable_model)
     config = PlanConfig(
@@ -741,7 +741,8 @@ def test_worldmodel_policy_no_warmstart_without_actionable():
 
     policy.get_action(info)
 
-    assert solver.received_init_action is None
+    assert solver.received_init_action.shape == (1, 5, 2)
+    assert solver.received_init_action.eq(0).all()
 
 
 ###########################