ReGal: A First Look at PPO-based Legal AI for Judgment Prediction and Summarization in India

AI summary

Overview

Research area: Natural Language Processing applied to Indian legal AI — specifically Court Judgment Prediction and Explanation (CJPE) and legal document summarization, using reinforcement learning (PPO) instead of purely supervised fine-tuning.

Technical level: Intermediate. Readers need passing familiarity with supervised fine-tuning, reinforcement learning from human/AI feedback, and n-gram and embedding-based text metrics.

Scope: A position-paper-style early exploration showing that a PPO-based RLAIF pipeline (ReGal), built on Llama-2-7B, currently underperforms supervised fine-tuned and proprietary models on Indian Supreme Court judgment prediction, explanation, and summarization, while documenting why.

What This Paper Is About

Most work on Indian legal NLP relies on supervised fine-tuning of pretrained language models, which depends on large annotated datasets and cannot incorporate feedback during generation. This paper asks whether reinforcement learning — specifically Proximal Policy Optimization (PPO) with AI-generated reward signals (RLAIF) — can instead refine both judgment predictions and their rationales, and whether the same optimization loop transfers to legal summarization. The answer the authors report is largely negative on metrics, but the paper's value lies in the diagnosis of why.

Key Contributions

1. A first PPO-based legal AI framework for India. The authors present ReGal (Reinforcement Learning-based Legal Reasoning), which combines multi-task instruction tuning with RLAIF using PPO, and state it is one of the first applications of PPO-based reinforcement learning to Indian legal judgment prediction and summarization. They also claim it is the first application of RLHF/RLAIF to both legal judgment prediction and summarization in the Indian legal domain.

2. Empirical and qualitative evidence of failure modes. The paper reports results on two datasets (PredEx and In-Abs) showing ReGal scoring below supervised and proprietary baselines, plus a dedicated hallucination analysis documenting fabricated precedent citations and hallucinated legal principles.

3. Ablations over base model and reward model. The authors test a smaller base model (Phi-3 Mini) and an un-fine-tuned pretrained LLaMA-2-7B, and swap the task-specific reward model for a generic InLegalBERT-pretrained one, isolating how initialization and reward alignment drive PPO outcomes.

4. A research agenda. The paper explicitly positions itself as groundwork, proposing improvements in reward modeling, domain-adaptive pretraining, human-in-the-loop feedback, and adaptive RL techniques.

Main Findings

• ReGal underperforms on judgment prediction and explanation. On the PredEx dataset, ReGal scored ROUGE-1 0.19, ROUGE-2 0.04, ROUGE-L 0.12, BLEU 0.01, METEOR 0.10, BERTScore 0.50, and BLANC 0.02 — against LLaMA-2 SFT at ROUGE-1 0.50 and BLEU 0.25, and Gemini Pro at ROUGE-1 0.31 and BERTScore 0.63.

• The same pattern holds on ILDC Expert. ReGal reached ROUGE-1 0.25, ROUGE-2 0.05, ROUGE-L 0.16, BLEU 0.01, METEOR 0.16, BERTScore 0.50, BLANC 0.03, compared with GPT-3.5 Turbo at ROUGE-1 0.54 and BERTScore 0.73, and LLaMA-2 SFT at ROUGE-1 0.49.

• PPO also trails simpler inference strategies. On PredEx inference, PPO scored ROUGE-1 0.30 versus SFT 0.42 and vanilla inference 0.39; DPO scored 0.38. On In-Abs summarization, PPO scored ROUGE-1 0.41, while vanilla inference scored 0.47 and SFT and DPO each scored 0.44.

• Eight candidate causes of underperformance are listed. Objective mismatch between the SFT starting point and the PPO objective; reward model limitations; legal text complexity; insufficient diversity in training data; PPO's dependence on the reward model without human-in-the-loop supervision; suboptimal hyperparameter selection; model size and architecture; and a domain pretraining gap relative to GPT-3.5 Turbo.

• Smaller or unadapted base models degrade further. Replacing LLaMA-2-7B with Phi-3 Mini dropped performance substantially across all metrics on both PredEx and In-Abs, with the model failing to produce coherent or factually grounded legal reasoning. Using pretrained LLaMA-2-7B without supervised fine-tuning also degraded lexical and semantic scores.

• An unaligned reward model makes things worse. Substituting the task-specific PredEx-fine-tuned legal classifier with an InLegalBERT-pretrained model that was not fine-tuned for the judgment-explanation task produced noisier, misaligned scores, incoherent or generic outputs, and further degradation in PPO performance.

• Hallucination is a central failure mode. The model fabricated precedent citations, invented legal principles, and asserted facts absent from the input. In one documented example, ReGal claimed the appellant's right to privacy was upheld under Article 21 based on reasoning not present in the original judgment. Hallucination appeared most when input facts were sparse or ambiguous and when PPO over-optimized for reward patterns from an imperfect reward model.

Methodology in Plain English

The authors built a two-stage pipeline on top of Llama-2-7B, chosen to match prior work in the Indian legal domain so that comparisons isolate the effect of reinforcement learning.

