readme_gen.py

"""
readme_gen.py — dynamically generated README for downloaded RoboGen datasets.
"""

from __future__ import annotations
import datetime
from typing import Dict, List


COLUMN_DOCS = {
    "state_0":  "Joint 0 — base rotation (rad)",
    "state_1":  "Joint 1 — shoulder pitch (rad)",
    "state_2":  "Joint 2 — elbow (rad)",
    "state_3":  "Joint 3 — wrist roll (rad)",
    "state_4":  "Joint 4 — wrist pitch (rad)",
    "state_5":  "Joint 5 — gripper (rad, 0=closed → 1=open)",
    "action_0": "Velocity command joint 0 — base (rad/s)",
    "action_1": "Velocity command joint 1 — shoulder (rad/s)",
    "action_2": "Velocity command joint 2 — elbow (rad/s)",
    "action_3": "Velocity command joint 3 — wrist roll (rad/s)",
    "action_4": "Velocity command joint 4 — wrist pitch (rad/s)",
    "action_5": "Velocity command joint 5 — gripper (rad/s)",
    "timestamp":       "Seconds since episode start (50 Hz → Δt=0.02 s)",
    "episode_index":   "Integer episode identifier (0-indexed)",
    "frame_index":     "Frame number within episode (0–49 for 50-frame episodes)",
    "task":            "Task label string",
    "use_for_training":"True for successful episodes only; False for failure episodes",
    "failure_type":    "'success' | 'grasp_slip' | 'velocity_spike' | 'torque_saturation'",
    "quality_score":   "Per-episode quality score (0–100) from HaptalAI scorer",
    "robot":           "Robot model string: 'SO-100' | 'SO-101' | 'Koch'",
}

FAILURE_DESCRIPTIONS = {
    "grasp_slip":        "Smooth trajectory until 60-70% of episode, then gripper opens "
                         "unintentionally (position discontinuity ≥ 0.18 rad) and contact "
                         "force collapses. Mimics inadequate grasp force.",
    "velocity_spike":    "1-2 isolated frames with joint velocity MAD z-score > 6.5 rad/s, "
                         "surrounded by normal motion. Mimics servo glitch or controller "
                         "communication dropout.",
    "torque_saturation": "One arm joint clamped at its angular limit for ≥ 3 consecutive frames "
                         "with near-zero velocity. Mimics joint hitting mechanical stop or "
                         "exceeding torque budget.",
}


def generate_readme(
    robot:         str,
    task:          str,
    n_episodes:    int,
    success_rate:  float,
    force_min:     float,
    force_max:     float,
    failures:      List[str],
    score:         float,
    band:          str,
    n_passed:      int,
    n_flagged:     int,
    mean_mismatch: float,
    failure_breakdown: Dict[str, int],
    scorer_used:   str,
) -> str:
    """Generate a complete README.md for the downloaded dataset."""

    task_display   = task.replace("_", " ").title()
    failures_list  = "\n".join(f"- **{f}**: {FAILURE_DESCRIPTIONS.get(f, f)}" for f in failures)
    col_table_rows = "\n".join(
        f"| `{col}` | {desc} |" for col, desc in COLUMN_DOCS.items()
    )
    fb_lines = "\n".join(f"- {k}: {v} episodes" for k, v in failure_breakdown.items()) or "- None"
    generated_at   = datetime.datetime.utcnow().strftime("%Y-%m-%d %H:%M UTC")
    band_emoji     = ""

    return f"""# RoboGen Synthetic Dataset — {robot} / {task_display}

> Generated by [HaptalAI RoboGen](https://huggingface.co/spaces/HaptalAI/robogen)
> on {generated_at}

---

## Dataset Summary

| Field | Value |
|---|---|
| Robot | **{robot}** |
| Task | **{task_display}** |
| Total episodes | **{n_episodes}** |
| Success rate (configured) | **{success_rate * 100:.0f}%** |
| Contact force range | **{force_min:.1f} – {force_max:.1f} N** |
| Frames per episode | **50** (50 Hz, Δt = 0.02 s) |
| Total rows | **{n_episodes * 50:,}** |

---

## Quality Score

| Metric | Value |
|---|---|
| Overall score | **{score:.1f} / 100** |
| Band | **{band}** |
| Episodes passed | **{n_passed}** |
| Episodes flagged | **{n_flagged}** |
| Mean mismatch rate | **{mean_mismatch:.4f}** |
| Scorer | `{scorer_used}` |

**Quality bands:**
- **Clean** (>= 80): suitable for policy training and augmentation
- **Review** (55-79): usable with caution; inspect flagged episodes
- **Flagged** (< 55): high anomaly rate; use for failure analysis only

---

## What the Dataset Contains

This dataset contains **{n_episodes} synthetic episodes** of a **{robot}** robot performing
the **{task_display}** task. Each episode is 50 frames (1 second at 50 Hz).

**Episode composition:**
- Success episodes (`use_for_training=True`): ~{success_rate * 100:.0f}% of total
- Failure episodes: ~{(1 - success_rate) * 100:.0f}% of total

### Failure types included
{failures_list}

### Failure breakdown in this dataset
{fb_lines}

---

## Column Reference

| Column | Description |
|---|---|
{col_table_rows}

---

## Physics Model

Joint trajectories are generated using **cubic spline interpolation** over
task-specific waypoints (approach → contact/grasp → lift/push → retract).
Velocities are the **analytical first derivative** of the position spline — not
independently sampled — ensuring physical consistency between state and action.

- Sensor noise: Gaussian σ_pos = 0.002 rad, σ_vel = 0.004 rad/s
- Contact force: spring-damper model during contact window (30–75% of episode)
- Episode variation: small Gaussian perturbations on target position (±2.5 cm equivalent)
- Joint limits: enforced per robot specification

---

## Recommended Use

Synthetic data is best used for:
1. **Policy bootstrapping** — pre-train before collecting real demonstrations
2. **Augmentation** — mix with real data to increase diversity and robustness
3. **Failure analysis / anomaly detection** — the labelled failure episodes are
   especially useful for training or evaluating anomaly detectors
4. **Simulation-to-real transfer research** — study domain gap with known ground truth

> **Do not** rely solely on synthetic data for safety-critical deployments.
> Always validate against real demonstrations before deploying to physical hardware.

---

## Validation & Benchmark

This dataset was generated and validated by **HaptalAI's misalignment failure benchmark
and physical failure scorer** — the same pipeline used to evaluate community SO-100
datasets. Calibrated thresholds:
- Velocity spike: MAD z-score > 6.5 rad/s
- Mismatch fraction: > 0.50 per episode → flagged

For questions, dataset requests, or benchmark access:
**aarav@haptal.ai**

---

*RoboGen is open source. Star us on GitHub and contribute at
[HaptalAI/robogen](https://github.com/aaravbedi/robogen).*
"""