Index

_

__call__() (nnx_ppo.algorithms.checkpointing.CheckpointCallback method) (nnx_ppo.networks.adapter.PPOAdapter method) (nnx_ppo.networks.containers.Concat method) (nnx_ppo.networks.containers.Parallel method) (nnx_ppo.networks.containers.Sequential method) (nnx_ppo.networks.containers.Splitter method) (nnx_ppo.networks.delay.Delay method) (nnx_ppo.networks.feedforward.Dense method) (nnx_ppo.networks.graph.graph.PopulationGraph method) (nnx_ppo.networks.graph.PopulationGraph method) (nnx_ppo.networks.normalizer.Normalizer method) (nnx_ppo.networks.recording.Recorder method) (nnx_ppo.networks.recurrent.LSTM method) (nnx_ppo.networks.sampling_layers.ActionSampler method) (nnx_ppo.networks.sampling_layers.NormalTanhSampler method) (nnx_ppo.networks.types.StatefulModule method) (nnx_ppo.networks.utils.Filter method) (nnx_ppo.networks.utils.Flattener method) (nnx_ppo.networks.utils.Map method) (nnx_ppo.networks.utils.Merge method) (nnx_ppo.networks.utils.Scale method) (nnx_ppo.networks.variational.AR1VariationalBottleneck method) (nnx_ppo.networks.variational.VariationalBottleneck method) __init__() (nnx_ppo.algorithms.config.DistillationConfig method) (nnx_ppo.algorithms.config.DistillationTrainConfig method) (nnx_ppo.algorithms.config.DistillationTrainResult method) (nnx_ppo.algorithms.config.EvalConfig method) (nnx_ppo.algorithms.config.PPOConfig method) (nnx_ppo.algorithms.config.TrainConfig method) (nnx_ppo.algorithms.config.TrainResult method) (nnx_ppo.algorithms.config.VideoConfig method) (nnx_ppo.algorithms.config.VideoData method) (nnx_ppo.algorithms.types.DistillationState method) (nnx_ppo.algorithms.types.DistillationTransition method) (nnx_ppo.algorithms.types.EnvState method) (nnx_ppo.algorithms.types.RLEnv method) (nnx_ppo.algorithms.types.TrainingState method) (nnx_ppo.algorithms.types.Transition method) (nnx_ppo.networks.adapter.PPOAdapter method) (nnx_ppo.networks.containers.Concat method) (nnx_ppo.networks.containers.Parallel method) (nnx_ppo.networks.containers.Sequential method) (nnx_ppo.networks.containers.Splitter method) (nnx_ppo.networks.delay.Delay method) (nnx_ppo.networks.feedforward.Dense method) (nnx_ppo.networks.graph.Connection method) (nnx_ppo.networks.graph.graph.PopulationGraph method) (nnx_ppo.networks.graph.Population method) (nnx_ppo.networks.graph.PopulationGraph method) (nnx_ppo.networks.normalizer.Normalizer method) (nnx_ppo.networks.recording.Recorder method) (nnx_ppo.networks.recurrent.LSTM method) (nnx_ppo.networks.sampling_layers.NormalTanhSampler method) (nnx_ppo.networks.types.PPONetworkOutput method) (nnx_ppo.networks.types.StatefulModuleOutput method) (nnx_ppo.networks.utils.Filter method) (nnx_ppo.networks.utils.Flattener method) (nnx_ppo.networks.utils.Map method) (nnx_ppo.networks.utils.Merge method) (nnx_ppo.networks.utils.Scale method) (nnx_ppo.networks.variational.AR1VariationalBottleneck method) (nnx_ppo.networks.variational.VariationalBottleneck method) (nnx_ppo.wrappers.episode_wrapper.EpisodeWrapper method) (nnx_ppo.wrappers.reward_scaling_wrapper.RewardScalingWrapper method)

A