In stage one, they fine-tuned the model with supervised instruction tuning on two tasks: judgment prediction plus explanation using the PredEx dataset, and abstractive summarization using the In-Abs dataset. This supervised model, denoted π^SFT, becomes the reference policy.

In stage two, they applied PPO. For judgment prediction, the reward model is a fine-tuned InLegalBERT classifier that assigns a binary reward — 1 if the predicted verdict is correct, 0 otherwise. For summarization, the reward model scores outputs on n-gram overlap with gold headnotes and coherence, using ROUGE-style matching and shallow semantic similarity. These reward models stand in for human feedback, which is what makes the setup RLAIF rather than RLHF.

PPO then updates the policy to maximize reward while penalizing deviation from the SFT reference policy through a KL term, with a clipping parameter ε of 0.1 that prevents overly large policy updates. Training ran on a rented NVIDIA A100 80GB GPU via Vast.ai at a total cost of approximately $100, with a learning rate of 1.41e-5, a maximum of 1 PPO epoch, batch size 4, mini-batch size 2, output lengths between 100 and 500 tokens, and mixed-precision training using GradScaler. Maximum new tokens at inference was 500.

Evaluation used both lexical metrics — ROUGE-1/2/L, BLEU, METEOR — and semantic metrics — BERTScore and BLANC.

Why This Matters

Impact on research. The paper is an honest negative result in a domain where negative results are rarely published. It argues that PPO's success in legal AI depends not just on the RL algorithm but on strong initialization and a precisely aligned reward function, and it documents specific failure modes — reward model misalignment, hallucination under sparse reward signals, domain pretraining gaps — that other researchers can target directly. It also contributes one of the first RLHF/RLAIF baselines for Indian legal NLP, giving future work a comparison point.

Real-world applications:

• Case screening and triage. A reliable judgment prediction component could help legal aid organizations prioritize which petitions or appeals to pursue, though this paper's results are not yet at that reliability.
• Headnote-style summarization. Abstractive summaries of long Supreme Court decisions, the In-Abs task, would help practitioners quickly grasp background, legal issues, arguments, and verdicts without reading thousands of words.
• Explainability and transparency. The rationale-generation subtask targets explanations that mimic judicial argumentation, which matters for trust in any AI-assisted legal tool.
• Legal education and research. Summaries and outcome predictions can support law students and researchers analyzing large corpora of judgments.

Industry relevance. Legal technology providers, court administration systems, and legal publishers have an interest in automating judgment analysis, but this paper's findings — especially the fabricated citations and invented Article 21 reasoning — set a clear caution: the technology is not yet safe for deployment without strict validation, legal oversight, and human-in-the-loop review. The authors state that outputs should not be considered substitutes for legal advice.

Future Directions

1. Build better-aligned reward models. The paper's ablations show that a task-specific, domain-tuned reward model is essential, and that a generic legal reward model disrupts learning. More expressive reward signals are the stated next step.

2. Integrate human-in-the-loop feedback. PPO's reliance on an imperfect proxy for human judgment is identified as a core limitation, so combining RLAIF with actual expert feedback is proposed to reduce hallucination and improve explanation quality.

3. Apply domain-adaptive pretraining. The "domain pretraining gap" against GPT-3.5 Turbo is listed as a factor, and the authors propose domain-specific pretraining strategies to close it.

4. Optimize PPO itself. Refined hyperparameters (learning rate, batch size, KL penalty coefficients) and adaptive RL techniques, along with alternative or larger model architectures, are proposed as paths to a more competitive framework.

Open questions the paper leaves unresolved: whether PPO can ever match SFT on long-document legal tasks, how to design hallucination-aware reward models, and whether the ReGal framework extends to other jurisdictions or high-stakes reasoning domains, which the authors claim it is designed to do.

Target Audience

Researchers working on legal NLP, reinforcement learning from human or AI feedback, and long-document generation will find the diagnostic detail most useful, particularly the reward-model ablation. Practitioners building legal AI products for the Indian market should read the hallucination analysis before considering deployment. Graduate students looking for a documented negative result and a concrete set of open problems in RLHF for high-stakes domains will also benefit, as will legal scholars interested in how automated systems handle Indian Supreme Court reasoning.

Authors’ abstract

This paper presents an early exploration of reinforcement learning methodologies for legal AI in the Indian context. We introduce Reinforcement Learning-based Legal Reasoning (ReGal), a framework that integrates Multi-Task Instruction Tuning with Reinforcement Learning from AI Feedback (RLAIF) using Proximal Policy Optimization (PPO). Our approach is evaluated across two critical legal tasks: (i) Court Judgment Prediction and Explanation (CJPE), and (ii) Legal Document Summarization. Although the framework underperforms on standard evaluation metrics compared to supervised and proprietary models, it provides valuable insights into the challenges of applying RL to legal texts. These challenges include reward model alignment, legal language complexity, and domain-specific adaptation. Through empirical and qualitative analysis, we demonstrate how RL can be repurposed for high-stakes, long-document tasks in law. Our findings establish a foundation for future work on optimizing legal reasoning pipelines using reinforcement learning, with broader implications for building interpretable and adaptive legal AI systems.