Index

A

action (kitae.algos.collections.dqn.DQN_tuple attribute) (kitae.algos.collections.ppo.PPO_tuple attribute) (kitae.algos.collections.sac.SAC_tuple attribute) (kitae.algos.collections.td3.TD3_tuple attribute) (kitae.buffer.Experience attribute) action_bias (kitae.modules.policy.PolicyNormalExternalStd attribute) action_clip() (in module kitae.operations.transformation) action_noise (kitae.algos.collections.td3.TD3Params attribute) action_scale (kitae.modules.policy.PolicyNormalExternalStd attribute) action_space (kitae.config.EnvConfig attribute) activation (kitae.modules.modules.MLP attribute)

add() (kitae.buffer.Buffer method) (kitae.interface.IBuffer method) AgentInfo (class in kitae.agent) AgentPyTree (class in kitae.modules.pytree) AgentSerializable (class in kitae.agent) algo_params (kitae.config.AlgoConfig attribute) algo_type (kitae.agent.OffPolicyAgent attribute) (kitae.agent.OnPolicyAgent attribute) AlgoConfig (class in kitae.config) AlgoParams (class in kitae.config) AlgoType (class in kitae.interface) alpha_state (kitae.algos.collections.sac.SACState attribute)

B

BaseAgent (class in kitae.agent) batch_size (kitae.config.UpdateConfig attribute) batchify() (in module kitae.buffer)

batchify_and_randomize() (in module kitae.buffer) Buffer (class in kitae.buffer) buffer (kitae.buffer.Buffer attribute) buffer_factory() (in module kitae.buffer)

C

calculate_gaes_targets() (in module kitae.operations.timesteps) check_env() (in module kitae.loops.train) clip_eps (kitae.algos.collections.ppo.PPOParams attribute) close() (kitae.envs.wrappers.vector.SubProcVecParallelEnv method) (kitae.envs.wrappers.vector.VecParallelEnv method)

close_extras() (kitae.envs.wrappers.vector.VecParallelEnv method) CloudPickleWrapper (class in kitae.envs.wrappers.vector) compute_td_targets() (in module kitae.operations.timesteps) config (kitae.agent.AgentInfo attribute) ConfigSerializable (class in kitae.config) conv_layer() (in module kitae.modules.modules)

D

default_run_name() (in module kitae.saving) dict_to_tuple() (in module kitae.algos.experience) done (kitae.buffer.Experience attribute) DoubleQValueContinuousOutput (class in kitae.modules.qvalue)

DQN (class in kitae.algos.collections.dqn) DQN_tuple (class in kitae.algos.collections.dqn) DQNParams (class in kitae.algos.collections.dqn) DQNState (class in kitae.algos.collections.dqn)

E

encoder_factory() (in module kitae.modules.encoder) encoder_state (kitae.algos.collections.ppo.PPOState attribute) entropy_coef (kitae.algos.collections.ppo.PPOParams attribute) env_cfg (kitae.config.AlgoConfig attribute) EnvConfig (class in kitae.config) EnvpoolCompatibility (class in kitae.envs.wrappers.compatibility) Experience (class in kitae.buffer) ExperiencePipeline (class in kitae.algos.experience) exploration (kitae.algos.collections.dqn.DQNParams attribute) explore() (kitae.agent.BaseAgent method) (kitae.algos.collections.dqn.DQN method) (kitae.algos.collections.sac.SAC method) (kitae.algos.collections.td3.TD3 method) (kitae.interface.IAgent method)

explore_factory() (in module kitae.algos.collections.dqn) (in module kitae.algos.collections.ppo) (in module kitae.algos.collections.sac) (in module kitae.algos.collections.td3) explore_general_factory() (in module kitae.algos.factory) extra (kitae.agent.AgentInfo attribute)

F

final_activation (kitae.modules.modules.MLP attribute)

flatten() (in module kitae.operations.transformation) fn_parallel() (in module kitae.algos.factory)

G

gae (kitae.algos.collections.ppo.PPO_tuple attribute) gamma (kitae.algos.collections.dqn.DQNParams attribute) (kitae.algos.collections.ppo.PPOParams attribute) (kitae.algos.collections.sac.SACParams attribute) (kitae.algos.collections.td3.TD3Params attribute)

get_images() (kitae.envs.wrappers.vector.VecParallelEnv method) get_log_prob() (in module kitae.modules.policy) get_wrapper_attr() (kitae.envs.wrappers.vector.VecParallelEnv method)

