Superhuman performance on sepsis MIMIC-III data by distributional reinforcement learning

Fig 1. Overall data flow. The data flow including the three main challenges: preprocessing, state-action space representation, and reinforcement learning.

Fig 2. Overview of the data. (a) Histogram for the percentage of missing values before matrix completion. Features with more than 90% missing values were not kept. It appears that missingness does not depend on sepsis diagnosis. (b) Histogram for the length (i.e. number of timesteps) of the episodes. (c) 3D representation of cluster centers, using a linear projection along the first three principal components of the data. The size of the points is affine in the population of the clusters. The color scale corresponds to the survival rate in each cluster. (d) A representation of some episodes using the linear projection along principal components. We represent in green two episodes where the patient survived and in red three episodes where the patient died within 90 days.

Fig 3. Comparison of the performance of stochastic policies on the training and test simulators for different numbers of states.

Fig 4. Examples of return distributions when following the clinicians’ policy. From left to right: negative outcome, bimodal distribution, slightly bimodal distribution, positive outcome.

Fig 5. Return distributions of the resulting stochastic policy for all initial states in gray and average return distribution (weighted by the numbers of occurrences as initial states) in blue. The average return distribution of the clinicians’ policy is shown in red for comparison. Left: training simulator; right: test simulator.

Fig 6. Stereographic representation of the data produced by the model. (a) Two-dimensional representation of the computed state values for the clinicians’ policy using PCA. Each dot is a state obtained through the method described in Section 2.3 with the size matching the number of corresponding observations in the dataset. The coloring corresponds to the computed state value. S1 and S2 are specific states from which stereographic representation is produced. (b) (Inverse) stereographic projection from the PCA plane to a sphere with center S1 and radius 1. The coloring corresponds to the state values. (c) The coloring is changed to fit the computed state-action values for some specified action. (d) The sphere is now centered at S2 and the coloring corresponds to a different action. The state-action values were smoothed using inverse distance weighing [33] instead using of piece-wise constant approximation.