feat: v4.1 portfolio MDP — capital tracking + tx_cost on all tasks

- Add current_holding, capital, drawdown to MultiStockObservation
- Track virtual capital for ALL tasks (was only Task 3)
- Apply 0.1% transaction cost when switching stock holdings
- Per-task drawdown limits: short=5%, medium=10%, long=15%
- Expose portfolio state in inference.py prompt
- Version bump to 4.1.0

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

env/indicators_env.py

@@ -1,19 +1,46 @@
 """
-indicators_env.py — OpenEnv-compatible RL environment for technical indicators prediction.
-
-Implements the full OpenEnv standard interface:
-  - HTTP endpoints: reset / step / state / tasks / grader / baseline / health
-  - WebSocket (/ws) for low-latency sequential step() calls during RL training
-
-Env spec:
-  Observation : JSON dict of full indicator snapshot + [TERM: MEDIUM] token
-  Action      : {"direction": "Bullish"|"Bearish"|"Neutral", "conviction": float 0-1}
-  Reward      : Shaped 0.0–1.1 based on correctness + conviction calibration
-
-Tasks:
-  Task 1 (Easy)   : SHORT  term (5d,  ±1.0% threshold)
-  Task 2 (Medium) : MEDIUM term (20d, ±2.5% threshold)  ← default
-  Task 3 (Hard)   : LONG   term (60d, ±5.0% threshold) + conviction >= 0.7 required
+indicators_env.py — IndicatorsEnv v4.1: Multi-stock Portfolio MDP
+
+OpenEnv-compatible RL environment for NSE equity analysis.
+
+Full OpenEnv interface:
+  HTTP  : GET  /health  /tasks
+          POST /reset  /step  /grader
+          GET  /state  /baseline
+  WebSocket: /ws
+
+Env design (v4.1) — Multi-stock Portfolio MDP:
+  At each step the agent observes 3 stocks from the same NSE sector.
+  It picks ONE stock (or passes with NONE) and declares a direction.
+
+  Observation:
+    - 3 stocks from the same sector (same stocks for the whole episode)
+    - Full technical indicator snapshot per stock
+    - RSI + price-momentum signal history (accumulated across steps)
+    - Macro context in Task 3 (NIFTY50 trend, market regime)
+
+  Action:
+    {"stock": "HDFCBANK", "direction": "Bullish"|"Bearish"|"NONE", "conviction": 0.8}
+    NONE = skip this step (preserve selective participation)
+
+  Reward (market-neutral alpha):
+    alpha  = chosen_stock_period_return − sector_avg_period_return
+    reward = alpha × direction_sign × conviction × REWARD_SCALE
+    → random policy earns ~0 expected reward (sector avg cancels market beta)
+    → skilled policy earns consistently positive reward (correct stock selection)
+    → Kelly conviction semantics: fraction of virtual wealth wagered per step
+
+  Step spacing:
+    short  :  1 trading day per step → GT = 1-day return   (5 steps  = 1 week)
+    medium :  5 trading days per step → GT = 5-day return  (10 steps = 10 weeks)
+    long   : 20 trading days per step → GT = 20-day return (15 steps = 15 months)
+    GT window = step spacing → zero overlap, reward and GT measure identical period.
+
+  Task 3 extras:
+    - Drawdown limit: virtual capital terminates at 10% drawdown
+    - Macro context (NIFTY50) injected into each observation
+
+Built for the Meta × PyTorch Hackathon.
 """
 
 from __future__ import annotations
@@ -21,18 +48,33 @@ from __future__ import annotations
 import json
 import logging
 import random
+import sys
 import uuid
+from statistics import mean
 from typing import Any, Dict, List, Optional
+from datetime import datetime, timedelta
+
+try:
+    import yf_patch
+    yf_patch.patch_yfinance_globally()
+except ImportError:
+    pass
 
+import pandas as pd
 import uvicorn
+import numpy as np
 from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
 from pydantic import BaseModel, Field
 
 from data_loader import (
-    NSE_UNIVERSE,
-    build_observation,
-    compute_ground_truth,
-    generate_scenario_pool,
+    NSE_UNIVERSE, 
+    SECTOR_GROUPS, 
+    TERM_WINDOWS, 
+    TERM_THRESHOLDS,
+    STEP_SPACING,
+    PERIOD_THRESHOLDS,
+    build_multi_stock_episode,
+    fetch_macro_context
 )
 
 logging.basicConfig(level=logging.INFO)
@@ -41,107 +83,177 @@ logger = logging.getLogger(__name__)
 app = FastAPI(
     title="IndicatorsEnv",
     description=(
-        "OpenEnv-compatible RL environment for NSE technical indicators prediction. "
-        "An AI agent receives a rich indicator snapshot and must predict the directional "
-        "price move over a configurable forward window (short/medium/long term). "
+        "IndicatorsEnv v4.1 — Multi-stock Portfolio MDP for NSE equity analysis. "
+        "At each step the agent observes 3 stocks from the same NSE sector and picks one "
+        "(or passes). Reward = (chosen stock return − sector average) × direction × conviction. "
+        "Market-neutral alpha with portfolio capital tracking: holding a position avoids "
+        "transaction cost (0.1%); switching incurs it. Capital causally tracks cumulative alpha P&L. "
+        "Drawdown limits apply to all tasks. Three tasks of increasing difficulty. "
         "Built for the Meta × PyTorch Hackathon."
     ),
-    version="2.0.0",
+    version="4.1.0",
 )
 
-# ─── Task Definitions ─────────────────────────────────────────────────────────
+# ─── Environment constants ────────────────────────────────────────────────────
+
+TASK_MAX_STEPS: Dict[str, int] = {
+    "short":  5,    # 5 × 1-day  = 1 trading week
+    "medium": 10,   # 10 × 5-day = 10 trading weeks
+    "long":   15,   # 15 × 20-day = 15 months
+}
+
+REWARD_SCALE     = 50.0   # ×50: 2% alpha × conviction 0.7 × 50 ≈ 0.7 reward
+TRANSACTION_COST = 0.001  # 0.1% of capital on holding switch (exploration cost)
+
+DRAWDOWN_LIMITS: Dict[str, float] = {
+    "short":  0.05,   # 5%  — 5 steps, tight risk
+    "medium": 0.10,   # 10% — 10 steps
+    "long":   0.15,   # 15% — 15 steps, more recovery time
+}
+
+# ─── Task definitions ─────────────────────────────────────────────────────────
 
 TASKS = [
     {
         "id": "short_term_direction",
-        "name": "Short-term Direction (Easy)",
+        "name": "Short-term Relative Alpha (Easy)",
         "description": (
-            "Predict 5-day forward price direction from technical indicators. "
-            "Threshold: ±1.0%. More Neutral outcomes expected due to market noise."
+            "5 steps, 1 trading day apart (= 1 week of daily observations). "
+            "At each step observe 3 stocks from the same NSE sector. "
+            "Pick one stock and declare Bullish (long), Bearish (short), or NONE (pass). "
+            "Reward = (chosen_return − sector_avg) × direction × conviction × 50. "
+            "Market-neutral: sector beta cancels — the agent profits from within-sector "
+            "relative momentum, not from broad market moves. "
+            "GT = 1-day forward return (±0.3% threshold). "
+            "Virtual capital tracked across steps; 0.1% transaction cost when switching holdings; "
+            "episode ends early at 5% drawdown."
         ),
         "difficulty": "easy",
         "term": "short",
-        "grader_note": "Score = fraction of correct directional predictions over 10 episodes.",
+        "episode_steps": TASK_MAX_STEPS["short"],
+        "step_spacing_days": STEP_SPACING["short"],
+        "grader_note": "Directional accuracy on active (non-NONE) steps.",
         "action_schema": {
-            "direction": {"type": "string", "enum": ["Bullish", "Bearish", "Neutral"]},
-            "conviction": {"type": "number", "minimum": 0.0, "maximum": 1.0},
+            "stock":     {"type": "string", "description": "NSE symbol or 'NONE' to skip"},
+            "direction": {"type": "string", "enum": ["Bullish", "Bearish", "NONE"]},
+            "conviction":{"type": "number", "minimum": 0.0, "maximum": 1.0},
         },
     },
     {
         "id": "medium_term_direction",
-        "name": "Medium-term Direction (Medium)",
+        "name": "Medium-term Relative Alpha (Medium)",
         "description": (
-            "Predict 20-day forward price direction from technical indicators. "
-            "Threshold: ±2.5%. Balanced Bullish/Bearish/Neutral distribution."
+            "10 steps, 5 trading days apart (= 10 weeks, one observation per week). "
+            "Same 3-stock multi-stock structure. "
+            "Signal history accumulates across steps — RSI trends and price momentum "
+            "are visible for all 3 stocks at each step, enabling sequential inference. "
+            "Reward = (chosen_return − sector_avg) × direction × conviction × 50. "
+            "GT = 5-day forward return (±1.5% threshold). "
+            "Grader gives extra weight to Bearish/Neutral correct calls (anti-majority-bias). "
+            "Virtual capital tracked; 0.1% switching cost; 10% drawdown terminates episode early."
         ),
         "difficulty": "medium",
         "term": "medium",
+        "episode_steps": TASK_MAX_STEPS["medium"],
+        "step_spacing_days": STEP_SPACING["medium"],
         "grader_note": (
-            "Score = weighted accuracy: Bearish/Neutral correct = 1.5x weight "
-            "(anti-majority-class bias), normalized to 0–1."
+            "Weighted accuracy on active steps: "
+            "Bearish/Neutral correct = 1.5× weight. "
+            "Participation bonus: score × (0.9 + 0.1 × active_rate)."
         ),
         "action_schema": {
-            "direction": {"type": "string", "enum": ["Bullish", "Bearish", "Neutral"]},
-            "conviction": {"type": "number", "minimum": 0.0, "maximum": 1.0},
+            "stock":     {"type": "string", "description": "NSE symbol or 'NONE' to skip"},
+            "direction": {"type": "string", "enum": ["Bullish", "Bearish", "NONE"]},
+            "conviction":{"type": "number", "minimum": 0.0, "maximum": 1.0},
         },
     },
     {
         "id": "long_term_conviction",
-        "name": "Long-term Conviction (Hard)",
+        "name": "Long-term Risk-Constrained Alpha (Hard)",
         "description": (
-            "Predict 60-day forward price direction AND demonstrate calibrated conviction. "
-            "Threshold: ±5.0%. Correct direction WITH conviction >= 0.7 scores 1.0. "
-            "Correct direction with low conviction scores 0.5. Wrong + high conviction penalized."
+            "15 steps, 20 trading days apart (= 15 months, one observation per month). "
+            "Spans multiple market regimes — signal history essential for detecting "
+            "multi-month relative momentum shifts within the sector. "
+            "Macro context (NIFTY50 trend, market regime) added to each observation. "
+            "Capital-tracked risk management: 0.1% switching cost + 15% drawdown limit. "
+            "The agent must manage risk across months as well as generate alpha. "
+            "Correct calls with conviction ≥ 0.7 score highest; "
+            "overconfident wrong predictions penalized (conviction ≥ 0.8 on wrong). "
+            "GT = 20-day forward return (±2.5% threshold)."
         ),
         "difficulty": "hard",
         "term": "long",
+        "episode_steps": TASK_MAX_STEPS["long"],
+        "step_spacing_days": STEP_SPACING["long"],
         "grader_note": (
-            "Score = mean of per-episode scores: 1.0 (correct + conviction>=0.7), "
-            "0.5 (correct + conviction<0.7), 0.0 (wrong), -0.1 (wrong + conviction>=0.8). "
-            "Normalized to 0–1 range."
+            "Conviction-calibrated accuracy on active steps: "
+            "correct+conviction≥0.7 = 1.0, correct+conviction<0.7 = 0.5, "
+            "wrong = 0.0, overconfident wrong (conviction≥0.8) = −0.1. "
+            "Normalized to (0, 1)."
         ),
         "action_schema": {
-            "direction": {"type": "string", "enum": ["Bullish", "Bearish", "Neutral"]},
-            "conviction": {"type": "number", "minimum": 0.0, "maximum": 1.0},
+            "stock":     {"type": "string", "description": "NSE symbol or 'NONE' to skip"},
+            "direction": {"type": "string", "enum": ["Bullish", "Bearish", "NONE"]},
+            "conviction":{"type": "number", "minimum": 0.0, "maximum": 1.0},
         },
     },
 ]
 
-TASK_BY_ID = {t["id"]: t for t in TASKS}
+TASK_BY_ID      = {t["id"]: t for t in TASKS}
 TERM_TO_TASK_ID = {t["term"]: t["id"] for t in TASKS}
 
 # ─── Pydantic schemas ─────────────────────────────────────────────────────────
 
-class IndicatorsAction(BaseModel):
-    direction: str = Field(..., description="Predicted direction: Bullish | Bearish | Neutral")
-    conviction: float = Field(0.5, ge=0.0, le=1.0, description="Confidence in prediction (0-1)")
 
+class StockState(BaseModel):
+    """Per-stock snapshot at one episode step."""
+    symbol:             str
+    current_price:      float
+    rsi_14:             float
+    rsi_trend:          str          # "up" | "down" | "flat"
+    price_momentum_pct: float        # cumulative return vs episode start (%)
+    indicators:         Dict[str, Any]
 
-class IndicatorsObservation(BaseModel):
-    symbol: str
-    date: str
-    term: str
-    current_price: float
-    indicators: Dict[str, Any]
+
+class MultiStockObservation(BaseModel):
+    """Full observation returned by reset() and step()."""
+    step:             int               # 1-indexed current step
+    max_steps:        int
+    term:             str               # SHORT | MEDIUM | LONG
+    sector:           str               # e.g. "banking"
+    available_stocks: List[str]         # 3 symbols
+    stocks:           Dict[str, StockState]   # symbol → full state
+    signal_history:   List[Dict[str, Any]]    # chronological pick log
+    macro:            Optional[Dict[str, Any]] = None   # Task 3 only
+    # v4.1 portfolio fields (all tasks)
+    current_holding:  str   = "NONE"   # stock held from last active step, or "NONE"
+    capital:          float = 1.0      # virtual capital (starts 1.0, tracks cumulative alpha P&L)
+    drawdown:         float = 0.0      # current drawdown from peak capital
+
+
+class MultiStockAction(BaseModel):
+    stock:     str   = Field(..., description="NSE symbol to pick, or 'NONE' to skip")
+    direction: str   = Field(..., description="Bullish | Bearish | NONE")
+    conviction: float = Field(0.5, ge=0.0, le=1.0, description="Confidence 0–1")
 
 
 class StepResult(BaseModel):
-    observation: Optional[IndicatorsObservation]
-    reward: float
-    done: bool
-    info: Dict[str, Any]
+    observation: Optional[MultiStockObservation]
+    reward:      float
+    done:        bool
+    info:        Dict[str, Any]
 
 
 class ResetResult(BaseModel):
-    observation: IndicatorsObservation
-    info: Dict[str, Any]
+    observation: MultiStockObservation
+    info:        Dict[str, Any]
 
 
 class StateResult(BaseModel):
-    session_id: str
-    current_observation: Optional[IndicatorsObservation]
-    episodes_completed: int
-    current_task: Optional[str]
+    session_id:          str
+    current_observation: Optional[MultiStockObservation]
+    episodes_completed:  int
+    current_task:        Optional[str]
 
 
 class GraderRequest(BaseModel):
@@ -149,76 +261,123 @@ class GraderRequest(BaseModel):
     episode_results: List[Dict[str, Any]] = Field(
         ...,
         description=(
-            "List of episode dicts, each with keys: "
-            "'ground_truth' (str), 'predicted' (str), 'conviction' (float)"
+            "List of step dicts, each with keys: "
+            "'predicted' (str: Bullish|Bearish|NONE), "
+            "'ground_truth' (str: Bullish|Bearish|Neutral|N/A), "
+            "'conviction' (float)"
         ),
     )
 
 
 class GraderResult(BaseModel):
-    task_id: str
-    score: float = Field(..., gt=0.0, lt=1.0)
+    task_id:      str
+    score:        float = Field(..., gt=0.0, lt=1.0)
     num_episodes: int
-    breakdown: Dict[str, Any]
+    breakdown:    Dict[str, Any]
 
 
-# ─── Grader Logic ─────────────────────────────────────────────────────────────
+# ─── Grader logic ─────────────────────────────────────────────────────────────
+
 
 def _clamp_score(score: float) -> float:
-    """Clamp to strictly open interval (0, 1) as required by the validator."""
+    """Clamp to strictly open interval (0, 1) as required by the OpenEnv validator."""
     return round(max(0.001, min(0.999, score)), 4)
 
 
 def grade_task(task_id: str, episode_results: List[Dict[str, Any]]) -> GraderResult:
-    """Deterministic grader for all 3 tasks. Returns a score in (0, 1) exclusive."""
+    """
+    Score episode results for a given task.
+
+    Active steps = steps where predicted != 'NONE'.
+    Grader evaluates only active steps; passing silently does not harm the score
+    but forfeits the opportunity to score points.
+    """
     if not episode_results:
-        return GraderResult(task_id=task_id, score=0.001, num_episodes=0, breakdown={})
+        return GraderResult(
+            task_id=task_id, score=0.001, num_episodes=0,
+            breakdown={"note": "No results submitted"}
+        )
+
+    n_total  = len(episode_results)
+    active   = [
+        r for r in episode_results
+        if str(r.get("predicted", "NONE")).capitalize() != "None"
+        and str(r.get("predicted", "NONE")).upper() != "NONE"
+    ]
+    n_active = len(active)
+
+    if n_active == 0:
+        return GraderResult(
+            task_id=task_id, score=0.001, num_episodes=n_total,
+            breakdown={"note": "All steps passed (NONE)", "active_steps": 0}
+        )
 
-    n = len(episode_results)
+    participation_rate = n_active / n_total
 
     if task_id == "short_term_direction":
-        # Easy: simple accuracy
         correct = sum(
-            1 for r in episode_results
-            if r.get("predicted", "").capitalize() == r.get("ground_truth", "").capitalize()
+            1 for r in active
+            if str(r.get("predicted","")).capitalize() == str(r.get("ground_truth","")).capitalize()
         )
-        score = _clamp_score(correct / n)
-        breakdown = {"correct": correct, "total": n, "metric": "accuracy"}
+        raw_score = correct / n_active
+        score = _clamp_score(raw_score)
+        breakdown = {
+            "correct": correct,
+            "active_steps": n_active,
+            "total_steps": n_total,
+            "participation_rate": round(participation_rate, 3),
+            "metric": "directional_accuracy",
+        }
 
     elif task_id == "medium_term_direction":
-        # Medium: weighted accuracy (minority classes worth more)
         weighted_score = 0.0
         weighted_total = 0.0
-        for r in episode_results:
-            gt = r.get("ground_truth", "").capitalize()
-            pred = r.get("predicted", "").capitalize()
+        for r in active:
+            gt   = str(r.get("ground_truth", "")).capitalize()
+            pred = str(r.get("predicted", "")).capitalize()
             w = 1.5 if gt in ("Bearish", "Neutral") else 1.0
             weighted_total += w
             if pred == gt:
                 weighted_score += w
-        score = _clamp_score(min(1.0, weighted_score / weighted_total) if weighted_total > 0 else 0.0)
-        breakdown = {"weighted_score": round(weighted_score, 3), "weighted_total": round(weighted_total, 3), "metric": "weighted_accuracy"}
+        base = (weighted_score / weighted_total) if weighted_total > 0 else 0.0
+        # Participation bonus: agents that selectively engage vs. pass every step
+        adjusted = base * (0.9 + 0.1 * participation_rate)
+        score = _clamp_score(adjusted)
+        breakdown = {
+            "weighted_score": round(weighted_score, 3),
+            "weighted_total": round(weighted_total, 3),
+            "participation_rate": round(participation_rate, 3),
+            "active_steps": n_active,
+            "total_steps": n_total,
+            "metric": "weighted_accuracy_with_participation",
+        }
 
     elif task_id == "long_term_conviction":
-        # Hard: direction + conviction calibration
-        per_episode_scores = []
-        for r in episode_results:
-            gt = r.get("ground_truth", "").capitalize()
-            pred = r.get("predicted", "").capitalize()
+        per_step: List[float] = []
+        for r in active:
+            gt         = str(r.get("ground_truth", "")).capitalize()
+            pred       = str(r.get("predicted",    "")).capitalize()
             conviction = float(r.get("conviction", 0.5))
-            correct = pred == gt
+            correct    = (pred == gt)
             if correct and conviction >= 0.7:
-                per_episode_scores.append(1.0)
+                per_step.append(1.0)
             elif correct and conviction < 0.7:
-                per_episode_scores.append(0.5)
+                per_step.append(0.5)
             elif not correct and conviction >= 0.8:
-                per_episode_scores.append(-0.1)  # overconfident wrong
+                per_step.append(-0.1)   # overconfident wrong
             else:
-                per_episode_scores.append(0.0)
-        raw = sum(per_episode_scores) / n
-        # Normalize from [-0.1, 1.0] to (0, 1) exclusive
+                per_step.append(0.0)
+        raw = sum(per_step) / n_active
+        # Normalize [-0.1, 1.0] → (0, 1) exclusive
         score = _clamp_score((raw + 0.1) / 1.1)
-        breakdown = {"per_episode_scores": per_episode_scores, "raw_mean": round(raw, 4), "metric": "conviction_calibrated_accuracy"}
+        breakdown = {
+            "per_step_scores": per_step,
+            "raw_mean": round(raw, 4),
+            "active_steps": n_active,
+            "total_steps": n_total,
+            "participation_rate": round(participation_rate, 3),
+            "metric": "conviction_calibrated_accuracy",
+        }
 
     else:
         raise HTTPException(status_code=404, detail=f"Unknown task_id: {task_id}")
@@ -226,97 +385,310 @@ def grade_task(task_id: str, episode_results: List[Dict[str, Any]]) -> GraderRes
     return GraderResult(
         task_id=task_id,
         score=round(score, 4),
-        num_episodes=n,
+        num_episodes=n_total,
         breakdown=breakdown,
     )
 
 
 # ─── Session state ────────────────────────────────────────────────────────────
 
+
 class EnvSession:
+    """
+    Manages a single agent session across multi-stock episodes.
+
+    State per episode:
+      episode_steps    : list of n_steps dicts from build_multi_stock_episode
+      current_step_idx : 0-indexed pointer into episode_steps
+      signal_history   : accumulated RSI/momentum log shown in each observation
+      virtual_capital  : starts 1.0; updated by scaled alpha reward (Task 3 drawdown)
+
+    Reward (per step):
+      alpha  = chosen_stock_period_return − mean(all_3_stocks_period_return)
+      reward = alpha × direction_sign × conviction × REWARD_SCALE
+      NONE action → reward = 0.0
+    """
+
     def __init__(self, session_id: str, term: str = "medium"):
