# Integrated Silicon Photonics for High-Speed Quantum Key Distribution

P. Sibson<EMAIL_ADDRESS>J. E. Kennard  Quantum Engineering
Technology Labs, H. H. Wills Physics Laboratory & Department of Electrical and
Electronic Engineering, University of Bristol, Merchant Venturers Building,
Woodland Road, Bristol, BS8 1UB, UK  S. Stanisic  Quantum Engineering
Technology Labs, H. H. Wills Physics Laboratory & Department of Electrical and
Electronic Engineering, University of Bristol, Merchant Venturers Building,
Woodland Road, Bristol, BS8 1UB, UK  Quantum Engineering Centre for Doctoral
Training, H. H. Wills Physics Laboratory & Department of Electrical and
Electronic Engineering, University of Bristol, Merchant Venturers Building,
Woodland Road, BS8 1UB, UK  C. Erven  J. L. O’Brien  M. G. Thompson  Quantum
Engineering Technology Labs, H. H. Wills Physics Laboratory & Department of
Electrical and Electronic Engineering, University of Bristol, Merchant
Venturers Building, Woodland Road, Bristol, BS8 1UB, UK

###### Abstract

Integrated photonics offers great potential for quantum communication devices
in terms of complexity, robustness and scalability. Silicon photonics in
particular is a leading platform for quantum photonic technologies, with
further benefits of miniaturisation, cost-effective device manufacture and
compatibility with CMOS microelectronics. However, effective techniques for
high-speed modulation of quantum states in standard silicon photonic platforms
have been limited. Here we overcome this limitation and demonstrate high-speed
low-error quantum key distribution modulation with silicon photonic devices
combining slow thermo-optic DC biases and fast (10 GHz bandwidth) carrier-
depletion modulation. The ability to scale up these integrated circuits and
incorporate microelectronics opens the way to new and advanced integrated
quantum communication technologies and larger adoption of quantum-secured
communications.

Figure 1: Integrated Silicon Photonic Devices for QKD: (a) Coherent-One-Way: A
balanced Mach-Zehnder interferometer (MZI) comprising two multi-mode
interference (MMI) devices acting as beam splitters, with phase modulation
from thermo-optic phase modulators (TOPM) and carrier-depletion modulators
(CDM) allows for the encoding of quantum information in path, or pulse
modulation. (b) Polarisation encoded BB84: Combining the two paths of the MZI
with a two-dimensional grating coupler allows for the conversion from path
encoded information to polarisation encoded information (P2P), suitable for
communication in free-space. (c) Time-bin encoded BB84: An unbalanced
asymmetric MZI (AMZI) allows for encoding in time by temporally separating
weak coherent pulses into two time intervals using an on-chip delay of 1.5 ns.
The extra loss this incurs is balanced by an MZI used as a tunable beam
splitter on the opposing arm. The last beam splitter in the AMZI is replaced
with another MZI that allows for the selection of time-bin $|0\rangle$ and
$|1\rangle$ states. DC offsets are provided by the TOPMs and fast modulation
by four CDMs. (d) An illustrated Bloch sphere highlighting the DC offset at
$|+\mathrm{i}\rangle$ (red dot), set by the TOPMs. Each CDM is only required
to modulate up to $\pi/2$ to permit the encoding of each BB84 state.

Quantum technologies are rapidly developing and have the potential to
revolutionise the fields of computing and telecommunications. They have major
implications for the security of many of our conventional cryptographic
techniques, which are known to be insecure against a quantum computer
QKDREVIEW . Fortunately, quantum key distribution (QKD) provides a highly
secure approach to sharing random encryption keys by transmitting single
photons Scarani2009 . Although QKD has advanced from simple proof-of-principle
experiments towards robust long-term demonstrations Yoshino:13 ; Korzh:13 ;
Dixon:15 ; sasaki2015quantum , it has still not obtained wide-scale adoption.

Integrated photonics provides a stable, compact, and robust platform to
implement complex photonic circuits amenable to mass-manufacture, and
therefore provides a compelling technology for optical quantum information
devices INTREVIEW . Silicon photonics, in particular, is a leading platform
for quantum photonic technologies with the promise of high density
integration, mature fabrication processing, and compatibility with
microelectronics lim2014review . Silicon has been used to demonstrate sources
of quantum light, manipulation and transmission of quantum information, and
integration with single photon detectors Silverston2016Silicon ; Pernice2012 ;
IntegratedSNSPDS . It has also been used in classical computing and
communications for modulation, transceivers doerr2015silicon , and a recent
demonstration of optical interconnects alongside electronic microprocessor
technology sun2015single .

The appeal of this platform has led to integrated photonic technologies
increasingly being deployed in the development of practical QKD systems.
Demonstrations include integrated “client” chips for Reference-Frame-
Independent QKD RFI , planar waveguide components in transmitters and
receivers tanaka2012high , and chip-to-chip QKD using GHz clocked indium
phosphide transmitters and silicon oxynitride receivers sibson2015chip .

However, high-speed modulation of quantum states in standard silicon photonic
fabrication has been limited. With no natural electro-optic non-linearity,
many silicon quantum photonic experiments instead utilise slow thermo-optic
phase modulators (TOPMs) for high-fidelity state preparation. Carrier
injection or carrier depletion modulators (CDMs) offer high-speed operation,
but incur phase dependent loss and saturation reed2010silicon , which are
detrimental in quantum applications where state preparation has stringent
requirements.

Here we show an approach to overcome the limitations of saturation and phase
dependent loss of high-speed CDMs in standard silicon photonic fabrication.
First we describe a combination of slow, but ideal, TOPMs alongside fast, but
non-ideal, CDMs utilised for QKD state preparation at GHz speeds. We then use
this technique to demonstrate three implementations of high-speed low-error
QKD (Figure 1): chip-to-chip Coherent One Way (COW) QKD, polarisation encoded
BB84, and time-bin encoded BB84 BB84  state preparation. We achieve estimated
asymptotic secret key rates of up to 916 kbps and quantum bit error rates
(QBER) as low as 1.01% over 20 km of fibre, experimentally demonstrating the
feasibility of high-speed QKD integrated circuits based on standard silicon
photonic fabrication.

Figure 1 shows the three different silicon photonic devices which use a
combination of TOPMs and CDMs to prepare and modulate QKD states. As
illustrated in Figure 2 (a), the TOPMs allow for almost ideal state
preparation with a phase relationship proportional to the square of the
voltage ($V^{2}$) and no change in transmission ($\Delta\mathrm{T}$). The
TOPMs were designed in silicon-on-insulator using doped resistive heating in
the waveguide slab. This design provides ohmic electrical characteristics
($\sim$6.14 k$\Omega$ for a 150 $\mu$m length) and a $2\pi$ voltage of
$\sim$24 V.

The TOPMs are limited in bandwidth ($\sim$kHz) and therefore inappropriate for
the fast modulation required in communications. Silicon also has no natural
$\chi^{(2)}$ non-linearity, and therefore the high-speed electro-optic effect
cannot be used. An approach to overcome this in standard silicon photonic
platforms is to use doped waveguide sections as shown in Figure 2 (b), and
employ carrier injection or depletion techniques.

Carrier depletion modulation induces a phase by reducing the carriers
occupying the region which overlaps with the optical mode. This is achieved by
reverse biasing a _p-n_ junction formed by doping _p_ and _n_ regions in the
core of the silicon waveguide. The depletion of carriers in the waveguide
decreases the absorption of the waveguide and induces an optical phase change
fang2012recent ; however, this effect saturates as the waveguide becomes fully
depleted of carriers. Figure 2 (b) plots the extracted relationship between
voltage, phase, and transmission for a 1.5 mm modulator ($\sim$5 dB loss at 0
V), illustrating that the induced phase saturates below $\pi$ over 8 V while
the transmission continues to increase.

The phase-dependent loss characteristics and saturation of the CDMs can
severely reduce operational fidelity. This is especially damaging for quantum
applications where requirements are often more stringent than many classical
applications. Here we describe a combination of thermal and carrier depletion
modulators to minimise this effect. By using the TOPMs to bias the circuits in
a favourable operating regime, and limiting the modulation depth required for
each individual CDM, we mitigate these negative characteristics for both pulse
modulation and state preparation.

Figure 2: Thermo-optic and Carrier-depletion Phase Modulation in Silicon
Photonics fabricated with standard doping processes absil2015imec . (a) Cross
section of the thermo-optic phase modulation waveguide with _p_ ++ doping in
the waveguide slab and intrinsic (_i_) silicon waveguide core, followed by the
power measured at the two outputs of an MZI, the fitted quadratic phase
($\phi$) relationship, and the change in transmission ($\Delta\mathrm{T}$) as
a function of the applied voltage (V). (b) Cross section of the carrier
depletion phase modulator with _p_ and _n_ doping in the waveguide core,
followed by the power measured at the two outputs of an MZI (with an
additional TOPM providing a $\pi/2$ offset or initially equal intensity
outputs), the fitted phase ($\phi$) relationship illustrating saturation, and
the change in transmission ($\Delta\mathrm{T}$) as a function of the applied
voltage (V).

Utilising this technique, we demonstrate the pulse modulation of coherent
light using a Mach-Zehnder interferometer (MZI) including both TOPMs and CDMs,
as illustrated in Figure 1 (a). The TOPMs provide a DC offset to minimise one
of the MZI output intensities. Consequently, small changes in one of the CDMs’
phases will cause a large change in intensity of this output arm with a high
extinction ratio of $\sim$25 dB without requiring a full $\pi$ phase change.
The -3 dB bandwidth of the CDM changes with the biasing conditions, and was
estimated at $\sim$10 GHz and produced $\sim$175 ps full-width-half-maximum
(FWHM) pulses. This allows operation of the devices for use in the Coherent-
One-Way (COW) QKD protocol.

COW is a distributed-phase-reference scheme stucki2005fast  that transmits
pulses in pairs, encoding $|0\rangle$ in the first bin and $|1\rangle$ in the
second. Here we modulate an external CW laser (1550 nm) to generate pulses in
these time-bins and generate a key from the unambiguous time of arrival of the
single photons in each time-bin pair. Security of the channel is maintained by
measuring the visibility from interfering successive photon pulses, when the
pattern of state $|1\rangle$ followed by state $|0\rangle$ occurs, at the
receiver branciard2008upper . A decoy state, with photon pulses in each time-
bin ($|0\rangle$ and $|1\rangle$) is included to increase the probability of
occupied successive pulses allowing a more accurate measurement of the
interference, and to detect photon-number-splitting attacks.

We perform chip-to-chip COW QKD using the integrated SiOxNy receiver device
from Ref. sibson2015chip  as illustrated in Figure 1 (a). The receiver uses a
reconfigurable thermo-optic MZI to route a larger proportion of the input
signal directly to a single photon detector for key generation. The fibre-
coupled superconducting nanowire single photon detectors are biased to
$\sim$40% system detection efficiency and dark count rates of $\sim$500 cps
per detector. A smaller proportion of the signal is routed to an asymmetric
MZI (AMZI) with a 580 ps delay line for the visibility measurement to verify
to security of the channel.

Extending this approach also allows the state preparation of the four BB84
states BB84 . Figure 1 (b) illustrates slow TOPMs inside and outside an MZI
that apply a static phase offset, while we use the non-ideal but fast CDMs in
both arms to apply small state-dependent modulations. As shown in Figure 1
(d), we start by preparing the state $|+\mathrm{i}\rangle$ (red dot on the
Bloch sphere), using the TOPMs as a DC offset, with near unit fidelity. We
then use four non-ideal CDMs (two inside and two outside of the MZI) to
prepare any of the four BB84 states
$\\{|0\rangle,|1\rangle,|+\rangle,|-\rangle\\}$ by shifting each of the fast
CDMs by $\pi/2$ respectively. This approach limits each CDM to a $\pi/2$
phase-shift thus minimising the phase dependent transmission across all four
BB84 states and allowing state preparation with non-ideal phase modulators.

