rlightning.policy.vla_policy.ppo_policy

class rlightning.policy.vla_policy.ppo_policy.VLAPPOPolicy(config, role_type, *args, **kwargs)

Bases: BasePolicy

compute_ppo_loss(batch, gradient_accum_steps=1)

construct_network(env_meta, *args, **kwargs)

Construct the neural network architecture.

Parameters:

• env_meta – Environment metadata for network configuration.

• *args – Variable positional arguments.

• **kwargs – Variable keyword arguments.

get_trainable_parameters()

Return a dict of module state dicts.

is_ready()

load_state_dict(state_dict, *, trainable_only: bool = True)

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)

optimizer_step()

postprocess = MockModule('torch.inference_mode')

rollout_step = MockModule('torch.inference_mode')

save_weights(save_dir: str, epoch: int)

setup_optimizer(optim_cfg)

Set up the optimizer for training.

Parameters:

optim_cfg – Optimizer configuration.

train()

Run a training step for the policy.

update_dataset(data)

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