module moplayground.learning.training

---

function setup_morlax

setup_morlax(ppo_params, sampling_params, network_params)

---

function setup_amor

setup_amor(ppo_params, sampling_params, network_params)

---

function train_policy

train_policy(
    algorithm,
    ppo_params,
    sampling_params,
    network_params,
    run_dir,
    env,
    eval_env,
    run=None,
    handle_params=None,
    progress_fn=None
)

Train a multi-objective policy on the given environment.

Sets up the GPU, builds the train function and network factory for algorithm, and runs training. Progress plots and checkpoints are written to run_dir. The caller creates run_dir and writes the run’s config.yaml before calling this function (see moplayground.config.save).

Args:

• algorithm: 'morlax' or 'amor'.
• ppo_params: PPO keyword arguments shared by both algorithms.
• sampling_params: Preference-sampling keyword arguments (alpha, k, sampling, warmup_frac).
• network_params: Keyword arguments for the algorithm’s network factory.
• run_dir: Existing run directory for checkpoints and progress plots.
• env: Training environment.
• eval_env: Evaluation environment used for periodic rollouts.
• run: (optional) Experiment-tracking handle (e.g. a wandb run) used to log progress and checkpoints.
• handle_params: (optional) Callable (ppo_params, sampling_params, network_params) -> (train_fn, network_factory). Defaults to the handler registered for algorithm in _ALGO_HANDLERS.
• progress_fn: (optional) Callback invoked each eval step as progress_fn(run, num_steps, metrics, save_dir, training_data) to log/plot training progress. Defaults to mop.utils.plotting.plot_mo_progress.

Returns: Tuple (make_inference_fn, params) — a factory that builds an inference function and the trained policy parameters.

---

function mo_wrapper

mo_wrapper(
    env: mujoco_playground._src.mjx_env.MjxEnv,
    episode_length: int = 1000,
    action_repeat: int = 1,
    randomization_fn=None
) → Wrapper

Multi-Objective Wrapper