trojai_rl.modelgen package


trojai_rl.modelgen.config module

class trojai_rl.modelgen.config.RLOptimizerConfig(algorithm: str = 'ppo', num_frames: int = 8000000, max_num_frames_rollout: int = 128, num_epochs: int = 1000, device: str = 'cuda', num_frames_per_test: int = 500000, learning_rate: float = 0.001)[source]

Bases: object

Defines configuration parameters for RL training

class trojai_rl.modelgen.config.RunnerConfig(train_env_factory: trojai_rl.datagen.environment_factory.EnvironmentFactory, test_env_factory: trojai_rl.datagen.environment_factory.EnvironmentFactory, trainable_model: <sphinx.ext.autodoc.importer._MockObject object at 0x7f20e44b3490>, optimizer: trojai_rl.modelgen.optimizer_interface.RLOptimizerInterface, parallel: bool = False, model_save_dir: str = '/tmp/models', stats_save_dir: str = '/tmp/model_stats', run_id: Any = None, filename: str = None, save_with_hash: bool = False, save_info: dict = None)[source]

Bases: object

Defines a runner configuration object, required to configure a Runner to train RL models

class trojai_rl.modelgen.config.TestConfig(environment_cfg: Any, count: int = 100, test_description: dict = None, agent_argmax_action: bool = False)[source]

Bases: object


trojai_rl.modelgen.optimizer_interface module

class trojai_rl.modelgen.optimizer_interface.RLOptimizerInterface[source]

Bases: abc.ABC

Object that performs training and testing of TrojAI RL models.

get_cfg_as_dict() → dict[source]

Return a dictionary with key/value pairs that describe the parameters used to train the model.

get_device_type() → str[source]

Return a string representation of the type of device used by the optimizer to train the model.

static load(fname: str)[source]

Load an optimizer from disk and return it :param fname: the filename where the optimizer is serialized :return: The loaded optimizer

save(fname: str) → None[source]

Save the optimizer to a file :param fname - the filename to save the optimizer to

test(model: Any, env_factory: trojai_rl.datagen.environment_factory.EnvironmentFactory) → trojai_rl.modelgen.statistics.TestStatistics[source]

Perform whatever tests desired on the model with clean data and triggered data, return a dictionary of results. :param model: (Any) Trained model :param env_factory: (EnvironmentFactory) :return: (Any, TestStatistics) a TestStatistics object

train(model: Any, env_factory: trojai_rl.datagen.environment_factory.EnvironmentFactory) -> (typing.Any, <class 'trojai_rl.modelgen.statistics.TrainingStatistics'>)[source]

Train the given model using parameters in self.training_params :param model: (Any) The untrained model :param env_factory: (EnvironmentFactory) :return: (Any, TrainingStatistics) trained model and TrainingStatistics object

trojai_rl.modelgen.runner module

class trojai_rl.modelgen.runner.Runner(runner_cfg: trojai_rl.modelgen.config.RunnerConfig)[source]

Bases: object

Defines a Runner object, which takes an environment specification, configuration for training, trains an actual model, and returns it.


Get a trained model and associated train and test statistics, then save.

trojai_rl.modelgen.runner.save_dict_to_json(d, fname)[source]

trojai_rl.modelgen.stable_baselines_optimizer module

trojai_rl.modelgen.statistics module

class trojai_rl.modelgen.statistics.BatchTrainingStatistics(batch_num: int, entropy: float, value: float, policy_loss: float, value_loss: float, grad_norm: float)[source]

Bases: object

Object which contains statistics of one batch of training

class trojai_rl.modelgen.statistics.TestStatistics(aggregated_results: Any, test_info: dict = None)[source]

Bases: object

This object mostly just takes care of saving test information, as the runner expects something like this.


Saves the statistics to disk :param fname: the filename to save to :return: None

class trojai_rl.modelgen.statistics.TrainingStatistics(train_info: dict = None, batch_statistics: Sequence[T_co] = None)[source]

Bases: object

Object which encapsulates all the Training Statistics that were captured during training

add_agent_run_stats(agent_run_statistics: dict)[source]
add_batch_stats(batch_statistics: Union[Sequence[T_co], trojai_rl.modelgen.statistics.BatchTrainingStatistics])[source]

Time should be in seconds.


Saves all batches and agent run stats :param fname: :return:


Saves the last batch statistics to disk :param fname: the filename to save to :return: None

trojai_rl.modelgen.statistics.save_dict_to_json(d, fname)[source]

trojai_rl.modelgen.torch_ac_optimizer module

trojai_rl.modelgen.utils module


Module contents