action_size (nnx_ppo.wrappers.episode_wrapper.EpisodeWrapper property) (nnx_ppo.wrappers.reward_scaling_wrapper.RewardScalingWrapper property) actions (nnx_ppo.networks.types.PPONetworkOutput attribute) ActionSampler (class in nnx_ppo.networks.sampling_layers) ACTOR_EXTRA (nnx_ppo.algorithms.types.LoggingLevel attribute) add_input() (nnx_ppo.networks.graph.graph.PopulationGraph method) (nnx_ppo.networks.graph.PopulationGraph method)

add_output() (nnx_ppo.networks.graph.graph.PopulationGraph method) (nnx_ppo.networks.graph.PopulationGraph method) add_population() (nnx_ppo.networks.graph.graph.PopulationGraph method) (nnx_ppo.networks.graph.PopulationGraph method) ALL (nnx_ppo.algorithms.types.LoggingLevel attribute) AR1VariationalBottleneck (class in nnx_ppo.networks.variational)

B

BASIC (nnx_ppo.algorithms.types.LoggingLevel attribute)

C

checkpoint_every_steps (nnx_ppo.algorithms.config.DistillationTrainConfig attribute) (nnx_ppo.algorithms.config.TrainConfig attribute) CheckpointCallback (class in nnx_ppo.algorithms.checkpointing) clip_range (nnx_ppo.algorithms.config.PPOConfig attribute) combine_advantages (nnx_ppo.algorithms.config.PPOConfig attribute) compute_metrics() (in module nnx_ppo.algorithms.metrics)

Concat (class in nnx_ppo.networks.containers) connect() (nnx_ppo.networks.graph.graph.PopulationGraph method) (nnx_ppo.networks.graph.PopulationGraph method) Connection (class in nnx_ppo.networks.graph) CRITIC_EXTRA (nnx_ppo.algorithms.types.LoggingLevel attribute) critic_loss_weight (nnx_ppo.algorithms.config.PPOConfig attribute)

D

data (nnx_ppo.algorithms.rollout.SlimState attribute) default_config() (in module nnx_ppo.algorithms.ppo) Delay (class in nnx_ppo.networks.delay) Dense (class in nnx_ppo.networks.feedforward) deterministic (nnx_ppo.networks.sampling_layers.ActionSampler attribute) discounting_factor (nnx_ppo.algorithms.config.PPOConfig attribute) distillation (nnx_ppo.algorithms.config.DistillationTrainConfig attribute) DistillationConfig (class in nnx_ppo.algorithms.config)

DistillationState (class in nnx_ppo.algorithms.types) DistillationTrainConfig (class in nnx_ppo.algorithms.config) DistillationTrainResult (class in nnx_ppo.algorithms.config) DistillationTransition (class in nnx_ppo.algorithms.types) done (nnx_ppo.algorithms.rollout.SlimState attribute) (nnx_ppo.algorithms.types.DistillationTransition attribute) (nnx_ppo.algorithms.types.EnvState property) (nnx_ppo.algorithms.types.Transition attribute)

E

enabled (nnx_ppo.algorithms.config.EvalConfig attribute) (nnx_ppo.algorithms.config.VideoConfig attribute) ENV_METRICS (nnx_ppo.algorithms.types.LoggingLevel attribute) env_states (nnx_ppo.algorithms.types.DistillationState attribute) (nnx_ppo.algorithms.types.TrainingState attribute) EnvState (class in nnx_ppo.algorithms.types) episode_length (nnx_ppo.algorithms.config.VideoConfig attribute) (nnx_ppo.algorithms.config.VideoData attribute) episode_reward (nnx_ppo.algorithms.config.VideoData attribute) EpisodeWrapper (class in nnx_ppo.wrappers.episode_wrapper)

eval (nnx_ppo.algorithms.config.DistillationTrainConfig attribute) (nnx_ppo.algorithms.config.TrainConfig attribute) eval_history (nnx_ppo.algorithms.config.DistillationTrainResult attribute) (nnx_ppo.algorithms.config.TrainResult attribute) eval_rollout() (in module nnx_ppo.algorithms.rollout) eval_rollout_for_render_scan() (in module nnx_ppo.algorithms.rollout) EvalConfig (class in nnx_ppo.algorithms.config) every_steps (nnx_ppo.algorithms.config.EvalConfig attribute) (nnx_ppo.algorithms.config.VideoConfig attribute) extract_activations() (in module nnx_ppo.networks.recording)

