trojai_rl.modelgen package

Subpackages

• trojai_rl.modelgen.architectures package
  • Submodules
  • trojai_rl.modelgen.architectures.atari_architectures module
  • Module contents

Submodules

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)

Bases: object

Defines configuration parameters for RL training

validate()

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)

Bases: object

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

validate()

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

Bases: object

get_argmax_action()

get_count()

get_description()

get_environment_cfg()

validate()

trojai_rl.modelgen.optimizer_interface module

class trojai_rl.modelgen.optimizer_interface.RLOptimizerInterface

Bases: abc.ABC

Object that performs training and testing of TrojAI RL models.

get_cfg_as_dict() → dict

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

get_device_type() → str

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

static load(fname: str)

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

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

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'>)

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)

Bases: object

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

run()

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

validate()

trojai_rl.modelgen.runner.save_dict_to_json(d, fname)

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)

Bases: object

Object which contains statistics of one batch of training

save(fname)

to_dict()

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

Bases: object

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

save(fname)

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

validate()

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

Bases: object

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

add_agent_run_stats(agent_run_statistics: dict)

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

add_train_time(time)

Time should be in seconds.

save_detailed_stats(fname)

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

save_summary(fname)

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)

trojai_rl.modelgen.torch_ac_optimizer module

trojai_rl.modelgen.utils module

trojai_rl.modelgen.utils.is_jsonable(arg)

Module contents