The technique can be applied to polarisation encoding as illustrated in Figure
1 (b). We couple an intensity modulated (IM) laser to launch 175 ps FWHM
optical pulses into a path encoded MZI. Applying static biases to the TOPMs in
the MZI prepares the $|+\mathrm{i}\rangle$ state, and using each of the CDMs
generates one of the four BB84 states as above. Finally this path encoded
state is combined on a 2D grating coupler that acts as a path-to-polarisation
converter wang2016chip .

A fibre-based receiver [shown under the transmitter chip in Figure 1 (b)] was
constructed to decode and measure the output polarisation states. It consisted
of a 50:50 beam splitter (BS) to passively choose the measurement basis, with
one arm connected to a polarisation beam-splitter (PBS) to measure $|0\rangle$
and $|1\rangle$, and one arm coupled to a PBS to measure $|+\rangle$ and
$|-\rangle$. Polarisation controllers (PC) were then used to set the
measurement bases in each arm. Finally, the outputs were coupled to
superconducting nanowire single photon detectors.

Figure 3: Estimated Secret Key Rates: The main data sets (squares, circles,
and diamonds) were collected by emulating a quantum channel with the use of a
variable optical attenuator and assuming standard fibre losses; however, the
data shown with asterisks was collected using a 20 km fibre spool as the
quantum channel. (a) Raw and secure key rates using the chip to implement the
COW QKD protocol, as illustrated in Figure 1 (a). The system operates with a
1.72 GHz clock-rate with a QBER of 1.01% and estimated secure key rate of 916
kbps over a 20 km fibre. (b) Raw and secure key rates using the chip to
produce polarisation encoded BB84 states, as illustrated in Figure 1 (b). We
measure a low QBER of 1.1% while the transmitter is operated with a 1 GHz
clock-rate which yields an estimated secure key rate of 329 kbps over a 20 km.
(c) Histogram measurements of the time-bin encoded BB84 state preparation and
measurement, as illustrated in Figure 1 (c).

Polarisation encoded systems are commonly used in free-space links owing to
the non-birefringent nature of the atmosphere. However, in fibre-based
networks the birefringence of the fibre induces a time-dependent polarisation
rotation on any transmitted qubits, requiring active compensation to maintain
reference frame alignment xavier2009experimental . Instead, encoding qubits in
time and phase, so called time-bin encoding, when transmitting over fibre
optics is a much more stable degree of freedom. Here we show the combined TOPM
and CDM approach for state preparation described above can be translated to
the equivalent time-bin encoded device.

A schematic of the integrated circuit used to produce time-bin qubits is
illustrated in Figure 1 (c). An external laser (1550 nm) is intensity
modulated (IM) to generate weak coherent pulses (350 ps FWHM), which are
coupled into the silicon chip containing an AMZI with a 1.5 ns delay line to
separate photon pulses into one of two time-bins. Here, $|0\rangle$ and
$|1\rangle$ is encoded by a photon in the first or second time-bin,
respectively, and $|+\rangle$ and $|-\rangle$ encoded by a photon in a
superposition of the first and second time-bin with $0$ and $\pi$ relative
phase, respectively. In future designs, the external LiNbO3 IM can be replaced
with an on-chip MZI as demonstrated in the COW protocol, as they showed
comparable performance with extinction ratios of 25-30 dB. Although the
modulation bandwidth and loss performance of the on-chip MZI is inferior to
the external LiNbO3 IM (10 GHz rather than 40 GHz bandwidth, and 5 dB loss
rather than 4 dB insertion loss), they are still sufficient for QKD and gain
the important advantage of monolithic integration in a standard silicon
photonic platform.

An MZI is included in the top arm and contains a TOPM to apply a loss
balancing correction to ensure each time-bin contains equal photon amplitudes
exiting the device. The inside of the AMZI includes TOPMs which can control
the DC relative phase between the separate pulses and CDMs to alternate
between the $|+\rangle$ and $|-\rangle$ states. The final beam splitter of the
AMZI is replaced with a further MZI, with the TOPMs biasing an equal
superposition of the two time-bins, and the two CDMs used to select the first
or second time-bin coupled into a single fibre output.

A second identical chip [pictured at the bottom right of Figure 1 (c)] is used
as a receiver circuit with a matched AMZI used to passively select basis
measurements before the output is fibre coupled and measured with
superconducting nanowire single photon detectors. The phase decoding AMZI
overlaps successive time-bins creating three possible time-slots within which
to detect photons. Phase information is interfered in the middle time-slot
allowing measurements in the $\\{|+\rangle,|-\rangle\\}$ basis, whereas time
of arrival information in the first and third time-slots allows measurement in
the $\\{|0\rangle,|1\rangle\\}$ basis.

Figure 3 illustrates the results from our three implementations of high-speed
low-error QKD with integrated silicon photonics. Figure 3 (a) shows the raw
and secret key rates, and the QBER from the system performing the COW QKD
protocol. Here, pulse modulation provides 175 ps FWHM pulses with a high
$\sim$25 dB extinction ratio between bright and empty pulses. The system
operates with a 1.72 GHz clock-rate (or 0.86 GHz system-rate, as one state is
sent every two clock cycles) with a QBER of 1.01% and estimated asymptotic
secure key rate of 916 kbps over a 20 km fibre, following the upper-bound
security proof against collective attacks of Branciard _et al._
branciard2008upper .

For the polarisation and time-bin BB84 QKD protocols, we model and fit the
data from our TOPMs and CDMs in Figure 2 (a) and (b), resulting in an expected
state preparation fidelity of 99.5%, which yields an expected QBER of
$\leq$1.1%. This is equivalent to a 19.5 dB extinction ratio between measuring
the state $|\psi\rangle_{i}$ and its orthogonal counterpart
$|\bar{\psi}\rangle_{i}$ (e.g. measuring $|1\rangle$ when preparing
$|0\rangle$ or measuring $|-\rangle$ when preparing $|+\rangle$).

Figure 3 (b) shows the measured raw and secret key rates as well as the QBER
from operating the transmitter for polarisation encoding and using the passive
fibre-based receiver detection scheme described. We measure a low QBER of 1.1%
while the transmitter is operated with a 1 GHz clock-rate which yields an
estimated asymptotic secure key rate of 329 kbps over a 20 km fibre using a
non-phase randomised weak coherent BB84 security proof without decoy states
against general attacks lo2007security . The rate and maximum secure distance
could be increased drastically with the addition of an extra integrated
intensity modulator to produce decoy states 2005Decoyproof .

Finally, measurements of the time-bin encoded states, using a second identical
chip as the receiver circuit as described above, are shown in Figure 3 (c).
Analysing these measurements, we observe a low QBER of $\sim$2.1% for this
proof-of-principle demonstration of time-bin encoded BB84 state preparation.
Future systems will benefit from a dedicated low-loss silicon receiver circuit
by minimising the fibre-to-chip coupling loss (currently $\sim$-4.5 dB) and
reducing the loss incurred in the on-chip delay; e.g., reducing the 1.5 ns
delay used here to the 600 ps used in COW would increase transmission by $>$3
dB. As we implement a passive basis selection scheme there is not a need for
fast modulation on the receiver, and removing the _p-n_ carrier depletion
modulators would further decrease absorption by $\sim$5 dB per modulator.

In conclusion, this work experimentally demonstrates the feasibility of high-
speed QKD transmitters in CMOS-based silicon photonic integrated circuits. In
particular, we show an approach to overcome the problems of high-fidelity
state preparation when using non-ideal fast modulation in standard silicon
photonic fabrication. Using a combination of slow, but ideal TOPMs alongside
high-bandwidth ($\sim$10 GHz), but non-ideal, CDMs we demonstrate QKD state
preparation and pulse modulation. We show three successful implementations:
time-bin encoded BB84 state preparation and measurement, polarisation encoded
BB84 (1 GHz clock-rate, 1.1% QBER, 329 kbps estimated asymptotic secret key
rate), and pulse modulation for COW QKD (1.72 GHz clock-rate, 1.01% QBER, 916
kbps estimated asymptotic secure key rate) over a 20 km fibre link.

The modulation bandwidth and estimated secret key rates generated in this
demonstration are comparable to previous integrated photonic QKD
demonstrations sibson2015chip , but with less functionality combined on a
single device and currently relying on off-chip laser sources. Future
generations of silicon based chips will benefit from the recent developments
of low-loss couplers cardenas2014high , low-loss delay lines cardenas2009low ,
integrated laser sources liang2010recent , and integrated single photon
detectors marsili2013detecting , allowing high-performing monolithically
integrated transmitter and receiver devices. Performance could be further
improved by combining the functionality demonstrated, such as integrated pulse
modulation, intensity modulators for decoy-state 2005Decoyproof , and
attenuation calibration, and by increasing the rates beyond the 10 GHz
bandwidth demonstrated here reed2010silicon . The ease of fabrication and
availability of silicon photonics (in comparison with other integrated
photonic platforms such as InP and SiOxNy) will also open routes to mass-
manufacture using standard CMOS fabrication tools and foundries.

Future demonstrations will require focus on the complete system for autonomous
QKD operation deployed in telecommunication networks. This includes the
development of real-time basis and bit selection using quantum random numbers,
active basis alignment, appropriate error reconciliation, privacy
amplification, and the use of finite-key analysis to qualify the security.

Ultimately, integrated silicon photonics will allow the manufacture of quantum
communication chips with electronic and photonic processing on a single
monolithic device and the reduced footprint will enable further multiplexing,
complexity and operation with single photon detection. The ability to scale up
these integrated circuits and incorporate microelectronics opens the way to
new and advanced integrated quantum communication technologies and larger
adoption of quantum-secured communications.

During the preparation of this manuscript, the authors became aware of a
complementary demonstration by Ma _et al._ Ma2016  in silicon using carrier
injection _p-i-n_ diode phase-shifters to encode BB84 states in polarisation
and micro-ring resonators for pulse modulation.

## Acknowledgments

The authors thank IMEC for the fabrication of the integrated photonic silicon
devices and LioniX for the fabrication of the silicon oxynitride receiver.

This work was supported by the Engineering and Physical Sciences Research
Council (EPSRC), The European Research Council (ERC), FP7 Action: Beyond the
Barriers of Optical Integration (BBOI), and the UK Quantum Communications Hub.
J.L.O’B. acknowledges a Royal Society Wolfson Merit Award and a Royal Academy
of Engineering Chair in Emerging Technologies. M.G.T acknowledges fellowship
support from the Engineering and Physical Sciences Research Council (EPSRC,
UK).