F

Filter (class in nnx_ppo.networks.utils) final_metrics (nnx_ppo.algorithms.config.DistillationTrainResult attribute) (nnx_ppo.algorithms.config.TrainResult attribute)

finalize() (nnx_ppo.networks.graph.graph.PopulationGraph method) (nnx_ppo.networks.graph.PopulationGraph method) Flattener (class in nnx_ppo.networks.utils) frames (nnx_ppo.algorithms.config.VideoData attribute)

G

gae() (in module nnx_ppo.algorithms.ppo) gae_lambda (nnx_ppo.algorithms.config.PPOConfig attribute)

GRAD_NORM (nnx_ppo.algorithms.types.LoggingLevel attribute) gradient_clipping (nnx_ppo.algorithms.config.DistillationConfig attribute) (nnx_ppo.algorithms.config.PPOConfig attribute)

I

info (nnx_ppo.algorithms.rollout.SlimState attribute) (nnx_ppo.algorithms.types.EnvState property) initialize_state() (nnx_ppo.networks.adapter.PPOAdapter method) (nnx_ppo.networks.containers.Concat method) (nnx_ppo.networks.containers.Parallel method) (nnx_ppo.networks.containers.Sequential method) (nnx_ppo.networks.delay.Delay method) (nnx_ppo.networks.graph.graph.PopulationGraph method) (nnx_ppo.networks.graph.PopulationGraph method) (nnx_ppo.networks.recording.Recorder method) (nnx_ppo.networks.recurrent.LSTM method) (nnx_ppo.networks.sampling_layers.NormalTanhSampler method) (nnx_ppo.networks.types.StatefulModule method) (nnx_ppo.networks.utils.Map method) (nnx_ppo.networks.utils.Merge method) (nnx_ppo.networks.variational.AR1VariationalBottleneck method) (nnx_ppo.networks.variational.VariationalBottleneck method)

J

JaxDataclass (class in nnx_ppo.jax_dataclass)

L

learning_rate (nnx_ppo.algorithms.config.DistillationConfig attribute) (nnx_ppo.algorithms.config.PPOConfig attribute) load_checkpoint() (in module nnx_ppo.algorithms.checkpointing) log_weight_stats() (in module nnx_ppo.algorithms.metrics) logging_level (nnx_ppo.algorithms.config.DistillationConfig attribute) (nnx_ppo.algorithms.config.EvalConfig attribute) (nnx_ppo.algorithms.config.PPOConfig attribute)

logging_percentiles (nnx_ppo.algorithms.config.DistillationConfig attribute) (nnx_ppo.algorithms.config.EvalConfig attribute) (nnx_ppo.algorithms.config.PPOConfig attribute) LoggingLevel (class in nnx_ppo.algorithms.types) loglikelihoods (nnx_ppo.networks.types.PPONetworkOutput attribute) LOSSES (nnx_ppo.algorithms.types.LoggingLevel attribute) LSTM (class in nnx_ppo.networks.recurrent)

M

