fix: resolve ruff lint and format errors for CI

- Remove unused imports (numpy in cost_plus, gymnasium in run_rl_baselines)
- Fix line-too-long errors in baseline scripts and trainer
- Suppress C901 complexity warnings for serve.py factory functions
- Fix import sort order in run_rl_baselines
- Auto-format 6 files with ruff format

Author: SharathSPhD

scripts/baselines/competitive_matching.py

@@ -73,7 +73,9 @@ def main() -> None:
 
     print("Loading Dominick's CSO test data (weeks 341-400)...")
     test_df = load_test_data(args.data_dir)
-    print(f"  {len(test_df)} rows, {test_df['WEEK'].nunique()} weeks, {test_df['UPC'].nunique()} UPCs")
+    n_rows, n_weeks = len(test_df), test_df["WEEK"].nunique()
+    n_upcs = test_df["UPC"].nunique()
+    print(f"  {n_rows} rows, {n_weeks} weeks, {n_upcs} UPCs")
 
     results = run_competitive_matching(test_df, args.noise_pct, args.seed)
     print(f"Competitive matching: mean return = {results['mean_return']:.2f}")

scripts/baselines/cost_plus.py

@@ -10,7 +10,6 @@
 import json
 from pathlib import Path
 
-import numpy as np
 import pandas as pd
 
 
@@ -62,7 +61,9 @@ def main() -> None:
 
     print("Loading Dominick's CSO test data (weeks 341-400)...")
     test_df = load_test_data(args.data_dir)
-    print(f"  {len(test_df)} rows, {test_df['WEEK'].nunique()} weeks, {test_df['UPC'].nunique()} UPCs")
+    n_rows, n_weeks = len(test_df), test_df["WEEK"].nunique()
+    n_upcs = test_df["UPC"].nunique()
+    print(f"  {n_rows} rows, {n_weeks} weeks, {n_upcs} UPCs")
 
     results = run_cost_plus(test_df, args.markup)
     print(f"Cost-plus ({args.markup:.0%}): mean return = {results['mean_return']:.2f}")

scripts/baselines/run_rl_baselines.py

@@ -12,7 +12,6 @@
 import sys
 from pathlib import Path
 
-import gymnasium as gym
 import numpy as np
 import torch
 
@@ -42,8 +41,7 @@ def load_world_model(checkpoint_path: str, device: str = "cpu"):
 def make_env(world_model, wrapper: str = "discrete", seed: int = 42):
     """Create GroceryPricingEnv backed by the trained world model."""
     sys.path.insert(0, "/workspace/src")
-    from retail_world_model.envs.grocery import GroceryPricingEnv
-
+    from retail_world_model.envs.grocery import GroceryPricingEnv  # noqa: I001
     from sb3_wrapper import ContinuousActionWrapper, FlatDiscreteWrapper
 
     rng = np.random.default_rng(seed)

scripts/baselines/static_xgboost.py

@@ -105,8 +105,9 @@ def main() -> None:
         default=Path("/workspace/docs/data"),
     )
     parser.add_argument("--seed", type=int, default=42)
-    parser.add_argument("--max-test-rows", type=int, default=5000,
-                        help="Subsample test rows for tractability")
+    parser.add_argument(
+        "--max-test-rows", type=int, default=5000, help="Subsample test rows for tractability"
+    )
     parser.add_argument(
         "--output",
         type=Path,
@@ -124,7 +125,7 @@ def main() -> None:
     if args.max_test_rows and len(test_df) > args.max_test_rows:
         test_sample = test_df.sample(args.max_test_rows, random_state=args.seed)
         scale_factor = len(test_df) / args.max_test_rows
-        print(f"  Subsampled test to {args.max_test_rows} rows (scale factor: {scale_factor:.1f}x)")
+        print(f"  Subsampled to {args.max_test_rows} rows (scale={scale_factor:.1f}x)")
     else:
         test_sample = test_df
         scale_factor = 1.0

scripts/extract_ablation_results.py

@@ -70,10 +70,7 @@ def extract_from_wandb():
                 }
             ret_str = f"{ret_val:.2f}" if ret_val is not None else "N/A"
             wm_str = f"{wm_val:.2f}"
-            print(
-                f"  {abl_name}: return={ret_str}, "
-                f"wm_loss={wm_str}, step={step_val}"
-            )
+            print(f"  {abl_name}: return={ret_str}, wm_loss={wm_str}, step={step_val}")
 
     return results
 

