Sumanth2377's picture
Convert models to Pydantic and enhance Blog writeup
16f4457
|
Raw
History Blame Contribute Delete
4.71 kB
# Building WeddingPlannerEnv: A Chaos-Aware RL Benchmark for Indian Weddings
*A submission for the AR'26 Meta OpenEnv Hackathon*
When we looked at the current landscape of Reinforcement Learning environments for LLMs, we noticed a pattern: they are either highly abstract (like GridWorld), perfectly deterministic (like Tic-Tac-Toe), or purely physical (like Mujoco). Real-world planning—especially in culturally nuanced, high-stakes contexts—is vastly underrepresented.
That's why we built **WeddingPlannerEnv**: a native OpenEnv-compatible simulation that challenges an AI to plan a 3-day Indian wedding. But we didn't just build a scheduling app. We built an adversarial, stochastic simulation designed to break brittle agent logic.
## The Problem: Long-Horizon Constraints
Planning an Indian wedding is a multidimensional, non-linear optimization problem. The agent must manage:
- A strict budget (often running into tens of lakhs).
- A sprawling, multi-day schedule without any overlapping ceremonies.
- A requirement to book over a dozen interdependent vendors (venues, caterers, decorators, priests, etc.).
But the hardest constraint isn't money; it's **time**. Ceremonies must align perfectly with *Muhurat* (auspicious astrological timing windows). If the venue is booked during *Rahu Kaal* (inauspicious time), the entire plan is invalid.
## Technical Architecture
To make this rigorous enough for research-grade RL training, we engineered the environment from the ground up using FastAPI and Pydantic:
### 1. Strongly Typed State & Action Spaces
We completely avoided loose JSON parsing. Every observation and action is strictly governed by Pydantic `BaseModel` schemas. The environment enforces strict payload validation at the `POST /step` boundary. If an LLM hallucinates an action type or a malformed date string, the Pydantic validator catches it and immediately issues a heavy negative reward, forcing the policy model to quickly learn exact schema compliance.
### 2. Composable Rubrics & Dense Rewards
Instead of a sparse `0/1` reward at the end of the episode, we built a highly dense, multi-dimensional reward function composited from specialized modules:
- **Coverage Reward (35%)**: Calculates the percentage of mandatory vendors booked.
- **Budget Reward (25%)**: Exponentially decays as the agent nears bankruptcy.
- **Muhurat Reward (20%)**: Uses intersection over union (IoU) to evaluate how well booked time slots overlap with astrological windows.
- **Chaos Resolution (10%)**: Rewards the agent for quickly neutralizing active stochastic events.
- **Guest UX (10%)**: Penalizes scheduling gaps or missing caterers that would cause family stress.
### 3. The Chaos Engine (Stochastic Adversary)
Real weddings don't go according to plan. Vendors cancel. Prices surge. Traffic delays the priest.
To simulate this, we built the **Chaos Engine**. It is a stochastic adversarial system that fires random `ChaosEvent` triggers mid-episode. A vendor might suddenly cancel, injecting an active conflict into the state observation. The agent must pause its planning, parse the `active_chaos` array, and use emergency actions like `resolve_conflict` or `negotiate` to recover its deposit and find a replacement within the remaining budget constraints.
## Curriculum Learning with GRPO
We fine-tuned `Qwen2.5-7B-Instruct` using **Group Relative Policy Optimization (GRPO)** and a 3-stage curriculum on Lightning AI:
1. **Stage 1 (Easy):** The agent learns the strict Pydantic output schema and basic vendor booking mechanics, quickly escaping the -500 penalty floor.
2. **Stage 2 (Medium):** Budgets tighten. The agent initially dips into negative rewards as it hits limits, but quickly adapts its policy to seek cheaper vendors and respect Muhurat windows.
3. **Stage 3 (Hard):** The Chaos Engine activates. The agent learns to detect conflicts and recover gracefully without falling into an infinite action loop.
## The Results
In a baseline zero-shot test, standard LLMs fail entirely—they hallucinate years (e.g., "2224"), blow the budget, or output invalid keys.
After our curriculum training, our agent achieved a stable **positive average score of ~24/100** on the Hard difficulty setting. While a score of 100 represents theoretical perfection, a recovery from -500 to a consistent +24 means the agent successfully mastered constraint avoidance under chaos. It learned the rules of a highly adversarial simulation from scratch.
We believe WeddingPlannerEnv represents a step forward in evaluating LLMs for complex, long-horizon planning tasks in the real world. You can try the live environment via our HuggingFace Space API and test the model yourself!