module medagentgym.args

This script is designed to manage and execute experiments in the MedAgentGym environment, supporting both sequential and parallel (async) execution modes over multiple biomedical-related tasks.

Features:

  • Flexible command-line interface for specifying experiment configuration
  • Supports multiple EHR task environments (e.g., MIMIC-III, BioCoder, MedAgentBench, etc.)
  • Seamless integration with different parallel execution backends (ray, joblib)
  • Experiment management: resume, skip completed, track results
  • Environment credential management via TOML
  • Automatic configuration and result directory setup
  • Logging and error handling for robust execution

This module defines base classes for handling arguments and configurations across the application. It includes classes for model descriptor arguments, general arguments, and configurations that encompasses dataset, model, and training settings.

class Command-line Arguments

Experiment-running arguments through command lines.

This class holds the arguments related to the experiment running of the model.

Attributes:

  • config_path (str): The path to the experiment config file (YAML or JSON). Default is None.
  • task (str): Name of the biomedical task to run (e.g., mimic_iii, biocoder, etc.). Default is None.
  • credential_path (str): Path to TOML file with environment credentials. Default is "./credentials.toml".
  • work_dir (str): Working directory for outputs and logs. Default is "./work_dir/".
  • result_dir_tag (str): Suffix tag for result directory. Default is None.
  • start_idx (int): Start index for experiment runs. Default is 0.
  • end_idx (int): End index (exclusive) for experiment runs. Default is -1 (indicating using the entire set).
  • num_steps (int): Number of steps per episode. Default is 30.
  • async_run (bool): Whether to run experiments in parallel. Default is False. The experiments will be run sequentially by default.
  • n_jobs (int): Number of run experiments in parallel. Default is 5. The argument only works when the --async_run argument is activated.
  • parallel_backend (str): Backend for parallelism (ray or joblib).
  • mode (str): Run mode (train or test).

function parse_arguments()

parse_arguments():

Handles all command-line inputs, allowing both config-based and CLI-based configuration. Command-line arguments override config file values.

function get_task_class(task):

Maps task names to the appropriate environment class dynamically, allowing flexible extension for new tasks.

function run_single_experiment(args, config, idx)

run_single_experiment(args, config, idx):

Runs a single experiment episode:

Args:

  • Instantiates environment and agent.
  • Resumes or skips completed runs.
  • Executes agent step loop with retry on error actions.
  • Records and saves conversation history and result

function sequential_run_experiments

sequential_run_experiments(args, config):

Runs experiments one-by-one for the specified index range.

function async_run_experiments

async_run_experiments(args, config, n_jobs, parallel_backend="ray")

Runs experiments in parallel, supporting either Ray or joblib as backend. Handles job initialization, collection, and logging.

function load_credentials

load_credentials()

Loads environment secrets from a TOML file and injects them as environment variables.