rlightning.policy.rsl_rl_policy

class rlightning.policy.rsl_rl_policy.RSLRLPolicy(config: PolicyConfig, role_type: PolicyRole)[source]

Bases: BasePolicy

construct_network(model_config: Config = None, env_meta: EnvMeta = None, *args, **kwargs)[source]

Create RSLRL algorithm instance, and register rsl_rl modules to the policy.

Parameters:

  • model_config (_type_, optional) – Model config given for initialization. Defaults to None.

  • env_meta (RSLRLVecEnvMeta, optional) – Environment configuration. Defaults to None.

  • *args – Additional positional arguments kept for interface compatibility.

  • **kwargs – Additional keyword arguments kept for interface compatibility.

get_trainable_parameters() → Dict[str, Dict[str, MockModule('torch.Tensor')]][source]

Return a dict of module state dicts.

init_eval(eval_config=None, env_meta: EnvMeta = None)[source]

Initialize evaluation mode, threading/async state, and weight updater.

init_train(train_config: TrainConfig, env_meta=None)[source]

Initialize training mode with rsl_rl algorithm.

load_state_dict(state_dict: Dict[str, Dict[str, MockModule('torch.Tensor')]])[source]

Load module state dicts by module name with warnings on mismatches.

postprocess(env_ret: EnvRet, policy_resp: PolicyResponse) → PolicyResponse[source]

Convert env transition to rsl_rl format and attach value targets.

rollout_step(env_ret: EnvRet) → PolicyResponse[source]

Sample action from policy or random policy before initialization.

setup_optimizer(optim_cfg)[source]

reconstruct optimizer for rsl_rl modules.

train()[source]

Pop a dataset from the buffer and run a single rsl_rl update.

update_dataset(data: MockModule('tensordict.TensorDict'))[source]

Assign the given dataset to policy for training.

class rlightning.policy.rsl_rl_policy.RSLRLVecEnvMeta(env_id: str = None, action_space: MockModule('gymnasium.spaces.Space') | None = None, observation_space: MockModule('gymnasium.spaces.Space') | None = None, num_envs: int | None = None, num_actions: int = None, get_observations: Callable = None)[source]

Bases: EnvMeta

get_observations: Callable = None
num_actions: int = None

Submodules