rlightning.policy.vla_policy

class rlightning.policy.vla_policy.VLAPPOPolicy(config, role_type, *args, **kwargs)[source]

Bases: BasePolicy

compute_ppo_loss(batch, gradient_accum_steps=1)[source]
construct_network(env_meta, *args, **kwargs)[source]

Construct the neural network architecture.

Parameters:

  • env_meta – Environment metadata for network configuration.

  • *args – Variable positional arguments.

  • **kwargs – Variable keyword arguments.

get_trainable_parameters()[source]

Return a dict of module state dicts.

is_ready()[source]
load_state_dict(state_dict, *, trainable_only: bool = True)[source]

Load parameters into the underlying model.

Parameters:

  • state_dict (dict) – A dict with key “model” that maps to the parameters to be loaded.

  • trainable_only (bool, optional) – If True, assumes state_dict only contains the trainable parameters and will load them with strict=False so that non-specified parameters remain unchanged. Defaults to True.

optimize(mini_batch)[source]
optimizer_step()[source]
postprocess = MockModule('torch.inference_mode')
rollout_step = MockModule('torch.inference_mode')
save_weights(save_dir: str, epoch: int)[source]
setup_optimizer(optim_cfg)[source]

Set up the optimizer for training.

Parameters:

optim_cfg – Optimizer configuration.

train()[source]

Run a training step for the policy.

update_dataset(data)[source]

Update the dataset in the policy by getting a batch from the buffer.

Submodules