Yalpha/AGRITECH-META

1"""2Baseline agents for AgriDecisionEnv v3.3Three policies: random, rule-based (state-adaptive), greedy.4"""5import random as _random6from models import Action, Observation7 8 9def random_policy(obs, rng=None):10    """Uniformly random action."""11    r = rng or _random.Random()12    return Action(13        crop       = r.choice(["rice", "wheat", "none"]),14        fertilizer = round(r.uniform(0.0, 1.0), 2),15        irrigation = round(r.uniform(0.0, 1.0), 2),16    )17 18 19def rule_based_policy(obs):20    """21    Adaptive heuristic:22    - Avoids monocropping23    - Adjusts inputs to nitrogen/moisture deficit24    - Reduces spend if budget is tight25    - Guards groundwater26    - Skips planting if nitrogen critically low27    """28    budget_factor = min(1.0, max(0.0, float(obs.budget) / 120.0))29 30    # Crop selection — monocrop avoidance checked BEFORE weather branch31    if float(obs.nitrogen) < 0.25:32        crop = "none"33    elif obs.last_crop == "wheat":34        # Must rotate away from wheat; use rice unless drought forces otherwise35        crop = "none" if obs.weather == "drought" else "rice"36    elif obs.last_crop == "rice":37        crop = "wheat"38    elif obs.weather == "drought":39        crop = "wheat"40    else:41        crop = "wheat"42 43    # Fertilizer: fill nitrogen deficit, stay under overuse threshold44    n_deficit  = max(0.0, 0.60 - float(obs.nitrogen))45    fertilizer = round(min(0.55, n_deficit * 1.6) * budget_factor, 2)46 47    # Irrigation: fill moisture deficit, stay under threshold48    m_deficit  = max(0.0, 0.50 - float(obs.moisture))49    irrigation = round(min(0.60, m_deficit * 1.8) * budget_factor, 2)50 51    # Groundwater guard — tiered caps to prevent depletion cascade52    gw = float(obs.groundwater)53    if gw < 0.15:54        irrigation = min(irrigation, 0.05)   # critical55    elif gw < 0.25:56        irrigation = min(irrigation, 0.12)   # low57    elif gw < 0.40:58        irrigation = min(irrigation, 0.25)   # moderate59 60    return Action(crop=crop, fertilizer=fertilizer, irrigation=irrigation)61 62 63def greedy_policy(obs):64    """65    Always plants rice at high inputs for maximum immediate yield.66    Ignores long-term consequences — useful as single-step upper bound.67    """68    fertilizer = 0.7069    irrigation = 0.7070    if float(obs.budget) < 40.0:71        fertilizer, irrigation = 0.30, 0.3072    return Action(crop="rice", fertilizer=fertilizer, irrigation=irrigation)73 74 75def run_episode(policy_fn, scenario="default", seed=42):76    """Run a full 5-step episode and return summary."""77    import sys, os78    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))79    from env import AgriEnv80 81    env = AgriEnv(scenario=scenario, seed=seed)82    obs = env.reset()83    rewards = []84 85    for _ in range(5):86        action = policy_fn(obs)87        obs, reward, done, info = env.step(action)88        rewards.append(reward)89        if done:90            break91 92    return {93        "total_reward": round(sum(rewards), 4),94        "avg_reward":   round(sum(rewards) / len(rewards), 4),95        "steps":        len(rewards),96        "final_budget": obs.budget,97        "final_soil":   obs.soil_quality,98        "rewards":      rewards,99    }100 101 102if __name__ == "__main__":103    import sys, os, random104    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))105 106    for name, fn in [107        ("random",     lambda o: random_policy(o, random.Random(0))),108        ("rule_based", rule_based_policy),109        ("greedy",     greedy_policy),110    ]:111        r = run_episode(fn)112        print(f"[{name:10s}] avg={r['avg_reward']:.4f}  soil={r['final_soil']:.4f}  budget={r['final_budget']:.1f}")