make_checkpoint_fn() (in module nnx_ppo.algorithms.checkpointing) make_mlp() (in module nnx_ppo.networks.factories) make_mlp_actor_critic() (in module nnx_ppo.networks.factories) make_mlp_layers() (in module nnx_ppo.networks.factories) Map (class in nnx_ppo.networks.utils) max_episode_length (nnx_ppo.algorithms.config.EvalConfig attribute) Merge (class in nnx_ppo.networks.utils) metrics (nnx_ppo.algorithms.rollout.SlimState attribute) (nnx_ppo.algorithms.types.DistillationTransition attribute) (nnx_ppo.algorithms.types.EnvState property) (nnx_ppo.algorithms.types.Transition attribute) (nnx_ppo.networks.types.StatefulModuleOutput attribute) mocap_pos (nnx_ppo.algorithms.rollout.SlimData attribute) mocap_quat (nnx_ppo.algorithms.rollout.SlimData attribute) module nnx_ppo.algorithms.callbacks nnx_ppo.algorithms.checkpointing nnx_ppo.algorithms.config nnx_ppo.algorithms.metrics nnx_ppo.algorithms.ppo nnx_ppo.algorithms.rollout nnx_ppo.algorithms.types nnx_ppo.jax_dataclass nnx_ppo.networks.adapter nnx_ppo.networks.containers nnx_ppo.networks.delay nnx_ppo.networks.factories nnx_ppo.networks.feedforward nnx_ppo.networks.graph nnx_ppo.networks.graph.graph nnx_ppo.networks.normalizer nnx_ppo.networks.recording nnx_ppo.networks.recurrent nnx_ppo.networks.sampling_layers nnx_ppo.networks.types nnx_ppo.networks.utils nnx_ppo.networks.variational nnx_ppo.wrappers.episode_wrapper nnx_ppo.wrappers.reward_scaling_wrapper

N

n_envs (nnx_ppo.algorithms.config.DistillationConfig attribute) (nnx_ppo.algorithms.config.EvalConfig attribute) (nnx_ppo.algorithms.config.PPOConfig attribute) n_epochs (nnx_ppo.algorithms.config.DistillationConfig attribute) (nnx_ppo.algorithms.config.PPOConfig attribute) n_minibatches (nnx_ppo.algorithms.config.DistillationConfig attribute) (nnx_ppo.algorithms.config.PPOConfig attribute) NETWORK_METRICS (nnx_ppo.algorithms.types.LoggingLevel attribute) network_output (nnx_ppo.algorithms.types.Transition attribute) network_states (nnx_ppo.algorithms.types.TrainingState attribute) networks (nnx_ppo.algorithms.types.TrainingState attribute) new_training_state() (in module nnx_ppo.algorithms.ppo) next_obs (nnx_ppo.algorithms.types.DistillationTransition attribute) (nnx_ppo.algorithms.types.Transition attribute) next_state (nnx_ppo.networks.types.StatefulModuleOutput attribute) nnx_ppo.algorithms.callbacks module nnx_ppo.algorithms.checkpointing module nnx_ppo.algorithms.config module nnx_ppo.algorithms.metrics module nnx_ppo.algorithms.ppo module nnx_ppo.algorithms.rollout module nnx_ppo.algorithms.types module nnx_ppo.jax_dataclass module nnx_ppo.networks.adapter module nnx_ppo.networks.containers module

nnx_ppo.networks.delay module nnx_ppo.networks.factories module nnx_ppo.networks.feedforward module nnx_ppo.networks.graph module nnx_ppo.networks.graph.graph module nnx_ppo.networks.normalizer module nnx_ppo.networks.recording module nnx_ppo.networks.recurrent module nnx_ppo.networks.sampling_layers module nnx_ppo.networks.types module nnx_ppo.networks.utils module nnx_ppo.networks.variational module nnx_ppo.wrappers.episode_wrapper module nnx_ppo.wrappers.reward_scaling_wrapper module NONE (nnx_ppo.algorithms.types.LoggingLevel attribute) normalize_advantages (nnx_ppo.algorithms.config.PPOConfig attribute) Normalizer (class in nnx_ppo.networks.normalizer) NormalizerStatistics (class in nnx_ppo.networks.normalizer) NormalTanhSampler (class in nnx_ppo.networks.sampling_layers)

O

obs (nnx_ppo.algorithms.types.DistillationTransition attribute) (nnx_ppo.algorithms.types.EnvState property) (nnx_ppo.algorithms.types.Transition attribute) observation_size (nnx_ppo.wrappers.episode_wrapper.EpisodeWrapper property) (nnx_ppo.wrappers.reward_scaling_wrapper.RewardScalingWrapper property)

