[FEATURE] Add Automatic Mixed Precision option

[mwulfman]

Proposal

I'm proposing adding an option to enable Automatic Mixed Precision (AMP) in order to improve training throughput with minimal code changes.

Results

To validate the implementation, I benchmarked training before and after enabling AMP with:

python mimickit/run.py \
  --mode train \
  --num_envs 2048 \
  --engine_config data/engines/isaac_lab_engine.yaml \
  --env_config data/envs/deepmimic_humanoid_env.yaml \
  --agent_config data/agents/deepmimic_humanoid_ppo_agent.yaml \
  --visualize false \

The only configuration change was enabling in the current branch:

     use_amp: True

Using the additional logging introduced in #107, I obtained the following results on an NVIDIA GeForce RTX 5070 Ti Laptop GPU.

Before Automatic Mixed Precision:

|   Samples_Per_Second |        3.89e+04 |

After Automatic Mixed Precision:

|   Samples_Per_Second |        6.86e+04 |

This corresponds to +70% increase in training throughput.

Changes done:

• Update MPOptimizer to support Automatic Mixed Precision (AMP) with both "fp16" and "bf16" modes
• Update the agent _update_actor and _update_critic methods to pass a loss_factory to optimizer.step, allowing the forward pass and loss computation to run under AMP autocast
• Update the configs while keeping the default behavior unchanged

[xbpeng]

Thanks for submitting this PR, that looks like a really nice speedup! I wonder if it is possible to simplify the control flow. Instead of using a loss factory and calling autocast in the optimizer, we could put autocast around _update_model instead:

MimicKit/mimickit/learning/base_agent.py

Line 238 in aeb7de9

train_info = self._update_model()

That way, all of the operations for model updates will be performed in mixed precision. This could help to simplify the code, and avoid the need for loss factories.

Also, "amp" conflicts with adversarial motion priors. So instead of using amp maybe we should go with mixed_precision.

I haven't used mixed precision before. But it seems like GradScaler is only used if amp_type == torch.float16. Is it not necessary for bf16?