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!