I

IActor (class in kitae.interface) IAgent (class in kitae.interface) IBuffer (class in kitae.interface) IndependentVariable (class in kitae.modules.modules)

init_fn (kitae.modules.modules.IndependentVariable attribute) init_params() (in module kitae.modules.modules) initial_alpha (kitae.algos.collections.sac.SACParams attribute) interact_keys() (kitae.agent.BaseAgent method) inverse_linear_interpolation() (in module kitae.operations.transformation)

J

jax_stack_experiences() (in module kitae.buffer)

K

kitae module kitae.actor module kitae.agent module kitae.algos module kitae.algos.collections module kitae.algos.collections.dqn module kitae.algos.collections.ppo module kitae.algos.collections.sac module kitae.algos.collections.td3 module kitae.algos.experience module kitae.algos.factory module kitae.buffer module kitae.config module kitae.envs module kitae.envs.make module kitae.envs.wrappers module kitae.envs.wrappers.compatibility module kitae.envs.wrappers.record_episode_statistics module kitae.envs.wrappers.vector module

kitae.interface module kitae.loops module kitae.loops.train module kitae.loops.update module kitae.modules module kitae.modules.encoder module kitae.modules.modules module kitae.modules.optimizer module kitae.modules.policy module kitae.modules.pytree module kitae.modules.qvalue module kitae.modules.value module kitae.operations module kitae.operations.loss module kitae.operations.timesteps module kitae.operations.transformation module kitae.saving module kitae.types module kitae.version module

L

layers (kitae.modules.modules.MLP attribute) learning_rate (kitae.config.UpdateConfig attribute) learning_rate_annealing (kitae.config.UpdateConfig attribute) length_queue (kitae.envs.wrappers.record_episode_statistics.ParallelRecordEpisodeStatistics attribute) linear_interpolation() (in module kitae.operations.transformation) linear_learning_rate_schedule() (in module kitae.modules.optimizer) log_prob (kitae.algos.collections.ppo.PPO_tuple attribute) (kitae.buffer.Experience attribute)

log_std_max (kitae.algos.collections.sac.SACParams attribute) (kitae.modules.policy.PolicyTanhNormal attribute) log_std_min (kitae.algos.collections.sac.SACParams attribute) (kitae.modules.policy.PolicyTanhNormal attribute) loss_policy_ppo() (in module kitae.operations.loss) loss_shannon_jensen_divergence() (in module kitae.operations.loss) loss_value_clip() (in module kitae.operations.loss)

M

make_double_q_value() (in module kitae.modules.qvalue) make_env() (in module kitae.envs.make) make_policy() (in module kitae.modules.policy) make_single_env() (in module kitae.envs.make) make_vec_env() (in module kitae.envs.make) max_buffer_size (kitae.buffer.Buffer attribute) (kitae.config.UpdateConfig attribute) max_grad_norm (kitae.config.UpdateConfig attribute) merge_n_first_dims() (in module kitae.algos.experience) MLP (class in kitae.modules.modules) module kitae kitae.actor kitae.agent kitae.algos kitae.algos.collections kitae.algos.collections.dqn kitae.algos.collections.ppo kitae.algos.collections.sac kitae.algos.collections.td3 kitae.algos.experience kitae.algos.factory kitae.buffer kitae.config kitae.envs kitae.envs.make kitae.envs.wrappers kitae.envs.wrappers.compatibility kitae.envs.wrappers.record_episode_statistics kitae.envs.wrappers.vector kitae.interface kitae.loops kitae.loops.train kitae.loops.update kitae.modules kitae.modules.encoder kitae.modules.modules kitae.modules.optimizer kitae.modules.policy kitae.modules.pytree kitae.modules.qvalue kitae.modules.value kitae.operations kitae.operations.loss kitae.operations.timesteps kitae.operations.transformation kitae.saving kitae.types kitae.version

N