## References

  * (1) H.-K. Lo, M. Curty, and K. Tamaki, “Secure quantum key distribution,” Nature Photonics, vol. 8, no. 8, pp. 595–604, 2014.
  * (2) V. Scarani, H. Bechmann-Pasquinucci, N. Cerf, M. Dušek, N. Lütkenhaus, and M. Peev, “The security of practical quantum key distribution,” Rev. Mod. Phys., vol. 81, pp. 1301–1350, Sept. 2009.
  * (3) K.-i. Yoshino, T. Ochi, M. Fujiwara, M. Sasaki, and A. Tajima, “Maintenance-free operation of WDM quantum key distribution system through a field fiber over 30 days,” Opt. Exp., vol. 21, no. 23, pp. 31395–31401, 2013.
  * (4) B. Korzh, N. Walenta, R. Houlmann, and H. Zbinden, “A high-speed multi-protocol quantum key distribution transmitter based on a dual-drive modulator,” Opt. Exp., vol. 21, no. 17, pp. 19579–19592, 2013.
  * (5) A. R. Dixon, J. F. Dynes, M. Lucamarini, B. Fröhlich, A. W. Sharpe, A. Plews, S. Tam, Z. L. Yuan, Y. Tanizawa, H. Sato, S. Kawamura, M. Fujiwara, M. Sasaki, and A. J. Shields, “High speed prototype quantum key distribution system and long term field trial,” Opt. Exp., vol. 23, no. 6, pp. 7583–7592, 2015.
  * (6) M. Sasaki, M. Fujiwara, R.-B. Jin, M. Takeoka, H. Endo, K.-I. Yoshino, T. Ochi, S. Asami, and A. Tajima, “Quantum Photonic Network: Concept, Basic Tools, and Future Issues,” IEEE Journal of Selected Topics in Quantum Electronics, vol. 21, no. 3, pp. 49–61, 2015.
  * (7) M. G. Thompson, A. Politi, J. C. Matthews, and J. L. O’Brien, “Integrated waveguide circuits for optical quantum computing,” IET circuits, devices & systems, vol. 5, no. 2, pp. 94–102, 2011.
  * (8) A. E.-J. Lim, J. Song, Q. Fang, C. Li, X. Tu, N. Duan, K. K. Chen, R. P.-C. Tern, and T.-Y. Liow, “Review of silicon photonics foundry efforts,” Selected Topics in Quantum Electronics, IEEE Journal of, vol. 20, no. 4, pp. 405–416, 2014.
  * (9) J. W. Silverstone, D. Bonneau, J. L. O’Brien, and M. G. Thompson, “Silicon Quantum Photonics,” IEEE Journal of Selected Topics in Quantum Electronics, vol. 22, no. 6, pp. 1–13, 2016.
  * (10) W. H. P. Pernice, C. Schuck, O. Minaeva, M. Li, G. N. Goltsman, A. V. Sergienko, and H. X. Tang, “High-speed and high-efficiency travelling wave single-photon detectors embedded in nanophotonic circuits.,” Nature communications, vol. 3, p. 1325, jan 2012.
  * (11) F. Najafi, J. Mower, N. C. Harris, F. Bellei, A. Dane, C. Lee, X. Hu, P. Kharel, F. Marsili, S. Assefa, K. K. Berggren, and D. Englund, “On-chip detection of non-classical light by scalable integration of single-photon detectors,” Nature communications, vol. 6, 2015.
  * (12) C. R. Doerr, “Silicon photonic integration in telecommunications,” Frontiers in Physics, vol. 3, p. 37, 2015.
  * (13) C. Sun, M. T. Wade, Y. Lee, J. S. Orcutt, L. Alloatti, M. S. Georgas, A. S. Waterman, J. M. Shainline, R. R. Avizienis, S. Lin, B. R. Moss, R. Kumar, F. Pavanello, A. H. Atabaki, H. M. Cook, A. J. Ou, J. C. Leu, Y.-H. Chen, K. Asanović, R. J. Ram, M. A. Popović, and V. M. Stojanović, “Single-chip microprocessor that communicates directly using light,” Nature, vol. 528, no. 7583, pp. 534–538, 2015.
  * (14) P. Zhang, K. Aungskunsiri, E. Martín-López, J. Wabnig, M. Lobino, R. W. Nock, J. Munns, D. Bonneau, P. Jiang, H. W. Li, A. Laing, J. G. Rarity, A. O. Niskanen, M. G. Thompson, and J. L. O’Brien, “Reference-Frame-Independent Quantum-Key-Distribution Server with a Telecom Tether for an On-Chip Client,” Phys. Rev. Lett., vol. 112, no. 13, p. 130501, 2014.
  * (15) A. Tanaka, M. Fujiwara, K.-i. Yoshino, S. Takahashi, Y. Nambu, A. Tomita, S. Miki, T. Yamashita, Z. Wang, M. Sasaki, and A. Tajima, “High-speed quantum key distribution system for 1-Mbps real-time key generation,” Quantum Electronics, IEEE Journal of, vol. 48, no. 4, pp. 542–550, 2012.
  * (16) P. Sibson, C. Erven, M. Godfrey, S. Miki, T. Yamashita, M. Fujiwara, M. Sasaki, H. Terai, M. G. Tanner, C. M. Natarajan, R. H. Hadfield, J. L. O’Brien, and M. G. Thompson, “Chip-based Quantum Key Distribution,” arXiv preprint arXiv:1509.00768, September 2015.
  * (17) G. T. Reed, G. Mashanovich, F. Gardes, and D. Thomson, “Silicon optical modulators,” Nature photonics, vol. 4, no. 8, pp. 518–526, 2010.
  * (18) C. H. Bennett, “Quantum cryptography: Public key distribution and coin tossing,” in Proceesings of IEEE International Conference on Computer SystemS and Signal Processing, pp. 175–179, 1984.
  * (19) Z. Fang and C. Z. Zhao, “Recent progress in silicon photonics: a review,” ISRN Optics, vol. 2012, 2012.
  * (20) P. P. Absil, P. De Heyn, H. Chen, P. Verheyen, G. Lepage, M. Pantouvaki, J. De Coster, A. Khanna, Y. Drissi, D. Van Thourhout, and J. Van Campenhout, “Imec iSiPP25G silicon photonics: a robust CMOS-based photonics technology platform,” in Proc. SPIE 9367, p. 93670V, International Society for Optics and Photonics, 2015.
  * (21) D. Stucki, N. Brunner, N. Gisin, V. Scarani, and H. Zbinden, “Fast and simple one-way quantum key distribution,” Appl. Phys. Lett., vol. 87, no. 19, p. 194108, 2005.
  * (22) C. Branciard, N. Gisin, and V. Scarani, “Upper bounds for the security of two distributed-phase reference protocols of quantum cryptography,” New J. Phys., vol. 10, no. 1, p. 013031, 2008.
  * (23) J. Wang, D. Bonneau, M. Villa, J. W. Silverstone, R. Santagati, S. Miki, T. Yamashita, M. Fujiwara, M. Sasaki, H. Terai, M. G. Tanner, C. M. Natarajan, R. H. Hadfield, J. L. O’Brien, and M. G. Thompson, “Chip-to-chip quantum photonic interconnect by path-polarization interconversion,” Optica, vol. 3, no. 4, pp. 407–413, 2016.
  * (24) G. Xavier, N. Walenta, G. V. De Faria, G. Temporão, N. Gisin, H. Zbinden, and J. Von der Weid, “Experimental polarization encoded quantum key distribution over optical fibres with real-time continuous birefringence compensation,” New Journal of Physics, vol. 11, no. 4, p. 045015, 2009\.
  * (25) H.-K. Lo and J. Preskill, “Security of Quantum Key Distribution using Weak Coherent States with Non-Random Phases,” Quantum Information and Computation, vol. 7, no. 5&6, pp. 431–458, 2007.
  * (26) X. Ma, B. Qi, Y. Zhao, and H.-K. Lo, “Practical decoy state for quantum key distribution,” Phys. Rev. A, vol. 72, p. 012326, July 2005\.
  * (27) J. Cardenas, C. B. Poitras, K. Luke, L.-W. Luo, P. A. Morton, and M. Lipson, “High coupling efficiency etched facet tapers in silicon waveguides,” IEEE Photonics Technol. Lett, vol. 26, no. 23, pp. 2380–2382, 2014.
  * (28) J. Cardenas, C. B. Poitras, J. T. Robinson, K. Preston, L. Chen, and M. Lipson, “Low loss etchless silicon photonic waveguides,” Optics express, vol. 17, no. 6, pp. 4752–4757, 2009.
  * (29) D. Liang and J. E. Bowers, “Recent progress in lasers on silicon,” Nature Photonics, vol. 4, no. 8, pp. 511–517, 2010.
  * (30) F. Marsili, V. B. Verma, J. A. Stern, S. Harrington, A. E. Lita, T. Gerrits, I. Vayshenker, B. Baek, M. D. Shaw, R. P. Mirin, and S. W. Nam, “Detecting single infrared photons with 93% system efficiency,” Nature Photonics, vol. 7, no. 3, pp. 210–214, 2013.
  * (31) C. Ma, W. D. Sacher, Z. Tang, J. C. Mikkelsen, Y. Yang, F. Xu, H.-K. Lo, and J. K. S. Poon, “Integrated Silicon Photonic Transmitter for Polarization-Encoded Quantum Key Distribution,” arXiv preprint arXiv:1606.04407, June 2016.


# Contextualized Policy Recovery: Modeling and Interpreting Medical Decisions
with Adaptive Imitation Learning

Jannik Deuschel
Karlsruhe Institute of Technology
Carnegie Mellon University
<EMAIL_ADDRESS>
&Caleb N. Ellington11footnotemark: 1
Carnegie Mellon University
<EMAIL_ADDRESS>
&Benjamin J. Lengerich
Massachusetts Institute of Technology
Broad Institute of MIT and Harvard
<EMAIL_ADDRESS>
&Yingtao Luo
Carnegie Mellon University
<EMAIL_ADDRESS>
&Pascal Friederich
Karlsruhe Institute of Technology
<EMAIL_ADDRESS>
&Eric P. Xing
Mohamed bin Zayed University of Artificial Intelligence
Carnegie Mellon University
Petuum, Inc.
<EMAIL_ADDRESS>
Equal contributionCorresponding author

###### Abstract

Interpretable policy learning seeks to estimate intelligible decision policies
from observed actions; however, existing models fall short by forcing a
tradeoff between accuracy and interpretability. This tradeoff limits data-
driven interpretations of human decision-making process. e.g. to audit medical
decisions for biases and suboptimal practices, we require models of decision
processes which provide concise descriptions of complex behaviors.
Fundamentally, existing approaches are burdened by this tradeoff because they
represent the underlying decision process as a universal policy, when in fact
human decisions are dynamic and can change drastically with contextual
information. Thus, we propose Contextualized Policy Recovery (CPR), which re-
frames the problem of modeling complex decision processes as a multi-task
learning problem in which complex decision policies are comprised of context-
specific policies. CPR models each context-specific policy as a linear
observation-to-action mapping, and generates new decision models on-demand as
contexts are updated with new observations. CPR is compatible with fully
offline and partially observable decision environments, and can be tailored to
incorporate any recurrent black-box model or interpretable decision model. We
assess CPR through studies on simulated and real data, achieving state-of-the-
art performance on the canonical tasks of predicting antibiotic prescription
in intensive care units ($+22\%$ AUROC vs. previous SOTA) and predicting MRI
prescription for Alzheimer’s patients ($+7.7\%$ AUROC vs. previous SOTA). With
this improvement in predictive performance, CPR closes the accuracy gap
between interpretable and black-box methods for policy learning, allowing
high-resolution exploration and analysis of context-specific decision models.

##  1 Introduction

Interpretable policy learning (Hüyük et al., 2022) seeks to recover an
underlying decision-making process from a dataset of demonstrated behavior,
and represent this process as an interpretable model that can be quantified,
audited, and intuitively understood. This approach has gained considerable
attention in the medical informatics community as a promising approach for
improving standards of care by detecting bias, explaining sub-optimal outcomes
(Lengerich et al., 2022), and quantifying regional (McKinlay et al., 2007) and
institutional (Westert et al., 2018) differences. Classic machine learning
algorithms for policy inference are based on inverse reinforcement learning
(Ng & Russell, 2000) or imitation learning (Bain & Sammut, 1999; Piot et al.,
2014; Ho & Ermon, 2016), and use black-box architectures such as recurrent
neural networks. These approaches have been applied in various medical
domains, most prominently oncology prognosis (Beck et al., 2011; Esteva et
al., 2017). However, black-box methods have been met with skepticism from the
medical community based on a lack of interpretability as well as an inability
to identify catastrophic failure modes and generalization issues (Laï et al.,
2020; Royal Society and Royal Society (Great Britain) Staff(2017), Great
Britain).