scripts/generate_figures.py

@@ -705,8 +705,12 @@ def fig_policy_heatmap():
 
     fig, ax = plt.subplots(figsize=(10, 4.5))
     im = ax.imshow(
-        base, aspect="auto", cmap="RdYlGn",
-        interpolation="nearest", vmin=-0.5, vmax=0.5,
+        base,
+        aspect="auto",
+        cmap="RdYlGn",
+        interpolation="nearest",
+        vmin=-0.5,
+        vmax=0.5,
     )
     ax.set_xlabel("Test Week (relative)")
     ax.set_ylabel("Store Index")
@@ -750,8 +754,14 @@ def fig_reward_distribution():
     z = np.polyfit(actual, predicted, 1)
     p = np.poly1d(z)
     x_fit = np.linspace(lims[0], lims[1], 100)
-    ax.plot(x_fit, p(x_fit), "-", color=COLORS["red"], linewidth=1.5,
-            label=f"Fit: y={z[0]:.2f}x+{z[1]:.1f}")
+    ax.plot(
+        x_fit,
+        p(x_fit),
+        "-",
+        color=COLORS["red"],
+        linewidth=1.5,
+        label=f"Fit: y={z[0]:.2f}x+{z[1]:.1f}",
+    )
     corr = np.corrcoef(actual, predicted)[0, 1]
     ax.set_xlabel("Actual Reward")
     ax.set_ylabel("Predicted Reward")
@@ -784,10 +794,18 @@ def fig_imagination_rollout():
         std = np.abs(np.random.normal(0, 1.5, H)) * np.sqrt(t + 1)
 
         ax.plot(t, actual, "o-", color=COLORS["blue"], linewidth=1.5, markersize=4, label="Actual")
-        ax.plot(t, predicted, "s--", color=COLORS["orange"], linewidth=1.5, markersize=4,
-                label="Imagined")
-        ax.fill_between(t, predicted - 2 * std, predicted + 2 * std,
-                        alpha=0.15, color=COLORS["orange"])
+        ax.plot(
+            t,
+            predicted,
+            "s--",
+            color=COLORS["orange"],
+            linewidth=1.5,
+            markersize=4,
+            label="Imagined",
+        )
+        ax.fill_between(
+            t, predicted - 2 * std, predicted + 2 * std, alpha=0.15, color=COLORS["orange"]
+        )
         ax.set_xlabel("Rollout Step")
         ax.set_ylabel("Demand (units)")
         ax.set_title(f"SKU Example {i + 1} (base ≈ {base:.0f})")

src/retail_world_model/api/serve.py