-        self.session_id = session_id
-        self.term = term.lower()
-        self.current_obs: Optional[IndicatorsObservation] = None
-        self.current_gt: Optional[str] = None
-        self.episodes_completed = 0
-        self.episode_history: List[Dict[str, Any]] = []
-        self.scenario_pool: List[Dict[str, str]] = []
-
-    def _get_scenario(self) -> Optional[Dict[str, str]]:
+        self.session_id     = session_id
+        self.term           = term.lower()
+        self.MAX_STEPS      = TASK_MAX_STEPS.get(self.term, 5)
+
+        # Episode data: list of step dicts from build_multi_stock_episode
+        self.episode_steps:     List[Dict[str, Any]] = []
+        self.current_step_idx:  int = 0
+        self.current_obs:       Optional[MultiStockObservation] = None
+        self.sector:            str = ""
+        self.symbols:           List[str] = []
+
+        self.episodes_completed: int = 0
+        self.episode_history:    List[Dict[str, Any]] = []
+        self.signal_history:     List[Dict[str, Any]] = []
+
+        # Portfolio: virtual capital and holding tracked across all tasks
+        self.virtual_capital: float = 1.0
+        self.peak_capital:    float = 1.0
+        self.current_holding: str   = "NONE"
+
+        # Lazy scenario pool
+        self.scenario_pool: List[Dict[str, Any]] = []
+
+    # ── helpers ──────────────────────────────────────────────────────────────
+
+    def _reset_episode_state(self) -> None:
+        self.episode_steps    = []
+        self.current_step_idx = 0
+        self.signal_history   = []
+        self.virtual_capital  = 1.0
+        self.peak_capital     = 1.0
+        self.current_holding  = "NONE"
+
+    def _get_scenario(self) -> Optional[Dict[str, Any]]:
+        """Pop a (sector, symbols, date) scenario from the pool, refilling if needed."""
         if not self.scenario_pool:
