docs/implementation_plan.md

Round 2: Hierarchical Multi-Agent Customer Support Environment

Goal

Upgrade the existing single-agent Customer Support RL Environment into a 3-level hierarchical multi-agent system that showcases:

• Multi-Agent Interactions (cooperation, oversight, hierarchy)
• Professional Tasks (enterprise workflows with Indian context)
• Instruction Following (role-specific instructions per agent)
• Self-Improvement (RL training with dense, non-gamable rewards)

The system will have:

1. Level 1 — Support Agent: Handles initial customer interaction
2. Level 2 — Supervisor Agent: Monitors Support Agent, gives feedback, enforces policy
3. Level 3 — Manager Agent: Handles escalations, resolves conflicts, makes final decisions

---

User Review Required

API Provider Decision: The plan uses your existing NVIDIA NIM keys (via OpenAI-compatible API) for both the customer simulator and LLM-as-Judge. If you want to use a different provider (OpenRouter, Bedrock, Claude), let me know and I'll adjust the API calls.

Backward Compatibility: The Round 1 single-agent /reset and /step endpoints will be preserved alongside the new hierarchical endpoints. Old inference.py tasks (easy/medium/hard/nightmare) still work unchanged. New hierarchical tasks are added separately.

LLM Cost During Development: The customer simulator and LLM-as-Judge both make API calls during /step. For training, you'll want to batch these or switch to a local model. During dev/testing, each episode will use ~5-15 API calls for the customer sim and ~3-8 for the judge. With 3 NVIDIA keys on failover, this should be manageable.

---

Open Questions

Q1: Should the Manager agent be LLM-driven during inference (like Support and Supervisor), or should it be a rule-based oracle that acts as the "gold standard" for training? Rule-based is simpler and more deterministic for RL; LLM-driven is more impressive for the demo.

My recommendation: LLM-driven for inference/demo, but with a deterministic fallback for training mode.

Q2: For the Hinglish customer simulator — should this only trigger on frustration, or should some tickets start in Hinglish from the beginning (to test multilingual capability)?

My recommendation: Start in English, degrade to Hinglish when frustration_level > 0.6. Add 2-3 tickets that start in Hinglish.

---

Architecture Overview

sequenceDiagram
    participant Client as Inference Client
    participant API as FastAPI Server
    participant Env as HierarchicalEnv
    participant L1 as Support Agent (L1)
    participant L2 as Supervisor (L2)
    participant L3 as Manager (L3)
    participant Cust as LLM Customer Sim
    participant Judge as LLM-as-Judge

    Client->>API: POST /reset?task=hierarchy_medium
    API->>Env: reset() → initial observation
    API-->>Client: {session_id, observation, active_role: "support_agent"}

    loop Hierarchical Step Loop
        Client->>API: POST /step (L1 action)
        API->>Env: step(L1 action)
        Env->>L2: Supervisor reviews L1 action
        alt Supervisor approves
            Env->>Cust: Generate customer reply
            Cust-->>Env: Customer response
        else Supervisor rejects / gives feedback
            Env-->>API: Observation with supervisor_feedback
            API-->>Client: {obs, reward, active_role: "support_agent", feedback: "..."}
        else Supervisor escalates to Manager
            Env->>L3: Manager intervenes
            L3-->>Env: Manager decision
            Env->>Cust: Customer reply based on manager action
        end
        Env->>Judge: Grade all agent actions
        Judge-->>Env: Role-specific rewards
        API-->>Client: {obs, reward, done, info}
    end

---

Proposed Changes

Component 1: Models (env/models.py)

[MODIFY] models.py

Expand the models to support the multi-agent hierarchy. Key additions:

New Enums:

• AgentRole — support_agent, supervisor, manager
• SupervisorDecision — approve, reject, feedback, escalate_to_manager
• ManagerDecision — override, approve_escalation, resolve_directly, send_back

Updated Action model:

• Add role: AgentRole field (who is taking this action)
• Add internal_note: Optional[str] (internal reasoning, not shown to customer)
• Add supervisor_decision: Optional[SupervisorDecision]
• Add manager_decision: Optional[ManagerDecision]
• Add feedback_to_agent: Optional[str] (supervisor/manager feedback)

Updated Observation model:

• Add active_role: AgentRole (which agent should act next)
• Add supervisor_feedback: Optional[str] (feedback from supervisor)
• Add manager_directive: Optional[str] (directive from manager)
• Add hierarchy_state: HierarchyState (nested model with internal comms)
• Add environment_event: Optional[str] (schema drift events)
• Add policy_context: str (current active policy, changes mid-episode)
• Add escalation_chain: List[str] (history of escalations)

New HierarchyState model:

• support_agent_actions: int
• supervisor_reviews: int
• manager_interventions: int
• current_phase: str (e.g., "support_handling", "supervisor_review", "manager_override")
• escalation_reason: Optional[str]

Updated Reward model:

• Add role-specific score fields: empathy_score, oversight_score, decision_quality_score
• Add role_reward: Dict[str, float] (per-role breakdown)

---

Component 2: Environment (env/environment.py)

[MODIFY] environment.py

Major refactor to support hierarchical step logic:

New class: HierarchicalCustomerSupportEnv (subclasses or replaces CustomerSupportEnv)

Hierarchical step logic:

1. Phase 1 — Support Agent acts: Client sends L1 action → env logs it
2. Phase 2 — Supervisor review: Environment auto-invokes supervisor logic:
   • If LLM-driven (inference mode): returns observation asking client for supervisor action
   • If rule-based (training mode): auto-evaluates based on policy rules
3. Phase 3 — Manager intervention (only if escalated): Same pattern as supervisor

Key methods:

• step_support(action) — L1 agent acts
• step_supervisor(action) — L2 reviews, decides approve/reject/escalate
• step_manager(action) — L3 resolves high-priority cases
• step(action) — unified entry point, routes based on action.role

Schema/Policy Drift:

• At random steps (configurable), inject environment_event into observation
• Examples: "Refund portal down", "New policy: max refund $50", "System outage: cannot query orders"
• Stored in self._active_policies dict, checked by reward engine

LLM Customer Simulator:

• Replace _FOLLOW_UPS with async LLM call via NVIDIA NIM
• Prompt template includes: persona, frustration level, conversation history, Hinglish trigger
• Frustration increases when tone is bad, decreases when empathetic
• When frustration > 0.6: 40% chance of Hinglish response

---

Component 3: LLM Customer Simulator (NEW)

[NEW] customer_simulator.py

Standalone module for the LLM-driven customer:

class CustomerSimulator:
    """LLM-driven customer that responds dynamically based on agent quality."""

    def __init__(self, api_key: str, base_url: str, model: str):
        ...

    async def generate_reply(
        self,
        persona: str,
        frustration_level: float,
        history: List[Message],
        ticket_context: str,
        use_hinglish: bool = False,
    ) -> str:
        """Generate contextual customer reply using LLM."""
        ...

    def _build_customer_prompt(self, ...) -> str:
        """Build the customer persona prompt with Hinglish instructions."""
        ...

Fallback: If LLM call fails, fall back to the existing _FOLLOW_UPS dict (graceful degradation).

---

Component 4: Reward Engine (env/reward_engine.py)

[MODIFY] reward_engine.py

Complete overhaul to hybrid dense reward system:

Overall Session Reward Components:

Component	Weight	Method
Resolution Quality	0.25	LLM-as-Judge
SLA Compliance	0.15	Rule-based (steps, timing)
Customer Satisfaction	0.15	Sentiment trajectory + LLM judge
Policy Adherence	0.15	LLM-as-Judge
Information Accuracy	0.10	Rule-based (regex patterns)
Efficiency	0.10	Rule-based (steps/max_steps)
Hierarchy Effectiveness	0.10	Rule-based (correct escalations, feedback quality)

Role-Specific Rewards:

Role	Metric	Weight	Method
Support Agent	Empathy & Tone	0.30	LLM-as-Judge
	Information Gathering	0.25	Rule-based
	Response Accuracy	0.25	LLM-as-Judge
	Efficiency	0.20	Rule-based
Supervisor	Oversight Quality	0.35	LLM-as-Judge (was the review correct?)
	Escalation Accuracy	0.30	Rule-based (should it have escalated?)
	Feedback Usefulness	0.20	LLM-as-Judge
	Speed of Review	0.15	Rule-based
Manager	Decision Quality	0.40	LLM-as-Judge
	Conflict Resolution	0.30	LLM-as-Judge
	Final Outcome	0.30	Rule-based (was it resolved?)

Penalties (Non-Gamable):

Penalty	Value	Trigger
Repetition	-0.15	TF-IDF cosine > 0.80
Policy Violation	-0.25	LLM detects violation of active policy
Unnecessary Escalation	-0.20	L1 escalates low-priority ticket
Unnecessary Manager Call	-0.20	L2 escalates when it shouldn't
Ignored Supervisor Feedback	-0.15	L1 repeats same mistake after feedback
Keyword Stuffing	-0.30	High keyword density without substance
Contradiction	-0.15	Claims done then asks for info

[NEW] llm_judge.py

class LLMJudge:
    """Async LLM-as-Judge for semantic reward evaluation."""

    RUBRIC_EMPATHY = """..."""
    RUBRIC_POLICY = """..."""
    RUBRIC_RESOLUTION = """..."""

    async def evaluate_empathy(self, message: str, context: str) -> float: ...
    async def evaluate_policy_adherence(self, action: Action, policy: str) -> float: ...
    async def evaluate_resolution(self, history: List[Message], ticket: dict) -> float: ...
    async def evaluate_supervisor_oversight(self, review: str, l1_action: Action) -> float: ...
    async def evaluate_manager_decision(self, decision: str, context: str) -> float: ...

Anti-Gaming Measures:

• LLM judge uses a strict rubric with negative examples
• Keyword density check: if resolution keywords > 5% of message, flag as stuffing
• Tone must be contextually appropriate (not just positive sentiment)
• Resolution must reference specific ticket details (not generic)

---

Component 5: Schema/Policy Drift (NEW)

[NEW] policy_engine.py

class PolicyEngine:
    """Manages dynamic policy changes and schema drift during episodes."""

    DRIFT_EVENTS = [
        {"trigger_step": 3, "event": "Refund portal is currently down. Do not promise immediate refunds.",
         "policy_change": {"can_refund": False}},
        {"trigger_step": 4, "event": "New policy: Maximum refund amount is now $50.",
         "policy_change": {"max_refund": 50}},
        {"trigger_step": 2, "event": "System outage: Order lookup service unavailable.",
         "policy_change": {"can_query_orders": False}},
    ]

    def check_drift(self, step: int, task: str) -> Optional[dict]: ...
    def get_active_policy(self) -> str: ...

---

Component 6: Graders

[NEW] task_hierarchy_easy.py

[NEW] task_hierarchy_medium.py

[NEW] task_hierarchy_hard.py

Each grader evaluates the full hierarchy:

• Was the Support Agent's initial response appropriate?
• Did the Supervisor make the right review decision?
• Was Manager intervention necessary and effective?
• Overall session quality

---

Component 7: OpenEnv Config

[MODIFY] openenv.yaml

Add new hierarchical tasks while keeping existing tasks:

tasks:
  # ... existing easy/medium/hard/nightmare ...
  - name: hierarchy_easy
    description: >
      Hierarchical multi-agent: Support Agent handles billing FAQ.
      Supervisor reviews and approves. No manager needed.
    max_steps: 8
  - name: hierarchy_medium
    description: >
      Hierarchical multi-agent: Support Agent handles technical issue.
      Supervisor may give feedback or request corrections.
      Mid-episode policy drift possible.
    max_steps: 12
  - name: hierarchy_hard
    description: >
      Hierarchical multi-agent: Critical SLA breach requiring all 3 levels.
      Support Agent must recognize urgency, Supervisor must escalate,
      Manager must make final decision. Schema drift guaranteed.
    max_steps: 15

action_space:
  type: ActionType
  values:
    - respond
    - escalate
    - close
    - request_info
    - supervisor_approve
    - supervisor_reject
    - supervisor_feedback
    - supervisor_escalate
    - manager_override
    - manager_resolve
    - manager_send_back

observation_space:
  # ... existing fields ...
  active_role: "support_agent | supervisor | manager"
  supervisor_feedback: "string | null"
  manager_directive: "string | null"
  environment_event: "string | null"
  policy_context: string
  escalation_chain: "list[string]"
  hierarchy_state:
    support_agent_actions: int
    supervisor_reviews: int
    manager_interventions: int
    current_phase: string

---

Component 8: Server

[MODIFY] app.py

• Import and register HierarchicalCustomerSupportEnv
• The existing /reset and /step endpoints continue to work for single-agent tasks
• For hierarchy_* tasks, /reset creates a HierarchicalCustomerSupportEnv
• /step auto-detects env type and routes accordingly
• No new endpoints needed — the hierarchy is managed inside the environment

---

Component 9: Inference

[MODIFY] inference.py

Add hierarchical inference mode:

• New HIERARCHY_TASKS list
• Role-specific system prompts:
  • SUPPORT_AGENT_PROMPT: Focus on empathy, info gathering, resolution
  • SUPERVISOR_PROMPT: Focus on reviewing L1 quality, policy compliance
  • MANAGER_PROMPT: Focus on high-stakes decisions, conflict resolution
• run_hierarchy_task(): Multi-turn loop that switches prompts based on active_role
• Existing run_task() unchanged for backward compatibility

---

Component 10: Ticket Store Updates

[MODIFY] ticket_store.py

Add hierarchy-specific tickets with Indian enterprise context:

• UPI payment failures
• Big Billion Days SLA breaches
• KYC document rejection loops
• Cross-border payment compliance issues

---

File Change Summary

File	Action	Description
env/models.py	MODIFY	Add AgentRole, hierarchy models, expand Action/Observation
env/environment.py	MODIFY	Add HierarchicalCustomerSupportEnv, keep original env
env/customer_simulator.py	NEW	LLM-driven customer with Hinglish support
env/llm_judge.py	NEW	LLM-as-Judge for semantic reward evaluation
env/policy_engine.py	NEW	Schema/policy drift management
env/reward_engine.py	MODIFY	Hybrid reward with role-specific scores + anti-gaming
env/ticket_store.py	MODIFY	Add hierarchy tickets with Indian context
env/graders/__init__.py	MODIFY	Register new hierarchy graders
env/graders/task_hierarchy_easy.py	NEW	Hierarchy easy grader
env/graders/task_hierarchy_medium.py	NEW	Hierarchy medium grader
env/graders/task_hierarchy_hard.py	NEW	Hierarchy hard grader
env/__init__.py	MODIFY	Export new classes
server/app.py	MODIFY	Support both env types in /reset and /step
inference.py	MODIFY	Add hierarchical inference with role-specific prompts
openenv.yaml	MODIFY	Add hierarchy tasks and expanded action/obs space
requirements.txt	MODIFY	Add aiohttp for async LLM calls
pyproject.toml	MODIFY	Add aiohttp dependency

---

Verification Plan

Automated Tests

# 1. Start the server
python -m server.app &

# 2. Run existing tests (backward compat)
pytest tests/ -v

# 3. Test hierarchy reset
curl -X POST http://localhost:7860/reset?task=hierarchy_easy

# 4. Test hierarchy step with L1 action
curl -X POST "http://localhost:7860/step?session_id=..." \
  -H "Content-Type: application/json" \
  -d '{"action_type": "respond", "message": "Hello, how can I help?", "role": "support_agent"}'

# 5. Run full hierarchy inference
python inference.py  # runs all tasks including hierarchy

Manual Verification

• Run a complete hierarchy_medium episode and verify all 3 agent levels are engaged
• Verify policy drift triggers mid-episode
• Verify Hinglish customer replies when frustration is high
• Verify LLM-as-Judge produces non-gamable scores
• Verify backward compatibility: old easy/medium/hard tasks still work identically

---

Next Steps for Unsloth Training

After this upgrade is working:

1. Create train_grpo.py: Wrap the HTTP API in a Gym-like interface for TRL
2. Generate trajectories: Run N episodes with base model, collect (state, action, reward) tuples
3. Train with GRPO: Use Unsloth + TRL GRPOTrainer on unsloth/Meta-Llama-3-8B-Instruct
4. Focus training: Train L1 (Support Agent) first, then L2 (Supervisor)
5. Generate plots: baseline vs. trained reward curves with matplotlib