n_agents (kitae.config.EnvConfig attribute) n_env_steps (kitae.config.TrainConfig attribute) n_envs (kitae.config.EnvConfig attribute) n_epochs (kitae.config.UpdateConfig attribute) name (kitae.modules.encoder.VectorEncoder attribute) (kitae.modules.encoder.VisionEncoder attribute) (kitae.modules.modules.IndependentVariable attribute) (kitae.modules.modules.MLP attribute) (kitae.modules.modules.PassThrough attribute) (kitae.modules.policy.PolicyCategorical attribute) (kitae.modules.policy.PolicyNormal attribute) (kitae.modules.policy.PolicyNormalExternalStd attribute) (kitae.modules.policy.PolicyOutput attribute) (kitae.modules.policy.PolicyTanhNormal attribute) (kitae.modules.qvalue.DoubleQValueContinuousOutput attribute) (kitae.modules.qvalue.QValueContinousOutput attribute) (kitae.modules.qvalue.QValueDiscreteOutput attribute) (kitae.modules.value.ValueOutput attribute)

next_observation (kitae.buffer.Experience attribute) normalize (kitae.algos.collections.ppo.PPOParams attribute) normalize_frames() (in module kitae.operations.transformation) num_outputs (kitae.modules.policy.PolicyCategorical attribute) (kitae.modules.policy.PolicyNormal attribute) (kitae.modules.policy.PolicyNormalExternalStd attribute) (kitae.modules.policy.PolicyOutput attribute) (kitae.modules.policy.PolicyTanhNormal attribute) (kitae.modules.qvalue.QValueDiscreteOutput attribute) numpy_stack_experiences() (in module kitae.buffer)

O

observation (kitae.algos.collections.dqn.DQN_tuple attribute) (kitae.algos.collections.ppo.PPO_tuple attribute) (kitae.algos.collections.sac.SAC_tuple attribute) (kitae.algos.collections.td3.TD3_tuple attribute) (kitae.buffer.Experience attribute) observation_space (kitae.config.EnvConfig attribute)

OFF_POLICY (kitae.interface.AlgoType attribute) OffPolicyAgent (class in kitae.agent) OffPolicyBuffer (class in kitae.buffer) ON_POLICY (kitae.interface.AlgoType attribute) OnPolicyAgent (class in kitae.agent) OnPolicyBuffer (class in kitae.buffer)

P

parallel (kitae.algos.experience.ExperiencePipeline attribute), [1] parallel_copy() (in module kitae.modules.modules) ParallelRecordEpisodeStatistics (class in kitae.envs.wrappers.record_episode_statistics) parent (kitae.modules.encoder.VectorEncoder attribute) (kitae.modules.encoder.VisionEncoder attribute) (kitae.modules.modules.IndependentVariable attribute) (kitae.modules.modules.MLP attribute) (kitae.modules.modules.PassThrough attribute) (kitae.modules.policy.PolicyCategorical attribute) (kitae.modules.policy.PolicyNormal attribute) (kitae.modules.policy.PolicyNormalExternalStd attribute) (kitae.modules.policy.PolicyOutput attribute) (kitae.modules.policy.PolicyTanhNormal attribute) (kitae.modules.qvalue.DoubleQValueContinuousOutput attribute) (kitae.modules.qvalue.QValueContinousOutput attribute) (kitae.modules.qvalue.QValueDiscreteOutput attribute) (kitae.modules.value.ValueOutput attribute) PassThrough (class in kitae.modules.modules) policy_output_factory() (in module kitae.modules.policy) policy_state (kitae.algos.collections.ppo.PPOState attribute) (kitae.algos.collections.sac.SACState attribute) (kitae.algos.collections.td3.TD3State attribute)

policy_update_frequency (kitae.algos.collections.td3.TD3Params attribute) PolicyActor (class in kitae.actor) PolicyCategorical (class in kitae.modules.policy) PolicyNormal (class in kitae.modules.policy) PolicyNormalExternalStd (class in kitae.modules.policy) PolicyOutput (class in kitae.modules.policy) PolicyTanhNormal (class in kitae.modules.policy) PPO (class in kitae.algos.collections.ppo) PPO_tuple (class in kitae.algos.collections.ppo) PPOParams (class in kitae.algos.collections.ppo) PPOState (class in kitae.algos.collections.ppo) preprocess_fn (kitae.modules.encoder.VectorEncoder attribute) (kitae.modules.encoder.VisionEncoder attribute) process_experience_factory() (in module kitae.algos.collections.dqn) (in module kitae.algos.collections.ppo) (in module kitae.algos.collections.sac) (in module kitae.algos.collections.td3) process_termination() (in module kitae.loops.train)

Q