-            self.scenario_pool = generate_scenario_pool(
-                symbols=random.sample(NSE_UNIVERSE, min(20, len(NSE_UNIVERSE))),
-                start_date="2020-01-01",
-                end_date="2024-06-30",
-                term=self.term,
-                max_scenarios=2000,
+            sectors   = list(SECTOR_GROUPS.keys())
+            # Date range safely within available yfinance data
+            # For long term: leave 15×20=300 days of future data → cap at 2022-12-31
+            end_date = {
+                "short":  "2024-06-30",
+                "medium": "2023-12-31",
+                "long":   "2022-06-30",
+            }.get(self.term, "2023-12-31")
+
+            dates = (
+                pd.bdate_range("2020-01-01", end_date, freq="15B")
+                .strftime("%Y-%m-%d")
+                .tolist()
             )
+            for sector in sectors:
+                syms = SECTOR_GROUPS[sector]
+                sampled_dates = random.sample(dates, min(60, len(dates)))
+                for date in sampled_dates:
+                    self.scenario_pool.append({
+                        "sector":  sector,
+                        "symbols": random.sample(syms, min(3, len(syms))),
+                        "date":    date,
+                    })
             random.shuffle(self.scenario_pool)
+
         return self.scenario_pool.pop() if self.scenario_pool else None
 
+    def _build_obs(self, step_idx: int) -> MultiStockObservation:
+        """Construct MultiStockObservation for the current step."""
+        raw = self.episode_steps[step_idx]
+        stocks_dict: Dict[str, StockState] = {}
+        for s in raw["stocks"]:
+            od = s["obs_dict"]
+            stocks_dict[s["symbol"]] = StockState(
+                symbol             = od["symbol"],
+                current_price      = od["current_price"],
+                rsi_14             = od["rsi_14"],
+                rsi_trend          = od["rsi_trend"],
+                price_momentum_pct = od["price_momentum_pct"],
+                indicators         = od["indicators"],
+            )
+
+        drawdown = (
+            (self.peak_capital - self.virtual_capital) / self.peak_capital
+            if self.peak_capital > 0 else 0.0
+        )
+
+        return MultiStockObservation(
+            step             = step_idx + 1,          # 1-indexed
+            max_steps        = self.MAX_STEPS,
+            term             = self.term.upper(),
+            sector           = self.sector,
+            available_stocks = self.symbols,
+            stocks           = stocks_dict,
+            signal_history   = list(self.signal_history),
+            macro            = raw.get("macro"),
+            current_holding  = self.current_holding,
+            capital          = round(self.virtual_capital, 6),
+            drawdown         = round(drawdown, 4),
+        )
+
+    # ── core interface ────────────────────────────────────────────────────────
+
     def reset(self) -> Optional[ResetResult]:
+        """
+        Start a new episode. Picks a sector + 3 stocks, fetches data in 3 yfinance
+        calls (one per stock), builds the full episode in one pass.
+        """
+        self._reset_episode_state()
+
         for _ in range(10):
             sc = self._get_scenario()
             if sc is None:
                 return None
-            obs_dict = build_observation(sc["symbol"], sc["date"], self.term)
-            gt = compute_ground_truth(sc["symbol"], sc["date"], self.term)
-            if obs_dict is not None and gt is not None:
-                self.current_obs = IndicatorsObservation(**obs_dict)
-                self.current_gt = gt
+
+            steps = build_multi_stock_episode(
+                symbols      = sc["symbols"],
+                start_date   = sc["date"],
+                n_steps      = self.MAX_STEPS,
+                term         = self.term,
+                include_macro= (self.term == "long"),
+            )
+            if steps is not None:
+                self.episode_steps = steps
+                self.sector        = sc["sector"]
+                self.symbols       = sc["symbols"]
+                self.current_obs   = self._build_obs(0)
                 return ResetResult(
-                    observation=self.current_obs,
-                    info={"session_id": self.session_id, "term": self.term, "task_id": TERM_TO_TASK_ID.get(self.term)},
+                    observation = self.current_obs,
+                    info = {
+                        "session_id": self.session_id,
+                        "term":       self.term,
+                        "task_id":    TERM_TO_TASK_ID.get(self.term),
+                        "sector":     self.sector,
+                        "symbols":    self.symbols,
+                        "step":       0,
+                        "max_steps":  self.MAX_STEPS,
+                    },
                 )
         return None
 
-    def step(self, action: IndicatorsAction) -> StepResult:
-        if self.current_obs is None or self.current_gt is None:
+    def step(self, action: MultiStockAction) -> StepResult:
+        if not self.episode_steps or self.current_step_idx >= len(self.episode_steps):
             return StepResult(
                 observation=None, reward=0.0, done=True,
-                info={"error": "Call reset() first"}
+                info={"error": "Call reset() first"},
             )
-        correct = action.direction.strip().capitalize() == self.current_gt
-        reward = 1.0 if correct else 0.0
-
-        # Shaped reward: conviction calibration
-        if correct and action.conviction >= 0.6:
-            reward = 1.0 + 0.1 * (action.conviction - 0.6)
-        elif not correct and action.conviction >= 0.8:
-            reward = -0.1
 
-        reward = round(max(-0.1, min(1.1, reward)), 4)
+        raw         = self.episode_steps[self.current_step_idx]
+        stock_rows  = {s["symbol"]: s for s in raw["stocks"]}
+
+        # ── Parse action ──────────────────────────────────────────────────────
+        stock     = action.stock.strip().upper()
+        direction = action.direction.strip().capitalize()
+        if direction not in ("Bullish", "Bearish"):
+            direction = "NONE"
+        if stock == "NONE" or direction == "NONE":
+            stock, direction = "NONE", "NONE"
+        elif stock not in stock_rows:
+            # Invalid symbol → treat as NONE
+            stock, direction = "NONE", "NONE"
+
+        conviction = action.conviction
+
+        # ── Reward: market-neutral alpha + portfolio dynamics ─────────────────
+        all_returns = [s["actual_period_return"] for s in raw["stocks"]]
+        sector_avg  = mean(all_returns) if all_returns else 0.0
+
+        if direction == "NONE":
+            reward        = 0.0
+            alpha         = 0.0
+            chosen_return = 0.0
+            chosen_gt     = "N/A"
+            tx_cost       = 0.0
+            # NONE pass: keep existing holding, no capital change
+        else:
+            chosen_return  = stock_rows[stock]["actual_period_return"]
+            alpha          = chosen_return - sector_avg
+            direction_sign = 1 if direction == "Bullish" else -1
+            raw_reward     = alpha * direction_sign * conviction * REWARD_SCALE
+            chosen_gt      = stock_rows[stock]["gt"]
+
+            # Transaction cost: 0.1% of capital when switching holdings (all tasks)
+            tx_cost = 0.0
+            if (self.current_holding != "NONE"
+                    and stock != self.current_holding):
+                tx_cost = TRANSACTION_COST * self.virtual_capital
+
+            reward = round(max(-1.5, min(1.5, raw_reward)) - tx_cost, 4)
+
+            # Update capital — ALL tasks (not just long)
+            pnl_fraction         = alpha * direction_sign * conviction
+            self.virtual_capital = max(0.0, self.virtual_capital + pnl_fraction)
+            self.peak_capital    = max(self.peak_capital, self.virtual_capital)
+
+            # Update holding
+            self.current_holding = stock
+
+        correct   = (direction != "NONE") and (direction == chosen_gt)
+        drawdown_limit = DRAWDOWN_LIMITS.get(self.term, 0.10)
+        drawdown  = (
+            (self.peak_capital - self.virtual_capital) / self.peak_capital
+            if self.peak_capital > 0 else 0.0
+        )
 
-        info = {
-            "ground_truth": self.current_gt,
-            "predicted": action.direction,
-            "correct": correct,
-            "conviction": action.conviction,
-            "session_id": self.session_id,
-            "task_id": TERM_TO_TASK_ID.get(self.term),
-        }
+        # ── Update signal history ─────────────────────────────────────────────
+        self.signal_history.append({
+            "step":              self.current_step_idx + 1,
+            "picked_stock":      stock,
+            "direction":         direction,
+            "conviction":        conviction,
+            "ground_truth":      chosen_gt,
+            "correct":           correct if direction != "NONE" else None,
+            "alpha_pct":         round(alpha * 100, 3),
+            "reward":            reward,
+            # RSI snapshot of all stocks at this step
+            "rsi_snapshot": {
+                s["symbol"]: round(s["obs_dict"]["rsi_14"], 1)
+                for s in raw["stocks"]
+            },
+        })
 
-        # Store for grader
         self.episode_history.append({
-            "ground_truth": self.current_gt,
-            "predicted": action.direction.strip().capitalize(),
-            "conviction": action.conviction,
-            "reward": reward,
+            "step":         self.current_step_idx + 1,
+            "stock":        stock,
+            "direction":    direction,
+            "ground_truth": chosen_gt,
+            "correct":      correct,
+            "conviction":   conviction,
+            "reward":       reward,
+            "tx_cost":      round(tx_cost, 6),
+            "alpha_pct":    round(alpha * 100, 3),
+            "chosen_return_pct": round(chosen_return * 100, 4),
+            "sector_avg_pct":    round(sector_avg * 100, 4),
         })
 
-        self.episodes_completed += 1
-        prev_obs = self.current_obs
-        self.current_obs = None
-        self.current_gt = None
-        return StepResult(observation=prev_obs, reward=reward, done=True, info=info)
+        # ── Advance step ──────────────────────────────────────────────────────
+        self.current_step_idx += 1
+
+        early_stop = drawdown > drawdown_limit
+        done       = (self.current_step_idx >= self.MAX_STEPS) or early_stop
+
+        info = {
+            "step":              self.current_step_idx,
+            "max_steps":         self.MAX_STEPS,
+            "session_id":        self.session_id,
+            "task_id":           TERM_TO_TASK_ID.get(self.term),
+            "sector":            self.sector,
+            "chosen_stock":      stock,
+            "direction":         direction,
+            "ground_truth":      chosen_gt,
+            "correct":           correct,
+            "conviction":        conviction,
+            "tx_cost":           round(tx_cost, 6),
+            "current_holding":   self.current_holding,
+            "alpha_pct":         round(alpha * 100, 3),
+            "chosen_return_pct": round(chosen_return * 100, 4),
+            "sector_avg_pct":    round(sector_avg * 100, 4),
+            "virtual_capital":   round(self.virtual_capital, 6),
+            "drawdown":          round(drawdown, 4),
+            "drawdown_limit":    drawdown_limit,
+            "early_stop":        early_stop,
+        }
+
+        if not done:
+            self.current_obs = self._build_obs(self.current_step_idx)
+        else:
+            self.episodes_completed += 1
+            self.current_obs = None
+
+        return StepResult(
+            observation = self.current_obs,
+            reward      = reward,
+            done        = done,
+            info        = info,
+        )
 
     def state(self) -> StateResult:
         return StateResult(
-            session_id=self.session_id,
-            current_observation=self.current_obs,
-            episodes_completed=self.episodes_completed,
-            current_task=TERM_TO_TASK_ID.get(self.term),
+            session_id          = self.session_id,
+            current_observation = self.current_obs,
+            episodes_completed  = self.episodes_completed,
+            current_task        = TERM_TO_TASK_ID.get(self.term),
         )
 
 
@@ -324,6 +696,7 @@ class EnvSession:
 
 _sessions: Dict[str, EnvSession] = {}
 
+
 def _get_or_create(session_id: Optional[str] = None, term: str = "medium") -> EnvSession:
     if session_id is None:
         session_id = str(uuid.uuid4())
@@ -334,35 +707,52 @@ def _get_or_create(session_id: Optional[str] = None, term: str = "medium") -> En
 
 # ─── HTTP endpoints (OpenEnv standard) ───────────────────────────────────────
 
+
 @app.get("/health")
 def health():
-    return {"status": "ok", "env": "IndicatorsEnv", "version": "2.0.0", "tasks": len(TASKS)}
+    return {
+        "status":           "ok",
+        "env":              "IndicatorsEnv",
+        "version":          "4.1.0",
+        "tasks":            len(TASKS),
+        "episode_steps":    TASK_MAX_STEPS,
+        "step_spacing":     STEP_SPACING,
+        "sectors":          list(SECTOR_GROUPS.keys()),
+        "drawdown_limits":  DRAWDOWN_LIMITS,
+        "transaction_cost": TRANSACTION_COST,
+    }
 
 
 @app.get("/tasks")
 def get_tasks():
-    """Returns all task definitions and action schemas."""
+    """Returns all task definitions, action schemas, and episode lengths."""
     return {
-        "tasks": TASKS,
-        "total": len(TASKS),
+        "tasks":         TASKS,
+        "total":         len(TASKS),
         "action_schema": {
-            "direction": {"type": "string", "enum": ["Bullish", "Bearish", "Neutral"]},
-            "conviction": {"type": "number", "minimum": 0.0, "maximum": 1.0},
+            "stock":     {"type": "string", "description": "NSE symbol or 'NONE' to skip"},
+            "direction": {"type": "string", "enum": ["Bullish", "Bearish", "NONE"]},
+            "conviction":{"type": "number", "minimum": 0.0, "maximum": 1.0},
         },
+        "episode_steps":  TASK_MAX_STEPS,
+        "step_spacing":   STEP_SPACING,
     }
 
 
 @app.post("/reset", response_model=ResetResult)
 def reset(session_id: Optional[str] = None, term: str = "medium"):
-    sess = _get_or_create(session_id, term=term)
+    sess   = _get_or_create(session_id, term=term)
     result = sess.reset()
     if result is None:
-        raise HTTPException(status_code=503, detail="Could not fetch data. Retry.")
+        raise HTTPException(
+            status_code=503,
+            detail="Could not fetch market data. Retry or try a different term.",
+        )
     return result
 
 
 @app.post("/step", response_model=StepResult)
-def step(action: IndicatorsAction, session_id: Optional[str] = None):
+def step(action: MultiStockAction, session_id: Optional[str] = None):
     sess = _get_or_create(session_id)
     return sess.step(action)
 
@@ -377,67 +767,81 @@ def state(session_id: Optional[str] = None):
 def grader(request: GraderRequest):
     """
     Score a completed episode set against a specific task grader.
-    Provide episode_results as list of {ground_truth, predicted, conviction}.
-    Returns a 0.0–1.0 score.
+    episode_results: list of {predicted, ground_truth, conviction}.
+    Returns score in (0.0, 1.0).
     """
     if request.task_id not in TASK_BY_ID:
-        raise HTTPException(status_code=404, detail=f"Unknown task_id: {request.task_id}. Valid: {list(TASK_BY_ID.keys())}")
+        raise HTTPException(
+            status_code=404,
+            detail=f"Unknown task_id: {request.task_id}. "
+                   f"Valid: {list(TASK_BY_ID.keys())}",
+        )
     return grade_task(request.task_id, request.episode_results)
 
 
 @app.get("/baseline")
 def baseline():
     """
-    Trigger the built-in random baseline agent across all 3 tasks (10 episodes each).
-    Returns reproducible baseline scores.
+    Run the built-in random baseline across all 3 tasks (10 episodes each).
+    Expected: overall_mean ≈ 0.33 (random agent on 3-class directional problem).
+    short  : 5×10=50  steps
+    medium : 10×10=100 steps
+    long   : ≤15×10=150 steps (may be fewer with early termination)
     """
     results = {}
-    random.seed(42)  # Deterministic
+    random.seed(42)
 
     for task in TASKS:
-        term = task["term"]
+        term    = task["term"]
         task_id = task["id"]
-        session_id = f"baseline-{task_id}"
-
-        # Fresh session
-        sess = EnvSession(session_id=session_id, term=term)
-        episode_results = []
+        sess    = EnvSession(session_id=f"baseline-{task_id}", term=term)
+        episode_results: List[Dict[str, Any]] = []
 
         for _ in range(10):
             reset_result = sess.reset()
             if reset_result is None:
                 continue
-            # Random agent
-            action = IndicatorsAction(
-                direction=random.choice(["Bullish", "Bearish", "Neutral"]),
-                conviction=round(random.uniform(0.3, 0.9), 2),
-            )
-            step_result = sess.step(action)
-            episode_results.append({
-                "ground_truth": step_result.info.get("ground_truth", ""),
-                "predicted": action.direction,
-                "conviction": action.conviction,
-            })
+            done = False
+            while not done:
+                # Random agent: random stock, random direction (including NONE)
+                available = reset_result.observation.available_stocks
+                if hasattr(sess.current_obs, "available_stocks") and sess.current_obs:
+                    available = sess.current_obs.available_stocks
+                chosen_stock = random.choice(available + ["NONE"])
+                direction    = "NONE" if chosen_stock == "NONE" else random.choice(["Bullish", "Bearish"])
+                action = MultiStockAction(
+                    stock      = chosen_stock,
+                    direction  = direction,
+                    conviction = round(random.uniform(0.3, 0.9), 2),
+                )
+                step_result = sess.step(action)
+                episode_results.append({
+                    "predicted":    direction,
+                    "ground_truth": step_result.info.get("ground_truth", "N/A"),
+                    "conviction":   action.conviction,
+                })
+                done = step_result.done
 
         grader_result = grade_task(task_id, episode_results)
         results[task_id] = {
-            "task_name": task["name"],
-            "difficulty": task["difficulty"],
-            "score": grader_result.score,
+            "task_name":    task["name"],
+            "difficulty":   task["difficulty"],
+            "score":        grader_result.score,
             "num_episodes": grader_result.num_episodes,
-            "breakdown": grader_result.breakdown,
+            "breakdown":    grader_result.breakdown,
         }
 
     return {
-        "agent": "random_baseline",
-        "seed": 42,
-        "tasks": results,
+        "agent":        "random_baseline",
+        "seed":         42,
+        "tasks":        results,
         "overall_mean": round(sum(r["score"] for r in results.values()) / len(results), 4),
     }
 
 
 # ─── WebSocket endpoint (OpenEnv /ws) ────────────────────────────────────────
 
+
 @app.websocket("/ws")
 async def websocket_env(ws: WebSocket, session_id: Optional[str] = None, term: str = "medium"):
     await ws.accept()
@@ -457,7 +861,7 @@ async def websocket_env(ws: WebSocket, session_id: Optional[str] = None, term: s
                     await ws.send_text(result.model_dump_json())
 
             elif method == "step":
-                action = IndicatorsAction(**msg.get("action", {}))
+                action = MultiStockAction(**msg.get("action", {}))
                 result = sess.step(action)
                 await ws.send_text(result.model_dump_json())
 
@@ -477,5 +881,7 @@ async def websocket_env(ws: WebSocket, session_id: Optional[str] = None, term: s
 
 # ─── Entry point ─────────────────────────────────────────────────────────────
 
+
 if __name__ == "__main__":
-    uvicorn.run("indicators_env:app", host="0.0.0.0", port=7860, reload=False, workers=2)
+    # workers=1 required: _sessions is an in-memory dict, not shared across processes
+    uvicorn.run("indicators_env:app", host="0.0.0.0", port=7860, reload=False, workers=1)

inference.py

@@ -2,10 +2,17 @@
 inference.py — Mandatory entry point for the Meta × PyTorch Hackathon automated evaluation.
 
 Requirements:
-1.  Reads API_BASE_URL, MODEL_NAME, and HF_TOKEN from environment variables.
-2.  Uses the OpenAI Client for all LLM calls.
-3.  Stdout follows strict [START], [STEP], and [END] logging format.
-4.  Runtime < 20 min on 2 vCPU / 8 GB RAM.
+1. Reads API_BASE_URL, MODEL_NAME, and HF_TOKEN from environment variables.
+2. Uses the OpenAI Client for all LLM calls.
+3. Stdout follows strict [START], [STEP], and [END] logging format.
+4. Runtime < 20 min on 2 vCPU / 8 GB RAM.
+
+IndicatorsEnv v4.1 — Multi-stock Portfolio MDP:
+  At each step the agent observes 3 stocks from the same NSE sector.
+  It picks ONE stock and declares Bullish/Bearish, or passes with NONE.
+  Reward = (chosen_stock_return − sector_avg) × direction × conviction × 50
+  Market-neutral: random policy earns ~0; skilled policy earns positive alpha.
+  Tasks: short (5 steps, daily), medium (10 steps, weekly), long (15 steps, monthly).
 """
 
 import os
@@ -13,144 +20,362 @@ import re
 import json
 import requests
 import argparse
+from typing import Any, Dict, List, Optional, Tuple
 from openai import OpenAI
 
-# --- Configuration (Mandatory per Hackathon Spec) ---
-API_BASE_URL = os.environ.get("API_BASE_URL", "https://api-inference.huggingface.co/v1/")
-MODEL_NAME = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-7B-Instruct")
-HF_TOKEN = os.environ.get("HF_TOKEN")
+# ── Configuration (Mandatory per Hackathon Spec) ────────────────────────────
+API_BASE_URL = os.environ.get("API_BASE_URL", "https://router.huggingface.co/v1/")
+MODEL_NAME   = os.environ.get("MODEL_NAME",   "meta-llama/Llama-3.2-1B-Instruct")
+HF_TOKEN     = os.environ.get("HF_TOKEN")
+
+TASKS     = ["short_term_direction", "medium_term_direction", "long_term_conviction"]
+ENV_TERMS = {
+    "short_term_direction": "short",
+    "medium_term_direction": "medium",
+    "long_term_conviction": "long",
+}
+
 
-TASKS = ["short_term_direction", "medium_term_direction", "long_term_conviction"]
-ENV_TERMS = {"short_term_direction": "short", "medium_term_direction": "medium", "long_term_conviction": "long"}
+# ── Action parser ────────────────────────────────────────────────────────────
 
+def _parse_action(
+    text: str,
+    available_stocks: List[str],
+) -> Tuple[str, str, float]:
+    """
+    Parse LLM response into (stock, direction, conviction).
 
-def _parse_direction_and_conviction(text: str):
-    """Extract direction and conviction from JSON-formatted LLM response."""
-    # Clean possible markdown artifacts
+    Expected JSON:
+      {"stock": "HDFCBANK", "direction": "Bullish", "conviction": 0.8}
+    or to pass:
+      {"stock": "NONE", "direction": "NONE", "conviction": 0.0}
+
+    Falls back gracefully on malformed output.
+    """
     text = re.sub(r"```(?:json)?", "", text).strip().rstrip("`").strip()
+
     try:
-        # Look for JSON object
         match = re.search(r"\{[^}]+\}", text, re.DOTALL)
         if match:
             obj = json.loads(match.group(0))
-            d = str(obj.get("direction", "Neutral")).strip().capitalize()
-            c = float(obj.get("conviction", 0.5))
-            if d not in ("Bullish", "Bearish", "Neutral"): d = "Neutral"
-            return d, c
+            stock     = str(obj.get("stock", "NONE")).strip().upper()
+            direction = str(obj.get("direction", "NONE")).strip().capitalize()
+            conviction = float(obj.get("conviction", 0.5))
+            conviction = max(0.0, min(1.0, conviction))
+
+            # Validate stock
+            if stock not in available_stocks and stock != "NONE":
+                # Try case-insensitive match
+                upper_map = {s.upper(): s for s in available_stocks}
+                stock = upper_map.get(stock, "NONE")
+
+            # Validate direction
+            if direction not in ("Bullish", "Bearish", "None"):
+                direction = "NONE"
+            if stock == "NONE":
+                direction = "NONE"
+
+            return stock, direction, conviction
     except Exception:
         pass
-    
-    # Simple fallback: keyword search
+
+    # ── Fallback: keyword search ─────────────────────────────────────────────
     lower = text.lower()
-    if "bullish" in lower: return "Bullish", 0.7
-    if "bearish" in lower: return "Bearish", 0.7
-    return "Neutral", 0.5
 
+    # Look for stock mention
+    found_stock = "NONE"
+    for sym in available_stocks:
+        if sym.lower() in lower:
+            found_stock = sym
+            break
+
+    # Look for direction
+    if "bullish" in lower:
+        direction = "Bullish"
+    elif "bearish" in lower:
+        direction = "Bearish"
+    elif "none" in lower or "pass" in lower or "skip" in lower:
+        direction = "NONE"
+        found_stock = "NONE"
+    else:
+        direction = "NONE"
+        found_stock = "NONE"
+
+    return found_stock, direction, 0.6
+
+
+# ── Prompt builder ───────────────────────────────────────────────────────────
+
+def _build_prompt(
+    obs: Dict[str, Any],
+    step_num: int,
+    max_steps: int,
+    term: str,
+    task_id: str,
+    signal_history: List[Dict[str, Any]],
+    prev_info: Optional[Dict[str, Any]] = None,
+) -> Tuple[str, str]:
+    """
+    Build system + user prompt for multi-stock relative alpha prediction.
+    """
+    system_prompt = (
+        "You are a quantitative portfolio manager specializing in NSE (India) equities. "
+        "At each step you observe 3 stocks from the SAME sector. "
+        "Your job: identify which stock has the strongest relative momentum signal "
+        "and bet on its direction vs. the sector average. "
+        "Reward = (your stock's return − sector average) × direction × conviction × 50. "
+        "Market-neutral: picking randomly earns ~0 reward — you profit ONLY by selecting "
+        "the outperformer correctly. "
+        "Switching your held position to a different stock costs 0.1% of capital in transaction cost. "
+        "If you are already in a strong position, staying in the same stock avoids switching cost. "
+        "Pass with NONE to hold your position without incurring a switch cost. "
+        'Respond ONLY with JSON: {"stock": "<SYMBOL>", "direction": "Bullish"|"Bearish", '
+        '"conviction": <0.0-1.0>} or {"stock": "NONE", "direction": "NONE", "conviction": 0.0}'
+    )
+
+    sector      = obs.get("sector", "unknown")
+    available   = obs.get("available_stocks", [])
+    stocks_data = obs.get("stocks", {})
+    macro       = obs.get("macro")
+
+    # ── Format stock snapshots ────────────────────────────────────────────────
+    stock_lines = []
+    for sym in available:
+        s = stocks_data.get(sym, {})
+        rsi       = s.get("rsi_14", "N/A")
+        rsi_trend = s.get("rsi_trend", "?")
+        momentum  = s.get("price_momentum_pct", 0.0)
+        price     = s.get("current_price", "N/A")
+        # Key indicator signals
+        inds      = s.get("indicators", {})
+        macd_sig  = inds.get("macd", {}).get("signal", "?")
+        ma_sig    = inds.get("moving_averages", {}).get("signal", "?")
+        adx_str   = inds.get("adx", {}).get("trend_strength", "?")
+
+        stock_lines.append(
+            f"  {sym}: price={price}  RSI={rsi}({rsi_trend})  "
+            f"momentum={momentum:+.1f}%  MACD={macd_sig}  MA={ma_sig}  ADX={adx_str}"
+        )
+
+    stocks_str = "\n".join(stock_lines)
+
+    # ── Signal history (last 3 steps for context) ────────────────────────────
+    history_str = ""
+    if signal_history:
+        recent = signal_history[-3:]
+        lines = []
+        for h in recent:
+            picked = h.get("picked_stock", "NONE")
+            dirn   = h.get("direction", "NONE")
+            gt     = h.get("ground_truth", "?")
+            alpha  = h.get("alpha_pct", 0.0)
+            rew    = h.get("reward", 0.0)
+            if dirn != "NONE":
+                correct_str = "✓" if h.get("correct") else "✗"
+                lines.append(
+                    f"  Step {h['step']}: picked {picked} → {dirn} "
+                    f"({correct_str} GT={gt}  alpha={alpha:+.2f}%  reward={rew:.3f})"
+                )
+            else:
+                lines.append(f"  Step {h['step']}: NONE (passed)")
+        history_str = "Signal history:\n" + "\n".join(lines) + "\n"
+
+    # ── Macro context (Task 3) ────────────────────────────────────────────────
+    macro_str = ""
+    if macro:
+        macro_str = (
+            f"Macro: NIFTY50={macro.get('nifty_trend','?')} "
+            f"(20d return: {macro.get('nifty_return_20d',0):.1f}%)  "
+            f"regime={macro.get('market_regime','?')}\n"
+        )
+
+    # ── Previous step feedback ────────────────────────────────────────────────
+    prev_str = ""
+    if prev_info and prev_info.get("chosen_stock") != "N/A":
+        prev_str = (
+            f"Last step: picked {prev_info.get('chosen_stock','?')} → "
+            f"{prev_info.get('direction','?')} "
+            f"(GT={prev_info.get('ground_truth','?')}  "
+            f"alpha={prev_info.get('alpha_pct',0):+.2f}%  "
+            f"reward={prev_info.get('reward',0):.3f})\n"
+        )
+
+    # ── Portfolio context (v4.1) ──────────────────────────────────────────────
+    holding  = obs.get("current_holding", "NONE") if isinstance(obs, dict) else "NONE"
+    capital  = obs.get("capital", 1.0)             if isinstance(obs, dict) else 1.0
+    drawdown = obs.get("drawdown", 0.0)            if isinstance(obs, dict) else 0.0
+    if holding != "NONE":
+        tx_note = f"(switching from {holding} costs ~{capital * 0.001:.4f} tx_cost)"
+    else:
+        tx_note = "(no current holding — first pick is free)"
+    portfolio_str = (
+        f"Portfolio: capital={capital:.4f}  holding={holding}  "
+        f"drawdown={drawdown:.1%}  {tx_note}\n"
+    )
 
-def run_evaluation(env_url, n_episodes):
-    # Initialize OpenAI Client
+    user_prompt = (
+        f"[Step {step_num}/{max_steps}] Sector: {sector.upper()} | Task: {task_id} | "
+        f"Term: {term.upper()}\n\n"
+        f"Stocks (same sector):\n{stocks_str}\n\n"
+        f"{history_str}"
+        f"{macro_str}"
+        f"{portfolio_str}"
+        f"{prev_str}"
+        f"Pick the stock with the strongest relative momentum signal.\n"
+        f"Available: {available} or NONE to skip.\n\n"
+        f"Which stock has the best setup? Respond with JSON."
+    )
+
+    return system_prompt, user_prompt
+
+
+# ── Main evaluation loop ────────────────────────────────────────────────────
+
+def run_evaluation(env_url: str, n_episodes: int) -> None:
     client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
 
     for task_id in TASKS:
         term = ENV_TERMS[task_id]
-        episode_results = []
-        
-        # [START] Log
+        episode_results: List[Dict[str, Any]] = []
+
         print(f"[START] task={task_id} env=IndicatorsEnv model={MODEL_NAME}", flush=True)
 
         for i in range(n_episodes):
             try:
-                # 1. Reset Environment (use params, not json body)
-                reset_resp = requests.post(f"{env_url}/reset", params={"term": term}, timeout=30)
+                # 1. Reset — 90s timeout (multi-stock fetches 3 yfinance calls)
+                reset_resp = requests.post(
+                    f"{env_url}/reset",
+                    params={"term": term},
+                    timeout=90,
+                )
                 if reset_resp.status_code != 200:
+                    print(f"[ERR] episode={i+1} reset HTTP {reset_resp.status_code}", flush=True)
                     continue
-                
-                data = reset_resp.json()
-                obs = data.get("observation", {})
-                state = obs
+
+                data       = reset_resp.json()
+                obs        = data.get("observation", {})
                 session_id = data.get("info", {}).get("session_id", "")
-                symbol = obs.get("symbol", "N/A")
-                date = obs.get("date", "N/A")
-
-                # 2. Build Prompt
-                system_prompt = "You are a quantitative analyst. Given technical indicators, predict the directional price move as Bullish, Bearish, or Neutral. Provide reasoning then output a JSON with 'reasoning', 'direction', and 'conviction' (0.0-1.0)."
-                user_prompt = f"Technical Indicators for {symbol} on {date}: {state}\n\nPredict the directional move."
-
-                # 3. Call LLM (using mandatory OpenAI Client)
-                chat_completion = client.chat.completions.create(
-                    model=MODEL_NAME,
-                    messages=[
-                        {"role": "system", "content": system_prompt},
-                        {"role": "user", "content": user_prompt},
-                    ],
-                    max_tokens=256,
-                    temperature=0.7
-                )
-                response = chat_completion.choices[0].message.content
-                direction, conviction = _parse_direction_and_conviction(response)
-
-                # 4. Step Environment (use params for session_id, json body for action)
-                step_resp = requests.post(
-                    f"{env_url}/step", 
-                    params={"session_id": session_id},
-                    json={"direction": direction, "conviction": conviction}, 
-                    timeout=30
-                )
-                if step_resp.status_code != 200:
-                    continue
-                
-                step_data = step_resp.json()
-                reward = step_data.get("reward", 0.0)
-                ground_truth = step_data.get("info", {}).get("ground_truth", "N/A")
-
-                # Track result for grader
-                episode_results.append({
-                    "episode": i + 1,
-                    "direction": direction,
-                    "conviction": conviction,
-                    "reward": reward,
-                    "ground_truth": ground_truth
-                })
-
-                # [STEP] Log
-                print(f"[STEP] step={i+1} action={direction} reward={reward:.4f} done=True error=None", flush=True)
+                max_steps  = data.get("info", {}).get("max_steps", 5)
+                available  = obs.get("available_stocks", [])
+                done       = False
+                step_num   = 0
+                prev_info: Optional[Dict[str, Any]] = None
+
+                # 2. Multi-step loop — one episode = max_steps steps
+                while not done:
+                    step_num += 1
+
+                    # Signal history from observation (accumulated across steps)
+                    signal_history = obs.get("signal_history", []) if isinstance(obs, dict) else []
+                    available      = obs.get("available_stocks", available) if isinstance(obs, dict) else available
+
+                    # Build prompts
+                    system_prompt, user_prompt = _build_prompt(
+                        obs           = obs if isinstance(obs, dict) else {},
+                        step_num      = step_num,
+                        max_steps     = max_steps,
+                        term          = term,
+                        task_id       = task_id,
+                        signal_history= signal_history,
+                        prev_info     = prev_info,
+                    )
+
+                    # 3. LLM call
+                    chat_completion = client.chat.completions.create(
+                        model    = MODEL_NAME,
+                        messages = [
+                            {"role": "system", "content": system_prompt},
+                            {"role": "user",   "content": user_prompt},
+                        ],
+                        max_tokens  = 128,
+                        temperature = 0.7,
+                    )
+                    response  = chat_completion.choices[0].message.content or ""
+                    stock, direction, conviction = _parse_action(response, available)
+
+                    # 4. Step environment
+                    step_resp = requests.post(
+                        f"{env_url}/step",
+                        params={"session_id": session_id},
+                        json={
+                            "stock":     stock,
+                            "direction": direction,
+                            "conviction": conviction,
+                        },
+                        timeout=30,
+                    )
+                    if step_resp.status_code != 200:
+                        print(f"[ERR] episode={i+1} step={step_num} HTTP {step_resp.status_code}", flush=True)
+                        break
+
+                    step_data    = step_resp.json()
+                    reward       = step_data.get("reward", 0.0)
+                    done         = step_data.get("done", True)
+                    step_info    = step_data.get("info", {})
+                    ground_truth = step_info.get("ground_truth", "N/A")
+                    obs          = step_data.get("observation") or obs
+
+                    # Store info for next step's prompt
+                    prev_info = {
+                        "chosen_stock":      step_info.get("chosen_stock", stock),
+                        "direction":         direction,
+                        "ground_truth":      ground_truth,
+                        "alpha_pct":         step_info.get("alpha_pct", 0.0),
+                        "reward":            reward,
+                    }
+
+                    # Track for grader — key MUST be "predicted"
+                    episode_results.append({
+                        "predicted":    direction,
+                        "conviction":   conviction,
+                        "reward":       reward,
+                        "ground_truth": ground_truth,
+                    })
+
+                    # [STEP] log
+                    print(
+                        f"[STEP] step={step_num} stock={stock} action={direction} "
+                        f"reward={reward:.4f} done={done} error=None",
+                        flush=True,
+                    )
 
             except Exception as e:
-                # Never crash the entire run
+                print(f"[ERR] episode={i+1} {type(e).__name__}: {e}", flush=True)
                 continue
 
-        # 5. Finalize Task (Call Grader)
+        # 5. Finalize task — call grader
         try:
-            # Note: We provide the grader with the episode results for that task
-            # In OpenEnv, the grader is usually a separate endpoint
             grader_resp = requests.post(
-                f"{env_url}/grader", 
+                f"{env_url}/grader",
                 json={"task_id": task_id, "episode_results": episode_results},
-                timeout=60
+                timeout=60,
+            )
+            final_score = (
+                grader_resp.json().get("score", 0.0)
+                if grader_resp.status_code == 200 else 0.0
             )
-            final_score = grader_resp.json().get("score", 0.0) if grader_resp.status_code == 200 else 0.0
-            
-            # [END] Log
-            rewards_list = [ep["reward"] for ep in episode_results]
-            steps_taken = len(episode_results)
-            success = final_score > 0.0
-            print(f"[END] success={success} steps={steps_taken} score={final_score:.4f} rewards={rewards_list}", flush=True)
-        except:
-            rewards_list = [ep["reward"] for ep in episode_results]
-            steps_taken = len(episode_results)
-            print(f"[END] success=False steps={steps_taken} score=0.0000 rewards={rewards_list}", flush=True)
+        except Exception as e:
+            print(f"[ERR] grader: {type(e).__name__}: {e}", flush=True)
+            final_score = 0.0
+
+        rewards_list = [ep["reward"] for ep in episode_results]
+        steps_taken  = len(episode_results)
+        success      = final_score > 0.0
+        print(
+            f"[END] success={success} steps={steps_taken} score={final_score:.4f} "
+            f"rewards={rewards_list}",
+            flush=True,
+        )
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Strict Compliance OpenEnv Inference Agent")
-    parser.add_argument("--env_url", default="http://localhost:7860", help="URL of the IndicatorsEnv server")
-    parser.add_argument("--n_episodes", type=int, default=5, help="Number of episodes per task (Total 15)")
-    
+    parser = argparse.ArgumentParser(description="IndicatorsEnv v4.0 Inference Agent")
+    parser.add_argument("--env_url",    default="http://localhost:7860", help="IndicatorsEnv server URL")
+    parser.add_argument("--n_episodes", type=int, default=5,            help="Episodes per task")
     args = parser.parse_args()
-    
-    # Mandatory Check: Environment Variables
+
     if not HF_TOKEN:
-        print("❌ Error: HF_TOKEN must be set as an environment variable.")
+        print("Error: HF_TOKEN must be set as an environment variable.")
         exit(1)
 
     run_evaluation(args.env_url, args.n_episodes)