Common Base Classes

class common.base_agent.BaseAgent(agent_id='Default Agent', log_path='./Brains', **kwargs)[source]
check_env(env)[source]

Check environment

Parameters:

env (vector environment) – vector env check for correctness

Raises:

Exception – raise exception if env check fails

Returns:

env check is successful or failed

Return type:

bool

load(path=None) None[source]

Load the model from the specified path

Parameters:

path (str) – model saved path. Defaults to None.

plot_results(steps: int, plot_name='chickai-train') None[source]

Generate reward plot for training

Parameters:
  • steps (int) – number of training steps

  • plot_name (str, optional) – Name of the reward plot. Defaults to “chickai-train”.

save(path: str | None = None) None[source]

Save agent prains to the specified path

Parameters:

path (str) – Path value to save the model

save_encoder_policy_network()[source]

Saves the policy and feature extractor of the agent’s model.

This method saves the policy and feature extractor of the agent’s model to the specified paths. It first checks if the model is loaded, and if not, it prints an error message and returns. Otherwise, it saves the policy as a pickle file and the feature extractor as a PyTorch state dictionary.

Returns:

None

set_feature_extractor_require_grad(model)[source]

Sets the requires_grad attribute of the parameters in the feature extractor of the given model to False.

Parameters:

model (torch.nn.Module) – The model whose feature extractor parameters need to have requires_grad set to False.

Returns:

The updated model with feature extractor parameters having requires_grad set to False.

Return type:

torch.nn.Module

test(env, eps, record_prefix='rest')[source]

Test the agent in the given environment for the set number of steps

Parameters:
  • env – gym environment wrapper

  • eps – number of test episodes

  • record_prefix (str, optional) – recording file name prefix

write_model_properties(model, steps)[source]

Writes the properties of the model to a JSON file.

Parameters:
  • model (object) – The model object.

  • steps (int) – The total number of timesteps.

Returns:

None

class common.base_experiment.Experiment(config)[source]
train_agents()[source]

Function to build training configuration, generate environments and start training