Figure 1: CPR uses patient-specific treatment contexts to estimate the agent’s
decision model at each timestep. Decision models are context-specific, and
each action probability is an interpretable linear combination of observed
features. In this way, CPR achieves exact model-based interpretability without
sacrificing representational capacity.

To achieve this desire for interpretable policies, there has been a recent
surge in transparent policy parametrizations for imitation learning. Recent
approaches include recurrent decision trees (Pace et al., 2022), visual
decision boundaries (Hüyük et al., 2022), high-level programming syntax (Verma
et al., 2018) or outcome preferences (Yau et al., 2020). These approaches are
generally more interpretable to clinicians, but their interpretability stems
from restrictive modeling architectures, sacrificing performance or imposing
obscure constraints that make real-world applications challenging. The primary
challenge is that human decisions are informed by a variety of factors
including patient background, medical history, lab tests, and more, and true
human decision processes are complex. Thus, compressing the decision policy
into a single universal observation-to-action mapping necessitates the use of
large nonparametric models (e.g. neural nets) that preclude direct
interpretability, or produce a model which fails to capture the complexities
of human decision-making. Succinctly, models must be both accurate and
interpretable to effectively support clinical decisions.

In this paper, we propose Contextual Policy Recovery (CPR). Instead of seeking
a universal policy that necessitates trading off accuracy against
interpretability, CPR embraces the wealth of contextual information that
guides human decision-making and reframes the problem of policy learning as
multi-task learning of interpretable _context-specific_ policies. CPR learns a
black-box generator function that encodes contextual information (e.g. the
history of observed symptoms and actions which have previously occurred in the
decision process) and generates linear observation-to-action mappings. With
this combination of black-box and glass-box components, CPR provides
interpretable context-specific decision functions at sample-specific
resolution. CPR does not sacrifice representational capacity and achieves
state-of-the-art performance in policy recovery. Finally, CPR is a modular
framework: both the model family of the context encoder and the interpretable
observation-to-action mapping model can be chosen according to the given task
at hand.

Contributions: Our work makes the following contributions to the personalized
modeling and medical machine learning communities:

  * •

We propose CPR, a framework for estimating time-varying and context-dependent
policies as linear observation-to-action mappings, operating in a fully
offline and partially observable environment. CPR enables interpretable and
personalized imitation learning, dynamically incorporating new information to
model decisions over the entire course of treatment.

  * •

We apply CPR to two canonical medical imitation learning tasks: predicting
antibiotic prescription in intensive care units and predicting MRI
prescription for Alzheimer’s patients. CPR matches the performance of black-
box models, generating decision models that recover best-practice treatment
policies under a continuum of patient contexts.

  * •

We simulate a heterogeneous Markov decision process and who that CPR
empiriclaly converges to the true decision model parameters, improving
representation, performance, and interpretability over black-box-only models.

##  2 Related work

Policy class | Contextualized | Global Tree | Global Aggregate
---|---|---|---
Representative model | CPR | POETREE | INTERPOLE
$\mathbb{P}(a_{t}|x_{0:t},a_{0:t-1})$ | $f_{h(x_{0:t-1},a_{0:t-1})}(x_{t})$ | $f_{\theta}(x_{t},h(x_{0:t-1}))$ | $f_{\theta}(h(x_{0:t},a_{0:t-1}))$
|  |  |
Adapts to observed actions | ✓ | ✗ | ✓
$x_{t}\rightarrow a_{t}$ glass-box | ✓ | ✓ | ✗
Table 1: Comparison of different policy learning algorithms. All listed model
classes are designed to provide interpretability beyond what a black-box
recurrent model would provide: CPR provides contextualized parametric
policies, POETREE (Pace et al., 2022) provides a global tree-based policy, and
INTERPOLE (Hüyük et al., 2022) provides interpretations of belief states in
$h$ as a summary of patient history.

We seek to learn, at each timestep, an interpretable parametrization of
observed behavior to understand how an agent’s action was taken in a partially
observable, offline environment expanding on the objective of classical
imitation learning, that solely seeks to replicate demonstrated behavior. CPR
combines the strengths of previous work (Table 1) by keeping the observation-
to-action mapping at each timestep interpretable while being able to adapt to
the full observed past, aligning learned policies closer with demonstrated
behavior. While POETREE (Pace et al., 2022) is able to carry over a hidden
state through time, this hidden state acts as an additive bias at each
timestep instead of adapting the underlying model parameters $\theta$, which
are static throughout time.

##### Imitation Learning

The classical approach for learning sequential decision-making processes
involves reinforcement learning, optimizing an agent’s reward signal $R$ in an
online environment. However, in some applications, such as clinical decision-
making, experiments with online policies would be both unethical and
impractical so only observational data is available. In this setting we do not
have access to the reward signal $R$ and instead focus on the inverse task of
replicating the observed behavior of an agent, known as imitation learning.
There are several approaches to tackle such problems, including simple
behavioral cloning, without taking interpretability into account, in which the
task is reduced to supervised learning mapping observations to actions (Bain &
Sammut, 1999; Piot et al., 2014; Sun et al., 2017). Other approaches are based
on distribution matching where adversarial training is used to match the
state-action distributions between demonstrator and learned policies (Ho &
Ermon, 2016; Jeon et al., 2018; Kostrikov et al., 2020). Adaptations of
inverse reinforcement learning have been proposed for partial observability
(Choi & Kim, 2011) and offline learning (Makino & Takeuchi, 2012). These
approaches solely utilize black-box models to maximize action-matching
performance, making it challenging to distill a transparent and tractable
description of the learned policies.

##### Interpretable Policy Learning

Recurrent neural networks depend on latent hidden representations for each
feed-forward pass in a time series, obfuscating popular post-hoc explanation
methods like LIME (Ribeiro et al., 2016) and SHAP (Lundberg & Lee, 2017).
Instead, recent interpretable policy learning methods have focused on adding
limiting assumptions, although this also limits real world applicability.
INTERPOLE (Hüyük et al., 2022) is a notable example, which parameterizes a
latent belief space that relates to decisions, but falls short of explaining
how this belief relates to prior observed information. Another notable method
is POETREE (Pace et al., 2022), which parameterizes policies as a recurrent
soft decision trees (Frosst & Hinton, 2017). While this approach results in a
human-interpretable model which maps between timeseries observations and
actions, POETREE requires significant post-processing to remove
uninterpretable components critical to training. This post-processing step
sacrifices performance, especially if the observation space is high-
dimensional. Further, these methods are often less accurate at decision
modeling than workhorse models like logistic regression (Table 1).

##  3 Methods

###  3.1 Preliminaries: Contextualized Modeling

Given a dataset consisting of targets $y\in Y$, observations $x\in X$ and
context $c\in C$, with the corresponding random variables denoted as
$\mathbf{Y}$, $\mathbf{X}$ and $\mathbf{C}$, we want to learn a model
$\mathbb{P}(\mathbf{Y}|x,c,\theta)$ predicting $y$ from $x$ and $c$. With
this, the probabilistic model is defined as follows:

$\displaystyle
y\sim\mathbb{P}(\mathbf{Y}|x,\theta),\quad\theta\sim\mathbb{P}(\theta|c),\quad\mathbb{P}(\mathbf{Y}|x,c)\sim\int_{\theta}\mathbb{P}(\mathbf{Y}|x,\theta)\mathbb{P}(\theta|c)d\theta$

as initially described by Al-Shedivat et al. (2020). This formulation was
generalized by Lengerich et al. , allowing us to model $\mathbb{P}(\theta|c)$
by any black-box model while keeping $\mathbb{P}(\mathbf{Y}|x,\theta)$ in a
simple model class parametrized by $\theta$.

Due to an increase in dataset complexity, heterogeneity, and size, sample-
specific inference has driven interest in many application areas (Ageenko et
al., 2010; Buettner et al., 2015; Fisher et al., 2018; Hart, 2016; Ng et al.,
2015). Contextualized modeling has been used in several different frameworks
to estimate the context-specific parameters $\theta$ using a meta-model
$\mathbb{P}(\theta|C)$ that relates contextual information $C$ to variation in
$\theta$. Typically, $\mathbb{P}(\theta|C)$ is a Diroc delta function
$\delta(\theta-h(C))$, where $h(C)$ is a deterministic context encoding. The
context encoder $h$ is parameterized as $h(C)=Z^{T}Q$, where $Q$ controls the
rank of the context-specific model space and $Z$ is a deterministic encoding
$Z=g_{\phi}(C)$. Many well-known contextualized models fit this framework: in
varying coefficient models (Hastie & Tibshirani, 1993), $g$ is a linear model
and $Q$ is the identity matrix, in contextualized explanation networks (Al-
Shedivat et al., 2020), $g$ is a deep neural network, and $\lVert Z\rVert=1$.
In CPR, $g$ is a differentiable, recurrent history encoder, and $Q$ is the
identity matrix.

###  3.2 Contextualized Policy Recovery

CPR builds on recent developments (Pace et al., 2022) in interpretable,
offline policy learning. Let dataset
$D=\\{(x_{1}^{i},a_{1}^{i}),...,(x_{T_{i}}^{i},a_{T_{i}}^{i})\\}_{i=1}^{N}$
consist of $N$ treatment trajectories, where each patient $i$ is observed over
$T_{i}$ discrete timesteps for symptoms $x\in\mathcal{X}$ and physician
actions $a\in\mathcal{A}$. The data is generated by an unknown policy of the
physician agent $\mathbb{P}(a_{t}|x_{1},a_{1},...,x_{t-1},a_{t-1},x_{t})$
where the action probability at time $t$ is a function of the agent’s current
state, which is defined by the current and past patient symptoms, as well as
past actions.

To recover a policy that is both tractable and interpretable, CPR makes a
practical and clinically-minded assumption that a physician places the highest
importance on the most current patient information when deciding an action.
While treatment history is important, this history is primarily useful for
placing the patient’s current disease presentation within a context of disease
progression and past treatment attempts. To represent this information
hierarchy, CPR leverages contextual and historical features to generate
context-specific policy models.

$\displaystyle
P(a_{t}|x_{1},a_{1},...,x_{t-1},a_{t-1},x_{t}):=f_{\theta_{t}}(a_{t}|x_{t})$
$\displaystyle\theta_{t}:=g(x_{1},a_{1},...,x_{t-1},a_{t-1})$

Where $f$ is an interpretable model class, e.g. logistic regression,
parameterized by a context-specific $\theta$, and $\theta$ is generated via a
historical context encoder $g$. The effects of current observation $x_{t}$ on
action probabilities $a_{t}$ are directly explained through the simple
context-specific model $f_{\theta}$. Furthermore, $g$ can take any functional
form without precluding the interpretability of $f$. The context-specific
policy models $f_{\theta_{t}^{i}}$ are generated for each patient $i$ at each
timepoint $t$, allowing us to investigate how previous actions, previous
symptoms, patient covariates, and treatment time influence the policy. CPR
flexibly allows the context encoder $g$ and the observation-to-action function
$f$ to be freely chosen, although they must be differentiable to allow for
joint optimization under an appropriate loss $\ell$.

$\min_{g}\frac{1}{N}\sum_{i}\frac{1}{T_{i}}\sum_{t}\frac{1}{|\mathcal{A}|}\sum_{\widehat{a}\in\mathcal{A}}\ell(a_{t}^{i},f_{g(x_{1}^{i},a_{1}^{i},...,x_{t-1}^{i},a_{t-1}^{i})}(\widehat{a},x_{t}^{i}))$

In our experiments, $g$ is parametrized by either a vanilla RNN or LSTM
(Hochreiter & Schmidhuber, 1997), $f$ is a logistic function, actions
$\mathcal{A}:=\\{0,1\\}$ are binary, and $\ell$ is binary cross-entropy loss.
Finally, CPR applies a lasso regularizer to $\theta$ to learn robust policy
parameters.

##  4 Experiments

We apply CPR to recover time-varying, context-specific decision models within
complex decision-making processes. First, we evaluate CPR on real MRI
prescription data for dementia patients and antibiotic prescription data in
intensive care units. Follow-up analysis of the contextualized models reveals
best-practice treatment plans for both common and outlier patients, while also
recovering unexpected and meaningful heterogeneity in physician policies. We
further ensure CPR’s ability to recover true policy models through a simulated
heterogeneous Markov decision process.

###  4.1 Medical Datasets

| ADNI MRI scans | MIMIC Antibiotics
---|---|---
Algorithm | AUROC | AUPRC | Brier $\downarrow$ | AUROC | AUPRC | Brier $\downarrow$
$\square$ Logistic regression | $0.66\pm 0.01$ | $0.86\pm 0.00$ | $0.16\pm 0.00$ | $0.57\pm 0.01$ | $0.80\pm 0.01$ | $0.20\pm 0.00$
$\square$ INTERPOLE $\dagger$ | $0.60\pm 0.04$ | $0.81\pm 0.08$ | $0.17\pm 0.05$ | NR | NR | NR
$\square$ INTERPOLE $\ddagger$ | $0.44\pm 0.04$ | $0.75\pm 0.09$ | $0.19\pm 0.07$ | $0.65\pm(\leq 0.04)$ | NR | $0.21\pm(\leq 0.04)$
$\square$ POETREE $\ddagger$ | $0.62\pm 0.01$ | $0.82\pm 0.01$ | $0.18\pm 0.01$ | $0.68\pm(\leq 0.04)$ | NR | $0.19\pm(\leq 0.04)$
$\square$ CPR-RNN (ours) | $\mathbf{0.72\pm 0.01}$ | $\mathbf{0.88\pm 0.01}$ | $\mathbf{0.15\pm 0.00}$ | $\mathbf{0.82\pm 0.00}$ | $\mathbf{0.90\pm 0.00}$ | $\mathbf{0.14\pm 0.00}$
$\square$ CPR-LSTM (ours) | $\mathbf{0.72\pm 0.01}$ | $\mathbf{0.88\pm 0.01}$ | $\mathbf{0.15\pm 0.00}$ | $\mathbf{0.82\pm 0.00}$ | $\mathbf{0.90\pm 0.00}$ | $\mathbf{0.14\pm 0.00}$
$\blacksquare$ RNN | $0.72\pm 0.01$ | $0.88\pm 0.01$ | $0.15\pm 0.00$ | $0.83\pm 0.00$ | $0.90\pm 0.00$ | $0.13\pm 0.00$
$\blacksquare$ LSTM | $0.71\pm 0.01$ | $0.88\pm 0.01$ | $0.15\pm 0.00$ | $0.84\pm 0.00$ | $0.91\pm 0.00$ | $0.13\pm 0.00$
Table 2: Action-matching performance of imitation learning algorithms. Bolded
values denote the best performance of interpretable models. Open source task-
agnostic models are reported as mean $\pm$ standard error for 10 bootstrap
runs. INTERPOLE and POETREE, which are task-specific or closed-source, are
based on prior reports: $\dagger$ reported by Hüyük et al. (2022), $\ddagger$
reported by Pace et al. (2022). NR: No values reported. $\square$:
interpretable method. $\blacksquare$: black-box method.

We apply CPR to two medical datasets for canonical imitation learning tasks,
ADNI and MIMIC-III. These datasets are a prime example of partially-observable
decision environments in which we are forced to learn from demonstrated
behavior, and where learned policies have the potential to improve clinical
operations. Empirical results show that CPR significantly outperforms other
interpretable baseline models, and even performing on-par with fully black-box
models for both EHR datasets. Low brier scores indicate that CPR is well
calibrated while achieving SOTA AUROC and AUPRC.

####  4.1.1 MIMIC Antibiotics

Figure 2: Exploration of contextualized policies generated by CPR for
predicting antibiotic prescription. (a) Contextualized policies identify prior
antibiotic prescription and (b) treatment time as drivers of treatment
heterogeneity. (c) CPR generates policies that evolve with time and treatment
history, revealing the context-specific importance of patient symptoms toward
future treatments.

We look at 4195 patients in the intensive care unit over up to 6 timesteps
extracted from the Medical Information Mart for Intensive Care III (Johnson et
al., 2016) dataset and predict antibiotic prescription based on 7 observations
- temperature, hematocrit, potassium, white blood cell count (WBC), blood
pressure, heart rate, and creatinine. We removed the hemoglobin feature used
in previous work by Pace et al. (2022) since it is highly correlated (>0.95)
with the hematocrit feature.

##### Contextualized policies reveal and explain heterogeneity in medical
decision processes

To see how decision functions change under different contexts, we compare them
in model space. UMAP embeddings of the coefficient vectors (Figure 2b) reveal
three distinct clusters of decision functions. The rightmost cluster contains
the initial model parameters $\theta_{0}$ for each trajectory. Since there is
no context that could differentiate the agent’s behavior at the initial visit,
these parameters are the same for all patients, and the contextualized models
recover the population estimator. Subsequent to this initial homogeneity,
heterogeneity in decision policies arises. The larger driver of this
heterogeneity is prior antibiotic prescription – patients that previously got
antibiotics are more likely to continue to receive antibiotics, while patients
that did not receive antibiotics are likely to continue to not receive
antibiotics. The lower cluster contains mostly ($99.8\%$) models in which the
patient did get antibiotics in the previous state $t-1$, while the upper
cluster contains patients ($99.3\%$) that did not get antibiotics in $t-1$.
This strong split is only recovered by contextualized policies; global
policies that ignore context fail to identify this heterogeneity (Figure 2a).
We train two models conditioned on their respective contexts, removing even
more variability by limiting observations to the second day in the ICU. The
global model only represents a small part of the population. Conditioning on
the main driver of model heterogeneity (whether or not a patient got
antibiotics in the previous visit) and training individual models for each
case yields models that look like an average over the contextualized models of
the respective clusters.

To uncover typical treatment regimes, we cluster patients into 5 subgroups
over the first 4 days of the ICU stay using hierarchical clustering (Fig. 12).
To identify the drivers of this heterogeneity in decision policies, we examine
the parameters of the decision function as a function of context (Figure 2c).
The most notable difference in parameters is the intercept value which is
positive for the group of patients that get treated with antibiotics and
negative for patients that did not receive antibiotic treatment. This is
consistent across timesteps, and it is likely that a patient is prescribed
antibiotics if they got it the day before. This aligns with medical protocols
that rarely suggest antibiotic treatments briefer 5 days (Guleria et al.,
2019).

Previous work by Pace et al. (2022) and Bica et al. (2021) described a
patient’s temperature and white blood cell count (WBC) as the main drivers of
antibiotic treatment decisions since these are known medical criteria to
counter infections (Masterton et al., 2008). Our results paint a more nuanced
picture, in which heterogeneous coefficients reflect heterogeneous priorities
when designing treatment plans. CPR identifies that the influence of
temperature changes based on prior antibiotic prescription. For patients who
have already been prescribed antibiotics, an infection has already been
detected and the doctor’s decision model shifts towards mitigating the risk of
possible side effects of the antibiotics treatment rather than strictly
considering the benefits of the new treatment. This shift in priorities is
supported by the change in the creatinine coefficient (Figure 2c). High serum
creatinine can be an indicator of impaired kidney function (Gounden et al.,
2023), a possible adverse effect of antibiotics (Khalili et al., 2013); as
such, a high creatinine level decreases the probability of continuing
antibiotic treatment. Finally, we see that increased potassium is associated
with the decision to begin antibiotics (Figure 2c). Electrolyte balance (in
which potassium plays a vital role) is an important factor in infectious
diseases, and reduced sodium (a 1+ ion competitor with potassium) in
particular is known as a marker of viral and bacterial infections (Królicka et
al., 2020).

Figure 3: CPR generates decision models for marginal groups with high
accuracy. Left: Using only a small subgroup of patients making up 7
observations in the training set, CPR identifies elevated creatinine as a
severe risk factor for kidney failure and reassigns patients to a non-
antibiotics treatment plan, while these patients would otherwise be likely to
receive treatment. Right: For the small subgroup of patients below 20 years of
age (with only 9 observations in the held-out set and 44/12 in the
train/validation set), CPR improves drastically in terms of cross-entropy
loss.

##### Contextualized Policies Reveal Outlier Patients

By relating modeling tasks through task-specific contexts, CPR learns to
generate context-specific models even when the number of samples per context
is as small as one, and similarly generalizes to generate models for unseen
contexts. Previously, we assessed these models in aggregate to reveal common
treatment trajectories and best practices (Figures 2b, 9). Here, we
demonstrate how personalized models also reveal small subgroups of patients or
even individuals with outlier policies. In particular, some patient
populations have higher dose tolerances and are amenable to more aggressive
treatment, while others display rare comorbidities and risk factors prohibit
common treatment options. Personalized treatment options are critical,
especially when more common treatment plans need to be avoided.

CPR identifies several of these outlier patients when in the MIMIC antibiotic
prescription dataset (Fig. 3). First, younger patients often have fewer
comorbidities and more robust immune systems, and physicians can be more
confident that antibiotics will not impose any adverse side effects if an
infection is suspected. We observe that contextualized policies recover this
case, and represent treatment for the under 20 age group much more accurately
when antibiotics are prescribed. Second, elevated creatinine is a rare side-
effect of antibiotics but a likely indicator of suboptimal Kidney function and
possible Kidney failure (Gounden et al., 2023). CPR identifies that patients
with elevated creatinine are immediately removed from antibiotics following an
initial prescription, placing them in a treatment cluster characterized by a
lack of antibiotics prescription that would otherwise be unlikely for these
patients. The patient-specific policies produced by CPR provide a novel view
of the treatment process, allowing us to easily identify these rare and
outlier effects in terms of context-specific policy parameters and errors,
revealing nuances in treatment decisions that were missed by prior works.

####  4.1.2 ADNI MRI Scans

Figure 4: Comparison of patient-specific model parameter distributions by age
and gender in visit $t=0$ after incorporating static contexts. Static contexts
help to personalize initial models when no history is available.

Following prior works by Hüyük et al. (2022) and Pace et al. (2022), we apply
CPR to 1605 patients from the Alzheimer’s Disease Neuroimaging Initiative
(ADNI). The canonical task is to predict at each visit whether a Magnetic
Resonance Image (MRI) scan is ordered for cognitive disorder diagnosis (noa,
2018). Patient observations consist of the Clinical Dementia Rating (CDR-SB)
on a severity scale (normal; questionable impairment; severe dementia)
(O’Bryant et al., 2008) and the MRI outcome of the previous visit falling into
4 categories: No MRI scan, below average, average and above average
hippocampal volume.

The small number of discrete context features and the checklist-like time-
independent diagnostic criteria for Alzheimer’s (O’Bryant et al., 2008)
provide a limited view of Alzheimer’s diagnosis that is unlikely to explain
any heterogeneity in clinical decisions. This is reinforced by the fact that a
single logistic regression outperforms all interpretable policy baselines
(Table 2). We introduce a new condition-specific logistic regression baseline,
where we learn a decision model for every set of unique context features at
each timestep. Indeed, this condition-specific model performs nearly as well
as CPR and black-box models, with an AUROC of 0.71. While CPR and the black-
box baselines can still capture dependencies on past actions and observations
(Fig. 8, 9), this seems to confer only marginal modeling improvements for the
canonical ADNI task.

