xanderjc / medkit-learn

The Medkit-Learn(ing) Environment: Medical Decision Modelling through Simulation (NeurIPS 2021) by Alex J. Chan, Ioana Bica, Alihan Huyuk, Daniel Jarrett, and Mihaela van der Schaar.

The Medkit-Learn(ing) Environment

Alex J. Chan, Ioana Bica, Alihan Huyuk, Daniel Jarrett, and Mihaela van der Schaar

The Medkit-Learn(ing) Environment, or Medkit, is a publicly available Python package providing simple and easy access to high-fidelity synthetic medical data.

Primarily, Medkit is a tool that supports: (1) a variety of realistic environment models—learned from actual data, to reflect real medical settings), thus allowing simulation of (2) a variety of expressive and customisable policy models that represent complex human decision-behaviours; as well as (3) ensuring that the environment and policy components are disentangled—hence independently controllable.

By fulfilling the above, Medkit seeks to enable advances in decision modelling to be validated more easily and robustly by enabling users to obtain batch datasets with known ground-truth policy parameterisations that simulate decision making behaviours with various degrees of Markovianity, bounded rationality, confounding, individual consistency and variation in practice.

Medkit is pip installable - to work with the latest version, we recommend cloning it, optionally creating a virtual env, and installing it (this will automatically install dependencies):

git clone https://github.com/XanderJC/medkit-learn.git

cd medkit-learn

pip install -e .

Alternatively, Medkit is available on PyPI, and can be installed simply with:

pip install medkit-learn

Example usage:

import medkit as mk

synthetic_dataset = mk.batch_generate(
    domain = "Ward",
    environment = "CRN",
    policy = "LSTM",
    size = 1000,
    test_size = 200,
    max_length = 10,
    scale = True
)

static_train, observations_train, actions_train = synthetic_dataset['training']
static_test,  observations_test,  actions_test  = synthetic_dataset['testing']

While medical machine learning is by necessity almost always entirely offline, we also provide an interface through which you can interact online with the environment should you find that useful. For example, you could train a custom RL policy on this environment with a specified reward function, then you can test inference algorithms on their ability to represent the policy.

env = mk.live_simulate(
    domain="ICU",
    environment="SVAE"
)

static_obs, observation, info = env.reset()
observation, reward, info, done = env.step(action)

Citing

If you use this software please cite as follows:

@inproceedings{chan2021medkitlearn,
        title={The Medkit-Learn(ing) Environment: Medical Decision Modelling through Simulation},
        author={Alex James Chan and Ioana Bica and Alihan H{\"u}y{\"u}k and Daniel Jarrett and Mihaela van der Schaar},
        booktitle={Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks},
        year={2021},
        url={https://openreview.net/forum?id=Ayf90B1yESX}
}

medkit-learn's Issues

Example code not working

import medkit as mk
synthetic_dataset = mk.batch_generate(
                                   domain = "Ward",
                                   environment = "CRN",
                                   policy = "LSTM",
                                   size = 1000,
                                   test_size = 200,
                                   max_length = 10,
                                   scale = True)

Gives an error:

Traceback (most recent call last):
  File "<stdin>", line 8, in <module>
  File "/home/gtennenholtz/medkit-learn/medkit/api.py", line 58, in batch_generate
    env = env_dict[environment](dom)
  File "/home/gtennenholtz/medkit-learn/medkit/environments/CounterfactualRNN.py", line 118, in __init__
    self.model = CRN_env(domain)
  File "/home/gtennenholtz/medkit-learn/medkit/environments/CounterfactualRNN.py", line 15, in __init__
    self.lstm_layers = self.hyper["lstm_layers"]
KeyError: 'lstm_layers'

Also:
env = mk.live_simulate(domain="ICU", environment="SVAE")

Gives an error:

Traceback (most recent call last):
  File "<stdin>", line 3, in <module>
  File "/home/gtennenholtz/medkit-learn/medkit/api.py", line 214, in live_simulate
    env = env_dict[environment](dom)
  File "/home/gtennenholtz/medkit-learn/medkit/environments/SequentialVAE.py", line 188, in __init__
    self.load_pretrained()
  File "/home/gtennenholtz/medkit-learn/medkit/bases/base_env.py", line 30, in load_pretrained
    self.model.load_state_dict(torch.load(path))
  File "/home/gtennenholtz/venv/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1224, in load_state_dict
    self.__class__.__name__, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for SVAE_env:
        Missing key(s) in state_dict: "lstm.weight_ih_l0", "lstm.bias_ih_l0", "lstm.weight_hh_l0", "lstm.bias_hh_l0", "lstm.weight_ih_l1", "lstm.bias_ih_l1", "lstm.weight_hh_l1", "lstm.bias_hh_l1", "lstm.layers.0.cell.ih.weight", "lstm.layers.0.cell.ih.bias", "lstm.layers.0.cell.hh.weight", "lstm.layers.0.cell.hh.bias", "lstm.layers.1.cell.ih.weight", "lstm.layers.1.cell.ih.bias", "lstm.layers.1.cell.hh.weight", "lstm.layers.1.cell.hh.bias".
        Unexpected key(s) in state_dict: "lstm.ih.weight", "lstm.ih.bias", "lstm.hh.weight", "lstm.hh.bias".
        size mismatch for encoder.linear1.weight: copying a param with shape torch.Size([128, 24]) from checkpoint, the shape in current model is torch.Size([128, 37]).
        size mismatch for decoder.series_cont_mean.weight: copying a param with shape torch.Size([23, 128]) from checkpoint, the shape in current model is torch.Size([37, 128]).
        size mismatch for decoder.series_cont_mean.bias: copying a param with shape torch.Size([23]) from checkpoint, the shape in current model is torch.Size([37]).
        size mismatch for decoder.series_cont_lstd.weight: copying a param with shape torch.Size([23, 128]) from checkpoint, the shape in current model is torch.Size([37, 128]).
        size mismatch for decoder.series_cont_lstd.bias: copying a param with shape torch.Size([23]) from checkpoint, the shape in current model is torch.Size([37]).
        size mismatch for decoder.series_bin.weight: copying a param with shape torch.Size([1, 128]) from checkpoint, the shape in current model is torch.Size([0, 128]).
        size mismatch for decoder.series_bin.bias: copying a param with shape torch.Size([1]) from checkpoint, the shape in current model is torch.Size([0]).