Meta-Gradient RL A2C

[RobvanGastel]

Hi!

I have been working with your package to test your meta-gradient RL example. As I have it implemented now, the algorithm converges on the cartpole environment, however, the meta-parameter gamma only trends downwards. Do I use torchopt incorrectly in the code sample below?

    agent = A2C(
        obs_space=env.observation_space,
        action_space=env.action_space,
        value_coeff=config["a2c"]["value_coeff"],
        writer=writer,
        ac_kwargs=config["actor_critic"],
        max_episode_steps=config["max_episode_steps"],
        device=config["device"],
    )

    # Set the meta-parameter
    gamma = nn.Parameter(
        -torch.log((1 / torch.tensor(config["gamma"])) - 1),
        requires_grad=True,
    )

    # Torchopt optimizers
    inner_optim = torchopt.MetaSGD(agent.ac, lr=config["inner_lr"])
    meta_optim = torchopt.SGD([gamma], lr=config["outer_lr"])

    for _ in range(config["epochs"]):

        for _ in range(config["inner_steps"]):
            data = agent.collect_rollouts(env, torch.sigmoid(gamma))
            loss = agent.optimize(data)

            inner_optim.step(loss)

        # Outer-loop
        data = agent.collect_rollouts(env, torch.sigmoid(gamma))
        meta_loss = agent.optimize(data)

        meta_optim.zero_grad()
        meta_loss.backward()

        # Log the gradient magnitude
        writer.add_scalar("A2C/grad_gamma", gamma.grad, agent.global_step)
        meta_optim.step()

        # Detach the graph
        torchopt.stop_gradient(agent.ac)
        torchopt.stop_gradient(inner_optim)

Any help would be much appreciated! Thank you!