-
Notifications
You must be signed in to change notification settings - Fork 0
Add a reinforcement learning based imaginary time evolution algorithm. #65
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull Request Overview
This PR adds a new reinforcement learning based imaginary time evolution (RLIM) algorithm to the qmb quantum many-body simulation library. The implementation provides an alternative optimization approach that combines reinforcement learning concepts with imaginary time evolution for quantum state optimization.
- Implements a new RLIM algorithm class with configurable parameters for sampling, learning rates, and optimization steps
- Integrates the new algorithm into the existing CLI framework through subcommand registration
- Provides comprehensive logging and TensorBoard monitoring for the optimization process
Reviewed Changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| qmb/rlim.py | Complete implementation of the RLIM algorithm with configuration class and main optimization loop |
| qmb/main.py | Registration of the new RLIM subcommand in the CLI interface |
| # pylint: disable=too-many-locals | ||
|
|
||
| model, network, data = self.common.main() | ||
| ref_network = network |
Copilot
AI
Jul 17, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Assigning ref_network = network creates a reference to the same object rather than an independent copy. This means both networks will be updated simultaneously during optimization, which may not be the intended behavior for a reference network that should remain stable.
| ref_network = network | |
| ref_network = type(network)() # Create a new instance of the same model class | |
| ref_network.load_state_dict(network.state_dict()) # Copy the parameters from the original network |
qmb/rlim.py
Outdated
| a = torch.outer(psi_src.detach().conj(), ref_psi_src) - torch.outer(psi_src.conj(), ref_psi_src.detach()) | ||
| b = torch.outer(hamiltonian_psi_dst.conj(), ref_psi_src) - torch.outer(psi_src.conj(), ref_hamiltonian_psi_dst) | ||
| diff = (a - self.evolution_time * b).flatten() |
Copilot
AI
Jul 17, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[nitpick] The variable name 'a' is not descriptive. Consider renaming it to something more meaningful like 'overlap_diff' or 'reference_term' to clarify its role in the loss calculation.
| a = torch.outer(psi_src.detach().conj(), ref_psi_src) - torch.outer(psi_src.conj(), ref_psi_src.detach()) | |
| b = torch.outer(hamiltonian_psi_dst.conj(), ref_psi_src) - torch.outer(psi_src.conj(), ref_hamiltonian_psi_dst) | |
| diff = (a - self.evolution_time * b).flatten() | |
| overlap_diff = torch.outer(psi_src.detach().conj(), ref_psi_src) - torch.outer(psi_src.conj(), ref_psi_src.detach()) | |
| b = torch.outer(hamiltonian_psi_dst.conj(), ref_psi_src) - torch.outer(psi_src.conj(), ref_hamiltonian_psi_dst) | |
| diff = (overlap_diff - self.evolution_time * b).flatten() |
qmb/rlim.py
Outdated
| b = torch.outer(hamiltonian_psi_dst.conj(), ref_psi_src) - torch.outer(psi_src.conj(), ref_hamiltonian_psi_dst) | ||
| diff = (a - self.evolution_time * b).flatten() |
Copilot
AI
Jul 17, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[nitpick] The variable name 'b' is not descriptive. Consider renaming it to something more meaningful like 'hamiltonian_term' or 'energy_term' to clarify its role in the loss calculation.
| b = torch.outer(hamiltonian_psi_dst.conj(), ref_psi_src) - torch.outer(psi_src.conj(), ref_hamiltonian_psi_dst) | |
| diff = (a - self.evolution_time * b).flatten() | |
| hamiltonian_term = torch.outer(hamiltonian_psi_dst.conj(), ref_psi_src) - torch.outer(psi_src.conj(), ref_hamiltonian_psi_dst) | |
| diff = (a - self.evolution_time * hamiltonian_term).flatten() |
| loss.energy = energy # type: ignore[attr-defined] | ||
| return loss | ||
|
|
||
| logging.info("Starting local optimization process") | ||
|
|
||
| for i in range(self.local_step): | ||
| loss: torch.Tensor = optimizer.step(closure) # type: ignore[assignment,arg-type] | ||
| energy: float = loss.energy # type: ignore[attr-defined] |
Copilot
AI
Jul 17, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Dynamically adding attributes to tensor objects is not a clean practice. Consider returning a tuple or using a dataclass to pass both loss and energy values instead of monkey-patching the tensor object.
| loss.energy = energy # type: ignore[attr-defined] | |
| return loss | |
| logging.info("Starting local optimization process") | |
| for i in range(self.local_step): | |
| loss: torch.Tensor = optimizer.step(closure) # type: ignore[assignment,arg-type] | |
| energy: float = loss.energy # type: ignore[attr-defined] | |
| return LossEnergy(loss=loss, energy=energy) | |
| logging.info("Starting local optimization process") | |
| for i in range(self.local_step): | |
| loss_energy: LossEnergy = optimizer.step(closure) # type: ignore[assignment,arg-type] | |
| loss, energy = loss_energy.loss, loss_energy.energy |
|
|
||
| for i in range(self.local_step): | ||
| loss: torch.Tensor = optimizer.step(closure) # type: ignore[assignment,arg-type] | ||
| energy: float = loss.energy # type: ignore[attr-defined] |
Copilot
AI
Jul 17, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Accessing the dynamically added energy attribute requires type ignore comments and makes the code fragile. This is a consequence of the monkey-patching approach on line 117.
Description
todo:
Checklist: