memorax.algorithms.PPO

Contents

memorax.algorithms.PPO#

class memorax.algorithms.PPO[source]#

Bases: object

PPO(cfg: memorax.algorithms.ppo.PPOConfig, env: gymnax.environments.environment.Environment, env_params: gymnax.environments.environment.EnvParams, actor_network: flax.linen.module.Module, critic_network: flax.linen.module.Module, actor_optimizer: optax._src.base.GradientTransformation, critic_optimizer: optax._src.base.GradientTransformation)

cfg: PPOConfig#

env: Environment#

env_params: EnvParams#

actor_network: Module#

critic_network: Module#

actor_optimizer: GradientTransformation#

critic_optimizer: GradientTransformation#

init(key)[source]#

Return type:: PPOState
Parameters:: key (Array)

warmup(key, state, num_steps)[source]#

Return type:

Parameters:

key (Array)
state (PPOState)
num_steps (int)

train(key, state, num_steps)[source]#

Return type:

Parameters:

key (Array)
state (PPOState)
num_steps (int)

evaluate(key, state, num_steps)[source]#

Return type:

Parameters:

key (Array)
state (PPOState)
num_steps (int)

__init__(cfg, env, env_params, actor_network, critic_network, actor_optimizer, critic_optimizer)#

Parameters:

cfg (PPOConfig)
env (Environment)
env_params (EnvParams)
actor_network (Module)
critic_network (Module)
actor_optimizer (GradientTransformation)
critic_optimizer (GradientTransformation)

Return type:

None