"""DreamWaQ training config — Manaro-Alpha aligned.

Actor [512,256,128], Critic [512,256,128], CENet [128,64].
lr=1e-3, vae_lr=1e-3, KL_weight=1.0, entropy=0.01.
"""
import motrix_envs.locomotion.go1.dreamwaq  # noqa: F401
from motrix_rl.registry import rlcfg
from motrix_rl.skrl.config import SkrlCfg


def _apply(cfg: SkrlCfg):
    # PPO hyperparams (exact DreamWaQ match)
    cfg.runner.agent.rollouts = 24          # num_steps_per_env
    cfg.runner.agent.learning_epochs = 5    # num_learning_epochs
    cfg.runner.agent.mini_batches = 4       # num_mini_batches
    cfg.runner.agent.learning_rate = 1e-3
    cfg.runner.agent.entropy_loss_scale = 0.01
    cfg.runner.agent.ratio_clip = 0.2
    cfg.runner.agent.discount_factor = 0.99
    cfg.runner.agent.lam = 0.95

    # Network (DreamWaQ: [512,256,128])
    cfg.runner.models.policy.hiddens = [512, 256, 128]
    cfg.runner.models.policy.hidden_activation = ["elu"]
    # initial_log_std: upstream uses 0.0 (std=1) in IsaacGym, but in MotrixSim with
    # action_scale=0.25 that gives ±0.25 rad jitter/step → robot falls immediately and
    # never learns to stand. -1.0 (std=0.37, ±0.09 rad) lets it stay upright early and
    # bootstrap walking. (MotrixSim<->IsaacGym reality-gap adjustment.)
    cfg.runner.models.policy.initial_log_std = -1.0
    cfg.runner.models.value.hiddens = [512, 256, 128]
    cfg.runner.models.value.hidden_activation = ["elu"]

    cfg.runner.memory.memory_size = 24
    cfg.runner.agent.experiment.checkpoint_interval = 200
    cfg.runner.agent.experiment.write_interval = 200


@rlcfg("go1-dreamwaq-walk")
class DreamWaQ(SkrlCfg):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        _apply(self)
