[Bounty] Update Dr. Agent model (PyHealth 2.0)

[joshuasteier]

Contributor

• Name: Joshua Steier

Contribution Type

Model

Description

Implements the Dr. Agent model from "Dr. Agent: Clinical predictive model via mimicked second opinions" (JAMIA 2020).

Key features:

• Dual-agent reinforcement learning with dynamic skip connections
• Primary agent (current visit) + second-opinion agent (historical context)
• Support for static demographic features
• GRU/LSTM cell options
• REINFORCE with optional baseline for variance reduction

Paper: https://doi.org/10.1093/jamia/ocaa074

Files to Review

File	Description
pyhealth/models/agent.py	Model implementation (AgentLayer + Agent)
tests/test_agent.py	Unit tests
examples/agent_mimic4.ipynb	Example notebook with MIMIC-IV

How to Test

# Run unit tests
python -m pytest tests/test_agent.py -v

# Or run the model directly
python pyhealth/models/agent.py

Checklist

• Code follows PEP8 (88-char line limit)
• Google-style docstrings with type hints
• Inherits from BaseModel (v2.0 style)
• Unit tests provided
• Example notebook provided
• Rebased with main branch

[Copilot]

Pull request overview

This PR implements the Dr. Agent model from "Dr. Agent: Clinical predictive model via mimicked second opinions" (JAMIA 2020) for PyHealth 2.0. The model uses dual reinforcement learning agents with dynamic skip connections to capture long-term dependencies in patient EHR sequences.

Key changes:

• New Agent and AgentLayer classes implementing the Dr. Agent architecture with REINFORCE-based policy gradient learning
• Comprehensive unit tests covering model initialization, forward/backward passes, and various configurations
• Example notebook demonstrating mortality prediction on MIMIC-IV data

Critical issues identified:

1. Missing tensor dimensionality reduction for NestedSequenceProcessor embeddings (4D → 3D pooling required)
2. Incorrect input schema in tests and examples ("sequence" should be "nested_sequence")
3. Binary classification reward computation doesn't provide feedback for negative class predictions
4. Mode resolution doesn't use BaseModel's _resolve_mode() method

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 14 comments.

File	Description
pyhealth/models/agent.py	Core model implementation with AgentLayer (dual RL agents) and Agent (BaseModel wrapper). Contains critical bugs in embedding handling and reward computation.
tests/core/test_agent.py	Unit tests for the model. Uses incorrect schema and doesn't test the temporal/nested sequence case.
examples/agent_mimic4.ipynb	Tutorial notebook for MIMIC-IV mortality prediction. Contains minor spelling issues and will fail with current model bugs.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.