Instead, we reformulate this canonical task to include a new source of
heterogeneity with clinical significance: patient age and gender (Castro-
Aldrete et al., ). Condition-specific logistic regression and both
interpretable baselines (Hüyük et al., 2022; Pace et al., 2022) are unable to
model changes in patient-specific policies over continuous static contexts
like age, but CPR is able to easily incorporate static as well as dynamic
contexts by encoding static contexts into the initial hidden state of the
context encoder $g$.

Figure 4 shows how the estimated policies at $t=0$ differ between four patient
subgroups. We find meaningful heterogeneity in the models generated by CPR,
where age dominates CDRSB coefficients and overall intercept, while gender
dominates hippocampal volume coefficients. Additionally, static contexts
substantially increase the action-matching performance of CPR to $0.763$
AUROC.

###  4.2 Simulations

Figure 5:  Comparing policy models learned by CPR and RNN in terms of the
Pearson’s correlation between estimated and true action probabilities and
context-specific policy coefficients. We choose default simulation arguments
$N=200$, $\sigma_{a}=0$, $T=15$, $\tau=4$, and $\sigma_{\theta}=0$, varying
each parameter individually. We hold out 15% of trajectories at random for
evaluation. Results are the mean and 95% confidence interval from three
randomly initialized and independently simulated data sets.

CPR incorporates both a deep learning component and a statistical modeling
component to introduce a novel mechanism of interpretability, the context-
specific linear policy. Our approach differs substantially from prior
interpretable methods by including a deep learning component. Naturally, we
wonder if CPR’s explicit linear policy representation is key to its
performance and interpretability, or if accurate and robust context-specific
linear policies can also be recovered from black-box policy models using post-
hoc interpretation methods. To test this, we simulate a heterogeneous, action-
dependent Markov decision process (MDP) and evaluate CPR versus black-box
baselines on their recovery of true simulation parameters: the true action
probability and the true coefficients of a context-specific linear policy
(Fig. 5). While CPR explicitly generates these context-specific linear
coefficients, black-box models implicitly model these coefficients as feature
gradients (i.e. linear coefficients in a first-order Taylor expansion). Akin
to popular post-hoc interpretability methods like LIME (Ribeiro et al., 2016),
we leverage the differentiability of RNNs $\Phi(x_{t},h)\rightarrow a_{t}$ to
recover the implicit context-specific linear policies $\theta$.

$\widehat{\theta}=\frac{\partial}{\partial x_{t}}\Phi(x_{t},h)$

We generate data from a heterogeneous MDP governed by a known context-specific
policy

$\displaystyle P(a_{t}=1|x_{t},x_{t-\tau},a_{t-\tau},t)=1/(1+\exp(-\theta\cdot
x_{t}+\epsilon_{a})),\quad\quad\theta=x_{t-\tau}\cdot(2a_{t-\tau}-1)+\frac{t}{T}+\epsilon_{\theta}$

where $T$ is the sequence length or maximum timestep, and $\tau$ is a time lag
between the current policy and a dependence on past observations $x_{t-\tau}$
and actions $a_{t-\tau}$. We simulate $N$ total sequences, drawing
observations $x_{t}\sim\operatorname{Unif}[-2,2]$, policy noise
$\epsilon_{\theta}\sim N(0,\sigma_{\theta}^{2})$, and decision noise
$\epsilon_{a}\sim N(0,\sigma_{a}^{2})$ at each timestep. On a known
heterogeneous and action-dependent MDP, CPR’s explicit policy representation
not only improves its representation of MDP parameters but increases overall
performance versus a black-box model with an unstructured policy
representation (Fig. 5).

##  5 Discussion

In this study, we propose contextualized policies as dynamic, interpretable,
and personalized linear decision models, each representing a single step in a
complex treatment process. By relating individual modeling tasks through
patient-specific histories and contexts we avoid the pitfalls of common
personalization methods that reduce statistical power (e.g. sample splitting
and subpopulation grouping). As a result, CPR matches the performance of
black-box models while retaining the interpretability of linear models. Post-
hoc analysis of patient-specifc models generated by CPR reveal rare covariates
with outsize effects on treatment decisions, as well as extremely subtle
effects in the general population. While we apply CPR in offline and partially
observable environments, CPR is directly portable to online policy inference
with only subtle training modifications. As a caveat, CPR makes an explicit
trade-off to enable context-specific interpretability that precludes
interpreting direct effects of historical features on current actions, which
may make it unsuitable for tasks where the effect of historical features are
direct and not mediated by a context-specific model. However, for medical
imitation learning, CPR’s policy representation is well-aligned with clinical
practice, where patient history contextualizes current patient symptoms and
treatment options. CPR is a step toward reinforcement learning agents that
explain as they think, and promises a general purpose platform for supporting
and improving complex human decisions across disparate domains.

## Acknowledgments

We thank the Alzheimer’s Disease Neuroimaging Initiative and Medical
Information Mart for Intensive Care for data access. We would like to thank
the INTERACT Program for making this collaboration possible, as well as the
Begabtenstiftung Informatik for providing funding.

## References

  * noa (2018) Dementia: assessment, management and support for people living with dementia and their carers.  _National Institute for Health and Care Excellence (NICE) Guideline No. 97_ , June 2018.  URL https://www.nice.org.uk/guidance/ng97.  URL: https://www.nice.org.uk/guidance/ng97.
  * Ageenko et al. (2010) Ilieva I. Ageenko, Kelley A. Doherty, and Adrian Paul Van Cleave.  Personalized lifetime financial planning tool, June 2010.  URL https://patents.google.com/patent/US20100161467/en.
  * Al-Shedivat et al. (2020) Maruan Al-Shedivat, Avinava Dubey, and Eric P. Xing.  Contextual Explanation Networks, September 2020.  URL http://arxiv.org/abs/1705.10301.  arXiv:1705.10301 [cs, stat].
  * Bain & Sammut (1999) Michael Bain and Claude Sammut.  A Framework for Behavioural Cloning.  In _Machine Intelligence 15, Intelligent Agents [St. Catherine’s College, Oxford, July 1995]_ , pp. 103–129, GBR, January 1999. Oxford University.  ISBN 978-0-19-853867-7.
  * Beck et al. (2011) Andrew H. Beck, Ankur R. Sangoi, Samuel Leung, Robert J. Marinelli, Torsten O. Nielsen, Marc J. van de Vijver, Robert B. West, Matt van de Rijn, and Daphne Koller.  Systematic Analysis of Breast Cancer Morphology Uncovers Stromal Features Associated with Survival.  _Science Translational Medicine_ , 3(108), November 2011.  ISSN 1946-6234, 1946-6242.  doi: 10.1126/scitranslmed.3002564.  URL https://www.science.org/doi/10.1126/scitranslmed.3002564.
  * Bica et al. (2021) Ioana Bica, Daniel Jarrett, Alihan Hüyük, and Mihaela van der Schaar.  Learning "What-if" Explanations for Sequential Decision-Making.  March 2021.  URL https://openreview.net/forum?id=h0de3QWtGG.
  * Buettner et al. (2015) Florian Buettner, Kedar N. Natarajan, F. Paolo Casale, Valentina Proserpio, Antonio Scialdone, Fabian J. Theis, Sarah A. Teichmann, John C. Marioni, and Oliver Stegle.  Computational analysis of cell-to-cell heterogeneity in single-cell RNA-sequencing data reveals hidden subpopulations of cells.  _Nature Biotechnology_ , 33(2):155–160, February 2015.  ISSN 1546-1696.  doi: 10.1038/nbt.3102.  URL https://www.nature.com/articles/nbt.3102.  Number: 2 Publisher: Nature Publishing Group.
  * (8) Laura Castro-Aldrete, Michele V. Moser, Guido Putignano, Maria Teresa Ferretti, Annemarie Schumacher Dimech, and Antonella Santuccione Chadha.  Sex and gender considerations in alzheimer’s disease: The women’s brain project contribution.  15\.  ISSN 1663-4365.  URL https://www.frontiersin.org/articles/10.3389/fnagi.2023.1105620.
  * Choi & Kim (2011) Jaedeug Choi and Kee-Eung Kim.  Inverse Reinforcement Learning in Partially Observable Environments.  _Journal of Machine Learning Research_ , 12(21):691–730, 2011.  ISSN 1533-7928.  URL http://jmlr.org/papers/v12/choi11a.html.
  * Esteva et al. (2017) Andre Esteva, Brett Kuprel, Roberto A. Novoa, Justin Ko, Susan M. Swetter, Helen M. Blau, and Sebastian Thrun.  Dermatologist-level classification of skin cancer with deep neural networks.  _Nature_ , 542(7639):115–118, February 2017.  ISSN 0028-0836, 1476-4687.  doi: 10.1038/nature21056.  URL http://www.nature.com/articles/nature21056.
  * Fisher et al. (2018) Aaron J. Fisher, John D. Medaglia, and Bertus F. Jeronimus.  Lack of group-to-individual generalizability is a threat to human subjects research.  _Proceedings of the National Academy of Sciences_ , 115(27):E6106–E6115, July 2018.  doi: 10.1073/pnas.1711978115.  URL https://www.pnas.org/doi/10.1073/pnas.1711978115.  Publisher: Proceedings of the National Academy of Sciences.
  * Frosst & Hinton (2017) Nicholas Frosst and Geoffrey Hinton.  Distilling a Neural Network Into a Soft Decision Tree, November 2017.  URL http://arxiv.org/abs/1711.09784.  arXiv:1711.09784 [cs, stat].
  * Gounden et al. (2023) Verena Gounden, Harshil Bhatt, and Ishwarlal Jialal.  Renal Function Tests.  In _StatPearls_. StatPearls Publishing, Treasure Island (FL), 2023.  URL http://www.ncbi.nlm.nih.gov/books/NBK507821/.
  * Guleria et al. (2019) Randeep Guleria, Anant Mohan, Atul Kulkarni, Karan Madan, and Vijay Hadda.  Guidelines for Antibiotic Prescription in Intensive Care Unit.  _Indian Journal of Critical Care Medicine_ , 23(S1):1–63, 2019.  ISSN 0972-5229, 1998-359X.  doi: 10.5005/jp-journals-10071-23101.  URL https://www.ijccm.org/doi/10.5005/jp-journals-10071-23101.
  * Hart (2016) S.A. Hart.  Precision Education Initiative: Moving Towards Personalized Education.  _Mind, brain and education : the official journal of the International Mind, Brain, and Education Society_ , 10(4):209–211, December 2016.  ISSN 1751-2271.  doi: 10.1111/mbe.12109.  URL https://www.ncbi.nlm.nih.gov/pmc/articles/PMC5476312/.
  * Hastie & Tibshirani (1993) Trevor Hastie and Robert Tibshirani.  Varying-Coefficient Models.  _Journal of the Royal Statistical Society. Series B (Methodological)_ , 55(4):757–796, 1993.  ISSN 0035-9246.  URL https://www.jstor.org/stable/2345993.  Publisher: [Royal Statistical Society, Wiley].
  * Ho & Ermon (2016) Jonathan Ho and Stefano Ermon.  Generative Adversarial Imitation Learning.  In _Advances in Neural Information Processing Systems_ , volume 29. Curran Associates, Inc., 2016.  URL https://papers.nips.cc/paper_files/paper/2016/hash/cc7e2b878868cbae992d1fb743995d8f-Abstract.html.
  * Hochreiter & Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber.  Long Short-Term Memory.  _Neural Computation_ , 9(8):1735–1780, November 1997.  ISSN 0899-7667.  doi: 10.1162/neco.1997.9.8.1735.  URL https://doi.org/10.1162/neco.1997.9.8.1735.
  * Hüyük et al. (2022) Alihan Hüyük, Daniel Jarrett, Cem Tekin, and Mihaela van der Schaar.  Explaining by Imitating: Understanding Decisions by Interpretable Policy Learning.  February 2022.  URL https://openreview.net/forum?id=unI5ucw_Jk.
  * Jarrett et al. (2021) Daniel Jarrett, Jinsung Yoon, Ioana Bica, Zhaozhi Qian, Ari Ercole, and Mihaela van der Schaar.  Clairvoyance: A pipeline toolkit for medical time series.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=xnC8YwKUE3k.
  * Jeon et al. (2018) Wonseok Jeon, Seokin Seo, and Kee-Eung Kim.  A Bayesian Approach to Generative Adversarial Imitation Learning.  In _Advances in Neural Information Processing Systems_ , volume 31. Curran Associates, Inc., 2018.  URL https://proceedings.neurips.cc/paper/2018/hash/943aa0fcda4ee2901a7de9321663b114-Abstract.html.
  * Johnson et al. (2016) Alistair E. W. Johnson, Tom J. Pollard, Lu Shen, Li-wei H. Lehman, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Leo Anthony Celi, and Roger G. Mark.  MIMIC-III, a freely accessible critical care database.  _Scientific Data_ , 3(1):160035, May 2016.  ISSN 2052-4463.  doi: 10.1038/sdata.2016.35.  URL https://www.nature.com/articles/sdata201635.  Number: 1 Publisher: Nature Publishing Group.
  * Khalili et al. (2013) Hossein Khalili, Samaneh Bairami, and Mona Kargar.  Antibiotics induced acute kidney injury: incidence, risk factors, onset time and outcome.  _Acta Medica Iranica_ , 51(12):871–878, 2013.  ISSN 1735-9694.
  * Kingma & Ba (2014) Diederik P. Kingma and Jimmy Ba.  Adam: A Method for Stochastic Optimization.  _CoRR_ , December 2014.  URL https://www.semanticscholar.org/paper/Adam%3A-A-Method-for-Stochastic-Optimization-Kingma-Ba/a6cb366736791bcccc5c8639de5a8f9636bf87e8.
  * Kostrikov et al. (2020) Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson.  Imitation Learning via Off-Policy Distribution Matching.  In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net, 2020.  URL https://openreview.net/forum?id=Hyg-JC4FDr.
  * Królicka et al. (2020) Anna L. Królicka, Adrianna Kruczkowska, Magdalena Krajewska, and Mariusz A. Kusztal.  Hyponatremia in Infectious Diseases-A Literature Review.  _International Journal of Environmental Research and Public Health_ , 17(15):5320, July 2020.  ISSN 1660-4601.  doi: 10.3390/ijerph17155320.
  * Laï et al. (2020) M.-C. Laï, M. Brian, and M.-F. Mamzer.  Perceptions of artificial intelligence in healthcare: findings from a qualitative survey study among actors in France.  _Journal of Translational Medicine_ , 18(1):14, December 2020.  ISSN 1479-5876.  doi: 10.1186/s12967-019-02204-y.  URL https://translational-medicine.biomedcentral.com/articles/10.1186/s12967-019-02204-y.
  * (28) Benjamin Lengerich, Caleb N. Ellington, Andrea Rubbi, Manolis Kellis, and Eric P. Xing.  Contextualized machine learning.  URL http://arxiv.org/abs/2310.11340.
  * Lengerich et al. (2022) Benjamin J Lengerich, Rich Caruana, Mark E Nunnally, and Manolis Kellis.  Death by round numbers: Glass-box machine learning uncovers biases in medical practice.  _medRxiv_ , pp. 2022–04, 2022.
  * Lundberg & Lee (2017) Scott Lundberg and Su-In Lee.  A Unified Approach to Interpreting Model Predictions.  ISBN: 1705.07874 Publication Title: arXiv [cs.AI], May 2017.  URL http://arxiv.org/abs/1705.07874.
  * Makino & Takeuchi (2012) Takaki Makino and Johane Takeuchi.  Apprenticeship learning for model parameters of partially observable environments.  In _Proceedings of the 29th International Coference on International Conference on Machine Learning_ , ICML’12, pp. 891–898, Madison, WI, USA, June 2012. Omnipress.  ISBN 978-1-4503-1285-1.
  * Masterton et al. (2008) R. G. Masterton, A. Galloway, G. French, M. Street, J. Armstrong, E. Brown, J. Cleverley, P. Dilworth, C. Fry, A. D. Gascoigne, Alan Knox, Dilip Nathwani, Robert Spencer, and Mark Wilcox.  Guidelines for the management of hospital-acquired pneumonia in the UK: report of the working party on hospital-acquired pneumonia of the British Society for Antimicrobial Chemotherapy.  _The Journal of Antimicrobial Chemotherapy_ , 62(1):5–34, July 2008.  ISSN 1460-2091.  doi: 10.1093/jac/dkn162.
  * (33) Leland McInnes, John Healy, Nathaniel Saul, and Lukas Großberger.  UMAP: Uniform manifold approximation and projection.  3(29):861.  ISSN 2475-9066.  doi: 10.21105/joss.00861.  URL https://joss.theoj.org/papers/10.21105/joss.00861.
  * McKinlay et al. (2007) J. B. McKinlay, C. L. Link, K. M. Freund, L. D. Marceau, A. B. O’Donnell, and K. L. Lutfey.  Sources of Variation in Physician Adherence with Clinical Guidelines: Results from a Factorial Experiment.  _Journal of General Internal Medicine_ , 22(3):289–296, March 2007.  ISSN 0884-8734, 1525-1497.  doi: 10.1007/s11606-006-0075-2.  URL http://link.springer.com/10.1007/s11606-006-0075-2.
  * Ng & Russell (2000) Andrew Y. Ng and Stuart J. Russell.  Algorithms for Inverse Reinforcement Learning.  In _Proceedings of the Seventeenth International Conference on Machine Learning_ , ICML ’00, pp. 663–670, San Francisco, CA, USA, June 2000. Morgan Kaufmann Publishers Inc.  ISBN 978-1-55860-707-1.
  * Ng et al. (2015) Kenney Ng, Jimeng Sun, Jianying Hu, and Fei Wang.  Personalized Predictive Modeling and Risk Factor Identification using Patient Similarity.  _AMIA Joint Summits on Translational Science proceedings. AMIA Joint Summits on Translational Science_ , 2015:132–136, 2015.  ISSN 2153-4063.
  * O’Bryant et al. (2008) Sid E. O’Bryant, Stephen C. Waring, C. Munro Cullum, James Hall, Laura Lacritz, Paul J. Massman, Philip J. Lupo, Joan S. Reisch, and Rachelle Doody.  Staging Dementia Using Clinical Dementia Rating Scale Sum of Boxes Scores.  _Archives of neurology_ , 65(8):1091–1095, August 2008.  ISSN 0003-9942.  doi: 10.1001/archneur.65.8.1091.  URL https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3409562/.
  * Pace et al. (2022) Alizée Pace, Alex Chan, and Mihaela van der Schaar.  POETREE: Interpretable Policy Learning with Adaptive Decision Trees.  January 2022.  URL https://openreview.net/forum?id=AJsI-ymaKn_.
  * Piot et al. (2014) Bilal Piot, Matthieu Geist, and Olivier Pietquin.  Boosted and reward-regularized classification for apprenticeship learning.  In _Proceedings of the 2014 international conference on Autonomous agents and multi-agent systems_ , AAMAS ’14, pp. 1249–1256, Richland, SC, 2014. International Foundation for Autonomous Agents and Multiagent Systems.  ISBN 978-1-4503-2738-1.
  * Ribeiro et al. (2016) Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin.  "Why Should I Trust You?": Explaining the Predictions of Any Classifier.  ISBN: 1602.04938 Publication Title: arXiv [cs.LG], February 2016.  URL http://arxiv.org/abs/1602.04938.
  * Royal Society and Royal Society (Great Britain) Staff(2017) (Great Britain) Royal Society (Great Britain) and Royal Society (Great Britain) Staff.  _Machine Learning: The Power and Promise of Computers That Learn by Example_.  Royal Society, 2017.  ISBN 978-1-78252-259-1.
  * Sun et al. (2017) Wen Sun, Arun Venkatraman, Geoffrey J. Gordon, Byron Boots, and J. Andrew Bagnell.  Deeply AggreVaTeD: Differentiable Imitation Learning for Sequential Prediction.  In _Proceedings of the 34th International Conference on Machine Learning_ , pp. 3309–3318. PMLR, July 2017.  URL https://proceedings.mlr.press/v70/sun17d.html.  ISSN: 2640-3498.
  * Verma et al. (2018) Abhinav Verma, Vijayaraghavan Murali, Rishabh Singh, Pushmeet Kohli, and Swarat Chaudhuri.  Programmatically Interpretable Reinforcement Learning.  In Jennifer G. Dy and Andreas Krause (eds.), _Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018_ , volume 80 of _Proceedings of Machine Learning Research_ , pp. 5052–5061. PMLR, 2018.  URL http://proceedings.mlr.press/v80/verma18a.html.
  * Westert et al. (2018) Gert P Westert, Stef Groenewoud, John E Wennberg, Catherine Gerard, Phil DaSilva, Femke Atsma, and David C Goodman.  Medical practice variation: public reporting a first necessary step to spark change.  _International Journal for Quality in Health Care_ , 30(9):731–735, November 2018.  ISSN 1353-4505, 1464-3677.  doi: 10.1093/intqhc/mzy092.  URL https://academic.oup.com/intqhc/article/30/9/731/4990397.
  * Yau et al. (2020) Herman Yau, Chris Russell, and Simon Hadfield.  What Did You Think Would Happen? Explaining Agent Behaviour through Intended Outcomes.  In _Advances in Neural Information Processing Systems_ , volume 33, pp. 18375–18386. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper/2020/hash/d5ab8dc7ef67ca92e41d730982c5c602-Abstract.html.

##  Appendix A Appendix

Code is available on
https://github.com/contextpr/contextualized_policy_replication. All model
embeddings are done with UMAP (McInnes et al., ).

###  A.1 Data

Each dataset is split up into a training set (70% of patients), validation set
(15% of patients) for hyperparameter tuning, and test set (15% of patients) to
report model performance.

####  A.1.1 ADNI

We follow the task described by Pace et al. (2022) and Hüyük et al. (2022),
taking the preprocessing code from Hüyük et al. (2022). Overall, the dataset
contains patients with at least two and a median of three visits. We exclude
patients without CDRSB measurements in one of their visits.

####  A.1.2 MIMIC

We follow the task of predicting antibiotic prescription in the intensive care
unit over up to 6 timesteps as described by Pace et al. (2022) adapting the
preprocessing code provided by Jarrett et al. (2021). Patient trajectories
with missing values in one of the observed features or non-consecutive
measurements, patients above $100$ years or below $1$ year of age, and short
stays below 3 days were eliminated. We end up with 4195 patients in our
dataset. Input features were standardized and the hemoglobin feature used in
previous work was removed since it is highly correlated with the hematocrit
feature ($>0.95$). See Table 3 for the results using all 8 input features.
Since measurements are taken as averages over each day in the ICU, we use the
measurements of the previous day as observations $x_{t}$ to predict $a_{t}$
ensuring that we don’t predict actions from measurements taken after the
patient was treated with antibiotics earlier in the day.

