Add support for initial alpha value for SAC

[varunagrawal]

I am trying to implement a brax version of A Walk In The Park and their config file uses init_temperature=0.1 which is the $\alpha$ value (or initial temperature) in SAC.

The current version of SAC in brax hard codes the $\alpha$ value to 1.0 by setting log_alpha = 0.0.
Thus this PR adds the option to specify the $\alpha$ as a paramater.

Additionally, I have added a detailed docstring for the SAC train function.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[btaba]

Thanks for the PR.

Can you comment on the practical usage and range of initial_alpha that has shown some improvement in your applications?

[varunagrawal]

Thanks for the PR.

Can you comment on the practical usage and range of initial_alpha that has shown some improvement in your applications?

Hi @btaba. I apologize for the late response.

I am trying to implement a brax version of A Walk In The Park and their config file uses init_temperature=0.1. Hence my need to add this as a parameter in sac.train.

[varunagrawal]

@btaba I updated the PR comment to describe what I am trying to do better and I merged in the latest HEAD. Can you please review this when you get a chance?