qvalue_factory() (in module kitae.modules.qvalue) qvalue_state (kitae.algos.collections.dqn.DQNState attribute) (kitae.algos.collections.sac.SACState attribute) (kitae.algos.collections.td3.TD3State attribute)

QValueContinousOutput (class in kitae.modules.qvalue) QValueDiscreteOutput (class in kitae.modules.qvalue)

R

rearrange_pattern (kitae.modules.encoder.VisionEncoder attribute) render() (kitae.envs.wrappers.vector.VecParallelEnv method) replace() (kitae.algos.collections.dqn.DQNState method) (kitae.algos.collections.ppo.PPOState method) (kitae.algos.collections.sac.SACState method) (kitae.algos.collections.td3.TD3State method) (kitae.modules.pytree.TrainState method) reset() (kitae.envs.wrappers.compatibility.EnvpoolCompatibility method) (kitae.envs.wrappers.compatibility.SubProcVecParallelEnvCompatibility method) (kitae.envs.wrappers.record_episode_statistics.ParallelRecordEpisodeStatistics method) (kitae.envs.wrappers.vector.SubProcVecParallelEnv method) (kitae.envs.wrappers.vector.VecParallelEnv method)

restore() (kitae.agent.BaseAgent method) resume() (kitae.agent.BaseAgent method) (kitae.interface.IAgent method) return_ (kitae.algos.collections.dqn.DQN_tuple attribute) return_queue (kitae.envs.wrappers.record_episode_statistics.ParallelRecordEpisodeStatistics attribute) reward (kitae.buffer.Experience attribute) rng (kitae.buffer.Buffer attribute) run() (kitae.algos.experience.ExperiencePipeline method) run_parallel() (kitae.algos.experience.ExperiencePipeline method) run_parallel_vectorized() (kitae.algos.experience.ExperiencePipeline method) run_single_pipe() (kitae.algos.experience.ExperiencePipeline method) run_vectorized() (kitae.algos.experience.ExperiencePipeline method)

S

SAC (class in kitae.algos.collections.sac) SAC_tuple (class in kitae.algos.collections.sac) SACParams (class in kitae.algos.collections.sac) SACState (class in kitae.algos.collections.sac) sample() (kitae.buffer.OffPolicyBuffer method) (kitae.buffer.OnPolicyBuffer method) (kitae.interface.IBuffer method) sample_and_log_prob() (in module kitae.modules.policy) save_frequency (kitae.config.TrainConfig attribute) saver (kitae.saving.SaverContext attribute) SaverContext (class in kitae.saving) scope (kitae.modules.encoder.VectorEncoder attribute) (kitae.modules.encoder.VisionEncoder attribute) (kitae.modules.modules.IndependentVariable attribute) (kitae.modules.modules.MLP attribute) (kitae.modules.modules.PassThrough attribute) (kitae.modules.policy.PolicyCategorical attribute) (kitae.modules.policy.PolicyNormal attribute) (kitae.modules.policy.PolicyNormalExternalStd attribute) (kitae.modules.policy.PolicyOutput attribute) (kitae.modules.policy.PolicyTanhNormal attribute) (kitae.modules.qvalue.DoubleQValueContinuousOutput attribute) (kitae.modules.qvalue.QValueContinousOutput attribute) (kitae.modules.qvalue.QValueDiscreteOutput attribute) (kitae.modules.value.ValueOutput attribute) seed (kitae.config.AlgoConfig attribute) select_action() (kitae.actor.PolicyActor method) (kitae.agent.BaseAgent method) (kitae.algos.collections.dqn.DQN method) (kitae.algos.collections.sac.SAC method) (kitae.algos.collections.td3.TD3 method) (kitae.interface.IActor method) (kitae.interface.IAgent method)

