Skip to content

Conversation

@hzhangxyz
Copy link
Member

@hzhangxyz hzhangxyz commented Jul 17, 2025

Description

todo:

  1. 自由地调整ref network
  2. 优化s'而非s
  3. 在mlp和attn上测试
  4. 为什么能量没下去?对比vmc?是不是log scale的loss function需要?

Checklist:

@hzhangxyz hzhangxyz requested a review from Copilot July 17, 2025 18:24

This comment was marked as outdated.

Copy link

Copilot AI left a 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
Copy link

Copilot AI Jul 17, 2025

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.

Suggested change
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

Copilot uses AI. Check for mistakes.
qmb/rlim.py Outdated
Comment on lines 107 to 109
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()
Copy link

Copilot AI Jul 17, 2025

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.

Suggested change
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()

Copilot uses AI. Check for mistakes.
qmb/rlim.py Outdated
Comment on lines 108 to 109
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()
Copy link

Copilot AI Jul 17, 2025

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.

Suggested change
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()

Copilot uses AI. Check for mistakes.
Comment on lines +117 to +124
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]
Copy link

Copilot AI Jul 17, 2025

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.

Suggested change
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

Copilot uses AI. Check for mistakes.

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]
Copy link

Copilot AI Jul 17, 2025

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.

Copilot uses AI. Check for mistakes.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants