← Back to list

Distress AI Coach: Teaching an LLM to Listen Before It Gives Advice

An OpenEnv environment for training small language models to handle difficult, long-horizon human conversations.

tej · 2026-04-26 07:46 · 0 claps · 4.9 min read
#reinforcement-learning #ai #knowledge-distillation
Open on Medium ↗
Wiki topics: LLM · Large Language Models FT · Fine-tuning & Adaptation AI · AI · General EDU · Education & Learning 🧠 · Mental Wellness

Distress AI Coach: Teaching an LLM to Listen Before It Gives Advice

An OpenEnv environment for training small language models to handle difficult, long-horizon human conversations.

Most supportive AI systems are good at sounding empathetic.

They say things like:

“I hear you.” “That sounds difficult.” “You’re not alone.”

That is useful, but it is not the hard part.

The hard part is what good therapists, coaches, and emotionally intelligent friends actually do:

They do not jump to advice immediately. They surface the real fear first. They remember what you said three turns ago. They help you leave the conversation with one concrete next step. And when the stakes are high, they de-escalate instead of just sounding kind.

That is what we wanted to train.

So we built Distress AI Coach — an OpenEnv environment where an LLM has to coach a simulated human through difficult real-world conversations.

This was built for the OpenEnv Hackathon India 2026, under Theme #2: Long-Horizon Planning & Instruction Following.

The idea

Distress AI Coach is an environment for training small LLMs to handle conversations that unfold over time.

The agent is not rewarded for sounding polite.

It is rewarded for what its replies cause later.

Does the person open up? Does the agent avoid premature advice? Does it remember unresolved fears? Does the conversation end with a concrete next step? Does it de-escalate when the situation is risky?

These are not one-turn behaviors.

They require memory, patience, state tracking, and delayed reward optimization.

The three conversation tasks

The environment includes three difficult conversation scenarios:

  1. Boundary reset with a manager The user needs to reset expectations with their boss without sounding defensive.
  2. Repair conversation after hurting someone The user needs to apologize, take responsibility, and rebuild trust.
  3. Boundary-setting with a volatile partner The user needs to plan a difficult conversation where escalation risk is real.

The third task is intentionally harder because the agent must protect the user logistically, not just emotionally.

A shallow model might say:

“Just be honest with them.”

A better model should ask:

“Do you feel physically safe having this conversation in person?”

That difference matters.

Why this is a long-horizon problem

Most chatbot evaluations are turn-level.

They ask whether a single response is helpful, polite, harmless, or empathetic.

But coaching is trajectory-level.

A reply that looks good now can hurt the conversation later.

For example, if the user says:

“I need to talk to my manager, but I keep avoiding it.”

A weak assistant may immediately generate a script:

“Here’s what you should say…”

That sounds useful, but it may be premature.

The real fear might be:

“I think I’ll be seen as incompetent.”

Or:

“My manager has retaliated before.”

Or:

“I don’t even know what outcome I want.”

A good coach should surface that first.

In this environment, the reward for not jumping to advice may only appear several turns later, when the seeker finally reveals the real concern.

That is why this is a long-horizon planning problem, not just an empathy benchmark.

What the agent sees

The agent gets only partial information.

Example observation:

seeker_utterance: "I keep putting off this conversation with my boss…"
turn: 3
stage_hint: "early — surface the real concern, don't script yet"
remaining_turns: 12
session_index: 1 / 2
durable_memory: [unresolved threads from prior sessions]

The agent does not see:

hidden stress level
hidden trust level
hidden openness level
the seeker's true unspoken fear
future reward implications of the current reply

So the agent has to infer the state of the conversation from behavior.

That makes it partially observable, closer to real human conversations.

What gets rewarded

We use a dense hybrid reward instead of a single end-of-episode success score.

The reward includes:

Immediate turn reward
- Is this reply useful right now?
- Is it specific?
- Does it avoid generic empathy?
Future-oriented trajectory reward
- Does this reply make the future conversation better?
- Does it increase openness?
- Does it avoid premature scripting?
Long-horizon continuity reward
- Does the agent remember unresolved threads?
- Does it use durable memory across sessions?
Anti-repetition penalty
- Punishes generic loops
- Punishes “I hear you” boilerplate
- Punishes filler like “as I said earlier”
Hard-gated success
- Final reveal
- Coherent plan
- De-escalation setup on hard tasks

This makes the environment harder to reward-hack.

A model that says:

“I hear you, that sounds hard.”

on every turn will score poorly because it does not move the conversation forward.

The training setup

We trained a small model using reward-guided distillation.

The policy model was:

Qwen2.5-0.5B-Instruct

The reward judge was:

Qwen2.5-3B-Instruct

The pipeline works like this:

  1. The 0.5B policy generates multiple candidate replies.
  2. The 3B judge scores those replies using the reward rubric.
  3. The best reply is selected.
  4. The 0.5B model is fine-tuned on those selected examples using LoRA / QLoRA.

In simple terms:

A larger model teaches a tiny model what good coaching behavior looks like.

The training command:

python train_trl.py \
  --model-name Qwen/Qwen2.5-0.5B-Instruct \
  --reward-backend llm_judge \
  --reward-model-name Qwen/Qwen2.5-3B-Instruct \
  --episodes-per-task 4 \
  --epochs 3 \
  --batch-size 1 \
  --gradient-accumulation-steps 8 \
  --max-length 512 \
  --use-4bit true

Results

The LoRA fine-tuned 0.5B model improved over the base model across the deterministic ESC environment.

The biggest gains appeared on the medium and hard tasks, where premature scripting hurts the most.

That is exactly where the reward was designed to punish shallow empathy and reward long-horizon coaching behavior.

The loss curve decreased steadily across epochs, and the before/after reward comparison showed clear improvement after fine-tuning.

Why a 0.5B model?

Most AI coach demos hide a large model behind an API.

We wanted the smallest model that could visibly improve.

A fine-tuned Qwen2.5–0.5B model is:

~250 MB on disk when 4-bit quantized
fast enough for phone-class inference
cheap to run after training
private by default

This matters for distress conversations.

If someone is planning a boundary conversation with a manager, partner, parent, or friend, that conversation should not need to leave their device.

The goal is not just intelligence.

The goal is deployable, private, low-cost intelligence.

What is in the repo

distress-ai-coach/
├── server/app.py              # OpenEnv-compliant FastAPI server
├── openenv.yaml               # OpenEnv manifest
├── src/
│   ├── env.py                 # Gym-style ESC environment
│   ├── seeker.py              # Deterministic seeker simulator
│   ├── tasks.py               # Conversation scenarios
│   ├── grader.py              # Hybrid reward implementation
│   ├── reward_backend.py      # Regression + LLM judge backends
│   └── training_utils.py      # Rollout + prompt utilities
├── train_reward_model.py      # QLoRA reward model training
├── train_trl.py               # Policy training with TRL + LoRA
├── gradio_app.py              # Demo frontend
└── results/                   # Loss, reward, before/after plots

Final thought

The future of supportive AI should not just be bigger models saying nicer things.

It should be smaller, private models that learn when to pause, when to ask, when to remember, and when not to give advice too early.

That is the core idea behind Distress AI Coach:

a tiny model trained to listen first.

Team and Acknowledgements

Deep Q-Lads (Sivateja and Gokul)

Reward shaping inspired by RLFF-ESC (Yang et al., 2025, arXiv:2508.12935). Built on OpenEnv, Hugging Face TRL, PEFT, and bitsandbytes.

Submitted to the OpenEnv Hackathon (India 2026) — Theme #2: Long-Horizon Planning & Instruction Following.


메타데이터
post_id
b1eb7e8b0fa9
slug
distress-ai-coach-teaching-an-llm-to-listen-before-it-gives-advice-b1eb7e8b0fa9
url
https://medium.com/@5ivatej/distress-ai-coach-teaching-an-llm-to-listen-before-it-gives-advice-b1eb7e8b0fa9
canonical_url
https://medium.com/@5ivatej/distress-ai-coach-teaching-an-llm-to-listen-before-it-gives-advice-b1eb7e8b0fa9
author_url
https://medium.com/@5ivatej
status
ok
fetched_at
2026-06-09 15:37:30