|  | MIMIC Antibiotics
---|---|---
Class | Algorithm | AUROC | AUPRC | Brier $\downarrow$
Interpretable | Logistic regression | $0.60$ | $0.79$ | $0.20$
INTERPOLE Hüyük et al. (2022)† | NR | NR | NR
INTERPOLE Hüyük et al. (2022) ‡ | $0.65$ | NR | $0.21$
POETREE Pace et al. (2022) ‡ | $0.68$ | NR | $0.19$
CPR-RNN (ours) | $0.79$ | $0.88$ | $0.15$
CPR-LSTM (ours) | $0.80$ | $0.88$ | $0.15$
Black-Box | RNN | $0.81$ | $0.89$ | $0.14$
LSTM | $0.82$ | $0.89$ | $0.14$
Table 3: Action-matching performance of different policy learning algorithms
on the MIMIC antibiotics task using all 8 input features. Listed performance
of INTERPOLE and POETREE are from prior reports of state-of-the-art
performance on these canonical datasets: † reported in Hüyük et al. (2022), ‡
reported in Pace et al. (2022). NR: No values reported.

###  A.2 Implementation

CPR implements $g$ as a recurrent neural network (either vanilla RNN or LSTM)
with a hidden state of dimension $k$ and $h$ as a neural network with one
hidden layer, again of size $k$, mapping to the output parameters $\theta$.
The initial context $c_{0}$ at timestep $t=0$ is set to all zeros.

Black-Box RNN and LSTM are implemented similarly, with $k$ being the
dimensionality of the hidden state and one hidden layer of size $k$ directly
mapping to the predicted actions $a_{t}$. The black-box model predicts
$P(a_{t}|x_{t},h_{t})$ by taking $[x_{t},a_{t-1}]$ as input at each timestep
and is optimized using the binary cross entropy loss.

###  A.3 Training

We train all models using the Adam optimizer Kingma & Ba (2014) and early
stopping on the validation set. The initial learning rate chosen for CPR is
$5\text{e-}4$ and $1\text{e-}4$ for the baseline RNNs. We select the
dimensions of the hidden state for both CPR and the baseline RNNs from
$[16,32,64]$. For CPR, $\lambda$ is chosen from $[0.0001,0.001,0.01,0.1]$. The
batch size is selected as $64$ for all models. Table 4 shows the optimal
hyperparameters chosen based on the validation set performance.

|  | ADNI MRI scans | MIMIC Antibiotics
---|---|---|---
Class | Algorithm | $\lambda$ | $k$ | $\lambda$ | $k$
Interpretable | CPR-RNN (ours) | $0.0001$ | $32$ | $0.0001$ | $32$
| CPR-LSTM (ours) | $0.001$ | $64$ | $0.0001$ | $32$
Black-Box | RNN | - | $64$ | - | $32$
| LSTM | - | $64$ | - | $64$
Table 4: Hyperparameters chosen for different models.

###  A.4 Additional Experiments

####  A.4.1 Simulations

Figure 6 shows that CPR is able to recover context-dependent threshold
decision boundaries where an agent takes an action if the observed value is
above or below a certain threshold. Here $x_{t}$ is sampled as
$x_{t}\sim\text{Unif}[0,1]$.

Figure 6: CPR recovers hard decision boundaries, generated by rule-based
decision making over time. Here, when the previous observation $x_{t-1}<0.5$,
the current action is taken if $x_{t}<0.5$. If $x_{t-1}\geq 0.5$, the current
action is reversed. The probabilistic models of these boundaries align closely
with the true step function.

We further simulate dynamic, history-dependent decision policies in a
homogeneous MDP and evaluate CPR based it’s recovery of true simulated true
parameters in these policies. We generate $n=2000$ patient trajectories of
length $T=9$ time steps. At each time step, we sample a random observation
variable $x_{t}\sim\text{Unif}[-1,1]$ and an agent which takes an action
$a_{t}\in[0,1]$. This action is determined by a true observation-to-action
mapping function which depends on the observed history and total treatment
time.

$P(a_{t}|x_{t},x_{t-1})=\sigma(w^{t}(x_{t-1})*x_{t}+b^{t}(t)),$ (1)

where $w^{t}(x_{t-1})=4x_{t-1}$ and $b^{t}(t)=\frac{t-5}{4}$. We simulate a
true contextual policy and assess recovery of simulation parameters with CPR
and a RNN. The probability of taking an action at each time $t$ depends on
both the absolute point in the time series, as well as the history of
observations immediately preceding any time point. We design this simulation
as a succinct way to demonstrate how simple decision models which depend on
both absolute and relative effects over a time series often combine to create
a universal policy that is exceptionally complex and difficult to capture in a
single model that is invariant to time or context. We evaluate CPR based on
it’s ability to imitate the observed time-dependent and context-dependent
actions and recover the simulated true parameters.

(a) Estimated coefficient $w^{t}$

(b) Estimated intercept $b^{t}$

Figure 7:  CPR recovers true policy coefficients in a homogeneous MDP

Our method is able to recover the true model parameters (marked in red) as
shown in Figure 7 with a slight bias in the upper and lower value range.

####  A.4.2 ADNI MRI scans

Bootstrapped Results We run $10$ bootstrap runs to get confidence estimates
for model performance on the ADNI (Table 5) dataset. Each bootstrap sample is
randomly split into a training, validation and test set.

|  | ADNI MRI scans
---|---|---
Class | Algorithm | AUROC | AUPRC | Brier $\downarrow$
Interpretable | Logistic regression | $0.66\pm 0.03$ | $0.86\pm 0.01$ | $0.16\pm 0.01$
INTERPOLE † | $0.60\pm 0.04$ | $0.81\pm 0.08$ | $0.17\pm 0.05$
INTERPOLE ‡ | $0.44\pm 0.04$ | $0.75\pm 0.09$ | $0.19\pm 0.07$
POETREE ‡ | $0.62\pm 0.01$ | $0.82\pm 0.01$ | $0.18\pm 0.01$
CPR-RNN (ours) | $\mathbf{0.72\pm 0.02}$ | $\mathbf{0.88\pm 0.02}$ | $\mathbf{0.15\pm 0.01}$
CPR-LSTM (ours) | $\mathbf{0.72\pm 0.02}$ | $\mathbf{0.88\pm 0.02}$ | $\mathbf{0.15\pm 0.01}$
Black-Box | RNN | $0.72\pm 0.02$ | $0.88\pm 0.02$ | $0.15\pm 0.01$
LSTM | $0.71\pm 0.02$ | $0.88\pm 0.02$ | $0.15\pm 0.01$
Table 5: Action-matching performance of different policy learning algorithms
on the ADNI MRI scans task. Bolded values in each column denote the best
performance of any interpretable model. Listed performance of INTERPOLE and
POETREE are from prior reports of state-of-the-art performance on these
canonical datasets: † reported by Hüyük et al. (2022), ‡ reported by Pace et
al. (2022).  Figure 8: Estimated decision model coefficients over 3 timesteps
for 3 patients with representative ADNI contexts. No confidence intervals are
available, as the discrete diagnostic features in ADNI produce only a discrete
number of possible trajectories.

Figure 8 shows how different contexts influence the agents decision function
over time. Most notably, a high CDRSB value during the first two visits
decreases the probability of ordering an MRI since this can be already seen as
a strong indicator of dementia, making a scan less informative noa (2018).
Patients A and B share the same decision model in $t=1$ since they both show
medium CDRSB and avg hippocampal volume in $t=0$. Afterward, their decision
function differs in $t=2$. Patient A’s hippocampal volume was measured as
"low" in $t=1$ leading to a lower overall probability of ordering an MRI in
$t=2$ indicated by a lower intercept and CDRSB coefficients. In contrast,
Patient B was again diagnosed with an "avg" hippocampal volume in visit $t=1$.
A low hippocampal volume can again be seen as a strong indicator of dementia
making a scan less informative. Patient C, in contrast, is diagnosed with
medium CDRSB and high hippocampal volume throughout all visits. This increases
the probability of ordering an MRI (higher intercept) since there is no clear
indication that would make an MRI obsolete.

Figure 9: Average estimated coefficients for different ADNI contexts, patients
over 70 vs patients under 70. Patient A and B from Figure 8 above. Older
patients are less likely to receive MRIs for both patient groups. Static
context is age and gender.

Static Contexts unrolled trough time Static Contexts such as age and gender do
not only influence the observation-to-action mapping at timestep $t=0$ but
uncover heterogeneity across subgroups throughout time. Figure 9 shows that
the difference in the intercept coefficient between patients below 70 and
above 70 years of age widens over time for both patients A and B indicating
that ordering an MRI is less likely for older patients. We can also see a
slight difference in the CDRSB_med coefficient with it being slightly negative
for old patients, reducing the probability of getting ordering an MRI if an
old patient in this patient group scored a medium result, compared to young
patients.

####  A.4.3 MIMIC Antibiotics

Bootstrapped Results We run $10$ bootstrap runs to get confidence estimates
for model performance on the MIMIC (Table 6) dataset. Each bootstrap sample is
randomly split into a training, validation and test set.

|  | MIMIC Antibiotics
---|---|---
Class | Algorithm | AUROC | AUPRC | Brier $\downarrow$
Interpretable | Logistic regression | $0.57\pm 0.03$ | $0.80\pm 0.03$ | $0.20\pm 0.01$
INTERPOLE † | NR | NR | NR
INTERPOLE ‡ | $0.65$ | NR | $0.21$
POETREE ‡ | $0.68$ | NR | $0.19$
CPR-RNN (ours) | $\mathbf{0.82\pm 0.01}$ | $\mathbf{0.90\pm 0.01}$ | $0.14\pm 0.01$
CPR-LSTM (ours) | $\mathbf{0.82\pm 0.01}$ | $\mathbf{0.90\pm 0.01}$ | $\mathbf{0.14\pm 0.00}$
Black-Box | RNN | $0.83\pm 0.01$ | $0.90\pm 0.01$ | $0.13\pm 0.01$
LSTM | $0.84\pm 0.01$ | $0.91\pm 0.01$ | $0.13\pm 0.00$
Table 6: Action-matching performance of different policy learning algorithms
on the MIMIC antibiotics task. Bolded values in each column denote the best
performance of any interpretable model. Listed performance of INTERPOLE and
POETREE are from prior reports of state-of-the-art performance on these
canonical datasets: † reported by Hüyük et al. (2022), ‡ reported by Pace et
al. (2022). NR: No values reported. Figure 10: Model coefficient embeddings
for MIMIC policy models. Figure 11: Estimated model parameters for three
remaining model trajectories (first two in Fig. 2c). Error bars are
coefficient standard deviations across patients in each trajectory group and
time point. Figure 12: MIMIC patient trajectory clusters, produced by
concatenating decision models over 4 consecutive time points into a trajectory
matrix and embedding them with UMAP.

Model Coefficients Figure 10 shows the heterogeneity in estimated model
coefficients. The main drivers of heterogeneity are the intercept, potassium
and creatinine coefficients.

Treatment Trajectories To uncover typical treatment trajectories, we look at
clusters of coefficients over 4 timesteps. We cluster patients into 5
subgroups using hierarchical clustering based on silhouette score ($0.566$) as
seen in Figure 12. Figure LABEL:fig:coefficients_line shows the two largest
clusters containing patients that get antibiotics throughout their stay
(cluster $1$, $1707$ patients) and patients that never get antibiotics
(cluster $5$, $384$ patients). The remaining three groups are plotted in
Figure 11. Patients that get treatment for the first two days fall into
cluster $3$ ($232$ patients). We can see that their treatment parametrization
changes significantly after treatment is stopped. Cluster $4$ contains
patients that get treated for the first day only ($240$ patients). Both
patients in cluster $3$ and $4$ share similar decision models in $t=1$ since
both were treated in $t=0$, it diverges in $t=2$ after treatment is stopped
for one group and share the decision parametrization in $t=3$ after both
groups were not treated in $t=2$. The remaining patients fall into cluster $4$
($375$ patients).

##  Appendix B Additional Plots

Figure 13: Coefficient values of antibiotic prescription models (MIMIC)
parametrized by context vs parameters of global model (dashed red line)