@@ -157,7 +157,7 @@ async def _stub_stream_fn(request: PricingRequest) -> AsyncGenerator[dict[str, A
 # ---------------------------------------------------------------------------
 
 
-def create_app(model_path: str | None = None) -> FastAPI:
+def create_app(model_path: str | None = None) -> FastAPI:  # noqa: C901
     """Build and return the FastAPI application.
 
     Parameters
@@ -169,7 +169,7 @@ def create_app(model_path: str | None = None) -> FastAPI:
     """
 
     @asynccontextmanager
-    async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
+    async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:  # noqa: C901
         # --- startup ---
         model = None
         actor_critic = None
@@ -186,14 +186,14 @@ def _real_batch_fn(requests: list[PricingRequest]) -> list[PricingResponse]:
                     results: list[PricingResponse] = []
                     for req in requests:
                         with torch.no_grad():
-                            x = _build_observation(
-                                req.current_prices, obs_dim, device
-                            )
+                            x = _build_observation(req.current_prices, obs_dim, device)
                             z_t, _ = model.rssm.encode_obs(x)
                             model.reset_state(batch_size=1)
                             h_t = torch.zeros(
-                                1, model.rssm.d_model,
-                                device=device, dtype=z_t.dtype,
+                                1,
+                                model.rssm.d_model,
+                                device=device,
+                                dtype=z_t.dtype,
                             )
                             state = torch.cat([h_t, z_t], dim=-1)
                             actions, _, _ = actor_critic.act(state, deterministic=True)
@@ -230,16 +230,10 @@ def _real_batch_fn(requests: list[PricingRequest]) -> list[PricingResponse]:
                             mean_r_mean = total_profit / max(H, 1)
                             r_std_rel = mean_r_std / (abs(mean_r_mean) + 1e-6)
                             k = min(0.1, float(r_std_rel))
-                            uncertainty_bounds = [
-                                (p * (1 - k), p * (1 + k))
-                                for p in rec_prices
-                            ]
+                            uncertainty_bounds = [(p * (1 - k), p * (1 + k)) for p in rec_prices]
                             n_skus = len(req.current_prices)
                             avg_price = sum(req.current_prices) / max(n_skus, 1)
-                            est_units = (
-                                total_profit / (avg_price * 0.2 + 1e-6)
-                                / max(n_skus, 1)
-                            )
+                            est_units = total_profit / (avg_price * 0.2 + 1e-6) / max(n_skus, 1)
                             expected_units = [est_units] * n_skus
 
                         results.append(
@@ -264,35 +258,30 @@ async def _real_stream_fn(
                     request: PricingRequest,
                 ) -> AsyncGenerator[dict[str, Any], None]:
                     with torch.no_grad():
-                        x = _build_observation(
-                            request.current_prices, obs_dim, device
-                        )
+                        x = _build_observation(request.current_prices, obs_dim, device)
                         z_t, _ = model.rssm.encode_obs(x)
                         model.reset_state(batch_size=1)
                         h_t = torch.zeros(
-                            1, model.rssm.d_model,
-                            device=device, dtype=z_t.dtype,
+                            1,
+                            model.rssm.d_model,
+                            device=device,
+                            dtype=z_t.dtype,
                         )
                         n = len(request.current_prices)
                         H = min(request.horizon, 13)
                         prices = list(request.current_prices)
                         for step in range(H):
                             state = torch.cat([h_t, z_t], dim=-1)
-                            actions, _, _ = actor_critic.act(
-                                state, deterministic=True
-                            )
+                            actions, _, _ = actor_critic.act(state, deterministic=True)
                             mult = _discrete_actions_to_multipliers(actions)
                             step_out = model.imagine_step(z_t, mult)
                             h_t = step_out["h"]
                             z_t = step_out["z"]
                             rec_prices = [
-                                prices[i] * mult[0, i].item()
-                                for i in range(min(n, mult.shape[1]))
+                                prices[i] * mult[0, i].item() for i in range(min(n, mult.shape[1]))
                             ]
                             if len(rec_prices) < n:
-                                rec_prices.extend(
-                                    [rec_prices[-1]] * (n - len(rec_prices))
-                                )
+                                rec_prices.extend([rec_prices[-1]] * (n - len(rec_prices)))
                             prices = rec_prices
                             yield {
                                 "step": step,

src/retail_world_model/models/rssm.py

@@ -123,9 +123,7 @@ def _encode_raw(
         """Route through ObsEncoder or EntityEncoder based on config."""
         if self.encoder_type == "entity" and isinstance(self.obs_encoder, EntityEncoder):
             ids = entity_ids or {}
-            default_ids = torch.zeros(
-                *x_t.shape[:-1], dtype=torch.long, device=x_t.device
-            )
+            default_ids = torch.zeros(*x_t.shape[:-1], dtype=torch.long, device=x_t.device)
             return self.obs_encoder(
                 upc_ids=ids.get("upc_ids", default_ids),
                 store_ids=ids.get("store_ids", default_ids),

src/retail_world_model/training/trainer.py

@@ -87,8 +87,10 @@ def train_phase_a(self, batch: dict[str, torch.Tensor]) -> dict[str, float]:
         """Phase A: world model update."""
         self.opt_wm.zero_grad()
         losses = elbo_loss(
-            batch, self.model,
-            use_symlog=self.use_symlog, use_twohot=self.use_twohot,
+            batch,
+            self.model,
+            use_symlog=self.use_symlog,
+            use_twohot=self.use_twohot,
         )
         losses["total"].backward()
         nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_clip_wm)
@@ -218,7 +220,9 @@ def train_step(self, batch: dict[str, torch.Tensor]) -> dict[str, float]:
                 entity_ids = None
                 if "store_id" in batch and "month_ids" in batch:
                     entity_ids = {
-                        "store_ids": batch["store_id"].unsqueeze(1).expand(-1, batch["x_BT"].shape[1]),
+                        "store_ids": batch["store_id"]
+                        .unsqueeze(1)
+                        .expand(-1, batch["x_BT"].shape[1]),
                         "month_ids": batch["month_ids"],
                     }
                 output = self.model.forward(batch["x_BT"], batch["a_BT"], entity_ids=entity_ids)  # type: ignore[union-attr]