optimizer (nnx_ppo.algorithms.types.DistillationState attribute) (nnx_ppo.algorithms.types.TrainingState attribute) output (nnx_ppo.networks.types.StatefulModuleOutput attribute)

P

Parallel (class in nnx_ppo.networks.containers) Population (class in nnx_ppo.networks.graph) PopulationGraph (class in nnx_ppo.networks.graph) (class in nnx_ppo.networks.graph.graph) ppo (nnx_ppo.algorithms.config.TrainConfig attribute)

ppo_loss() (in module nnx_ppo.algorithms.ppo) ppo_step() (in module nnx_ppo.algorithms.ppo) PPOAdapter (class in nnx_ppo.networks.adapter) PPOConfig (class in nnx_ppo.algorithms.config) PPONetworkOutput (class in nnx_ppo.networks.types)

Q

qpos (nnx_ppo.algorithms.rollout.SlimData attribute)

qvel (nnx_ppo.algorithms.rollout.SlimData attribute)

R

record_activations_rollout() (in module nnx_ppo.algorithms.rollout) Recorder (class in nnx_ppo.networks.recording) regularization_loss (nnx_ppo.networks.types.StatefulModuleOutput attribute) render_kwargs (nnx_ppo.algorithms.config.VideoConfig attribute) replace() (nnx_ppo.jax_dataclass.JaxDataclass method) reset() (nnx_ppo.algorithms.types.RLEnv method) (nnx_ppo.wrappers.episode_wrapper.EpisodeWrapper method) (nnx_ppo.wrappers.reward_scaling_wrapper.RewardScalingWrapper method) reset_state() (nnx_ppo.networks.adapter.PPOAdapter method) (nnx_ppo.networks.containers.Concat method) (nnx_ppo.networks.containers.Parallel method) (nnx_ppo.networks.containers.Sequential method) (nnx_ppo.networks.delay.Delay method) (nnx_ppo.networks.graph.graph.PopulationGraph method) (nnx_ppo.networks.graph.PopulationGraph method) (nnx_ppo.networks.recording.Recorder method) (nnx_ppo.networks.recurrent.LSTM method) (nnx_ppo.networks.types.StatefulModule method) (nnx_ppo.networks.utils.Map method) (nnx_ppo.networks.utils.Merge method) (nnx_ppo.networks.variational.AR1VariationalBottleneck method) (nnx_ppo.networks.variational.VariationalBottleneck method)

reward (nnx_ppo.algorithms.types.EnvState property) rewards (nnx_ppo.algorithms.types.DistillationTransition attribute) (nnx_ppo.algorithms.types.Transition attribute) RewardScalingWrapper (class in nnx_ppo.wrappers.reward_scaling_wrapper) RLEnv (class in nnx_ppo.algorithms.types) rng_key (nnx_ppo.algorithms.types.DistillationState attribute) (nnx_ppo.algorithms.types.TrainingState attribute) rollout_extras (nnx_ppo.algorithms.types.Transition attribute) (nnx_ppo.networks.types.StatefulModuleOutput attribute) rollout_length (nnx_ppo.algorithms.config.DistillationConfig attribute) (nnx_ppo.algorithms.config.PPOConfig attribute) ROLLOUT_OBS (nnx_ppo.algorithms.types.LoggingLevel attribute) ROLLOUT_STATS (nnx_ppo.algorithms.types.LoggingLevel attribute)

S

Scale (class in nnx_ppo.networks.utils) seed (nnx_ppo.algorithms.config.DistillationTrainConfig attribute) (nnx_ppo.algorithms.config.TrainConfig attribute) Sequential (class in nnx_ppo.networks.containers) single_transition() (in module nnx_ppo.algorithms.rollout) SlimData (class in nnx_ppo.algorithms.rollout) SlimState (class in nnx_ppo.algorithms.rollout) Splitter (class in nnx_ppo.networks.containers) StatefulModule (class in nnx_ppo.networks.types) StatefulModuleOutput (class in nnx_ppo.networks.types)