serializable_dict (kitae.agent.BaseAgent attribute) serialize() (kitae.agent.AgentSerializable static method) (kitae.config.ConfigSerializable static method) setup() (kitae.modules.policy.PolicyNormal method) (kitae.modules.policy.PolicyNormalExternalStd method) (kitae.modules.policy.PolicyTanhNormal method) shape (kitae.modules.modules.IndependentVariable attribute) shared_encoder (kitae.config.UpdateConfig attribute) should_update() (kitae.agent.OffPolicyAgent method) (kitae.agent.OnPolicyAgent method) (kitae.interface.IAgent method) skip_steps (kitae.algos.collections.dqn.DQNParams attribute) (kitae.algos.collections.sac.SACParams attribute) (kitae.algos.collections.td3.TD3Params attribute) stack() (in module kitae.envs.wrappers.vector) stack_and_merge_n_first_dims() (in module kitae.algos.experience) start_step (kitae.algos.collections.dqn.DQNParams attribute) (kitae.algos.collections.sac.SACParams attribute) (kitae.algos.collections.td3.TD3Params attribute) step (kitae.agent.OffPolicyAgent attribute) step() (kitae.envs.wrappers.record_episode_statistics.ParallelRecordEpisodeStatistics method) (kitae.envs.wrappers.vector.VecParallelEnv method) step_async() (kitae.envs.wrappers.vector.SubProcVecParallelEnv method) (kitae.envs.wrappers.vector.VecParallelEnv method) step_wait() (kitae.envs.wrappers.vector.SubProcVecParallelEnv method) (kitae.envs.wrappers.vector.VecParallelEnv method) SubProcVecParallelEnv (class in kitae.envs.wrappers.vector) SubProcVecParallelEnvCompatibility (class in kitae.envs.wrappers.compatibility)

T

target (kitae.algos.collections.ppo.PPO_tuple attribute) (kitae.algos.collections.sac.SAC_tuple attribute) (kitae.algos.collections.td3.TD3_tuple attribute) target_noise_clip (kitae.algos.collections.td3.TD3Params attribute) target_noise_std (kitae.algos.collections.td3.TD3Params attribute) target_params (kitae.modules.pytree.TrainState attribute) task_name (kitae.config.EnvConfig attribute) tau (kitae.algos.collections.sac.SACParams attribute) (kitae.algos.collections.td3.TD3Params attribute) TD3 (class in kitae.algos.collections.td3) TD3_tuple (class in kitae.algos.collections.td3) TD3Params (class in kitae.algos.collections.td3)

TD3State (class in kitae.algos.collections.td3) tile_images() (in module kitae.envs.wrappers.vector) train() (kitae.agent.BaseAgent method) (kitae.interface.IAgent method) train_cfg (kitae.config.AlgoConfig attribute) train_state_ddpg_factory() (in module kitae.algos.collections.td3) train_state_dqn_factory() (in module kitae.algos.collections.dqn) train_state_ppo_factory() (in module kitae.algos.collections.ppo) train_state_sac_factory() (in module kitae.algos.collections.sac) TrainConfig (class in kitae.config) TrainState (class in kitae.modules.pytree) transforms (kitae.algos.experience.ExperiencePipeline attribute), [1] tuple_to_dict() (in module kitae.algos.experience)

U

unserialize() (kitae.agent.AgentSerializable static method) (kitae.agent.BaseAgent class method) (kitae.config.ConfigSerializable static method) update() (kitae.agent.BaseAgent method) (kitae.algos.collections.td3.TD3 method) (kitae.saving.SaverContext method) update_alpha_factory() (in module kitae.algos.collections.sac) update_cfg (kitae.config.AlgoConfig attribute) update_epoch() (in module kitae.loops.update) update_policy_factory() (in module kitae.algos.collections.sac) (in module kitae.algos.collections.td3)

update_policy_value_factory() (in module kitae.algos.collections.ppo) update_qvalue_factory() (in module kitae.algos.collections.dqn) (in module kitae.algos.collections.sac) (in module kitae.algos.collections.td3) update_step_factory() (in module kitae.algos.collections.dqn) (in module kitae.algos.collections.ppo) (in module kitae.algos.collections.sac) (in module kitae.algos.collections.td3) UpdateConfig (class in kitae.config)

V

value (kitae.algos.collections.ppo.PPO_tuple attribute) value_coef (kitae.algos.collections.ppo.PPOParams attribute) value_state (kitae.algos.collections.ppo.PPOState attribute) ValueOutput (class in kitae.modules.value)

VecParallelEnv (class in kitae.envs.wrappers.vector) VectorEncoder (class in kitae.modules.encoder) vectorized (kitae.algos.experience.ExperiencePipeline attribute), [1] vectorized_train() (in module kitae.loops.train) VisionEncoder (class in kitae.modules.encoder)

W

worker() (in module kitae.envs.wrappers.vector) wrap_multiple_parallel_env() (in module kitae.envs.make)

wrap_single_env() (in module kitae.envs.make) wrap_single_parallel_env() (in module kitae.envs.make)