step (nnx_ppo.algorithms.config.VideoData attribute) step() (nnx_ppo.algorithms.types.RLEnv method) (nnx_ppo.wrappers.episode_wrapper.EpisodeWrapper method) (nnx_ppo.wrappers.reward_scaling_wrapper.RewardScalingWrapper method) steps_taken (nnx_ppo.algorithms.types.DistillationState attribute) (nnx_ppo.algorithms.types.TrainingState attribute) student (nnx_ppo.algorithms.types.DistillationState attribute) student_output (nnx_ppo.algorithms.types.DistillationTransition attribute) student_rollout_extras (nnx_ppo.algorithms.types.DistillationTransition attribute) student_states (nnx_ppo.algorithms.types.DistillationState attribute)

T

teacher_rollout_extras (nnx_ppo.algorithms.types.DistillationTransition attribute) teacher_states (nnx_ppo.algorithms.types.DistillationState attribute) THROUGHPUT (nnx_ppo.algorithms.types.LoggingLevel attribute) time (nnx_ppo.algorithms.rollout.SlimData attribute) total_iterations (nnx_ppo.algorithms.config.DistillationTrainResult attribute) (nnx_ppo.algorithms.config.TrainResult attribute) total_steps (nnx_ppo.algorithms.config.DistillationConfig attribute) (nnx_ppo.algorithms.config.DistillationTrainResult attribute) (nnx_ppo.algorithms.config.PPOConfig attribute) (nnx_ppo.algorithms.config.TrainResult attribute) train_ppo() (in module nnx_ppo.algorithms.ppo)

TrainConfig (class in nnx_ppo.algorithms.config) training_state (nnx_ppo.algorithms.config.DistillationTrainResult attribute) (nnx_ppo.algorithms.config.TrainResult attribute) TrainingState (class in nnx_ppo.algorithms.types) TrainResult (class in nnx_ppo.algorithms.config) Transition (class in nnx_ppo.algorithms.types) tree_flatten() (nnx_ppo.jax_dataclass.JaxDataclass method) tree_unflatten() (nnx_ppo.jax_dataclass.JaxDataclass class method) tree_where() (in module nnx_ppo.algorithms.rollout) truncated (nnx_ppo.algorithms.types.DistillationTransition attribute) (nnx_ppo.algorithms.types.Transition attribute)

U

unroll_env() (in module nnx_ppo.algorithms.rollout) unstack_trajectory() (in module nnx_ppo.algorithms.rollout) update_statistics() (nnx_ppo.networks.adapter.PPOAdapter method) (nnx_ppo.networks.containers.Concat method) (nnx_ppo.networks.containers.Parallel method) (nnx_ppo.networks.containers.Sequential method) (nnx_ppo.networks.graph.graph.PopulationGraph method) (nnx_ppo.networks.graph.PopulationGraph method) (nnx_ppo.networks.normalizer.Normalizer method) (nnx_ppo.networks.recording.Recorder method) (nnx_ppo.networks.types.StatefulModule method) (nnx_ppo.networks.utils.Map method) (nnx_ppo.networks.utils.Merge method)

V

value_estimates (nnx_ppo.networks.types.PPONetworkOutput attribute) VariationalBottleneck (class in nnx_ppo.networks.variational) video (nnx_ppo.algorithms.config.DistillationTrainConfig attribute) (nnx_ppo.algorithms.config.TrainConfig attribute)

VideoConfig (class in nnx_ppo.algorithms.config) VideoData (class in nnx_ppo.algorithms.config)

W

wandb_video_fn() (in module nnx_ppo.algorithms.callbacks) weight_decay (nnx_ppo.algorithms.config.DistillationConfig attribute) (nnx_ppo.algorithms.config.PPOConfig attribute)

WEIGHTS (nnx_ppo.algorithms.types.LoggingLevel attribute) with_recording() (in module nnx_ppo.networks.recording)

X

xfrc_applied (nnx_ppo.algorithms.rollout.SlimData attribute)