API changes

GFNOrg · Oct 1, 2024 · 4434e5f · 4434e5f
1 parent 738b062
commit 4434e5f
Show file tree

Hide file tree

Showing 7 changed files with 20 additions and 18 deletions.
diff --git a/tutorials/examples/train_box.py b/tutorials/examples/train_box.py
@@ -230,8 +230,9 @@ def main(args):  # noqa: C901
         if iteration % 1000 == 0:
             print(f"current optimizer LR: {optimizer.param_groups[0]['lr']}")
 
+        # Sampling on-policy, so we save logprobs for faster computation.
         trajectories = gflownet.sample_trajectories(
-            env, save_logprobs=True, n_samples=args.batch_size
+            env, save_logprobs=True, n=args.batch_size
         )
 
         training_samples = gflownet.to_training_samples(trajectories)
@@ -241,7 +242,7 @@ def main(args):  # noqa: C901
 
         loss.backward()
         for p in gflownet.parameters():
-            if p.ndim > 0 and p.grad is not None:  # We do not clip logZ grad
+            if p.ndim > 0 and p.grad is not None:  # We do not clip logZ grad.
                 p.grad.data.clamp_(-10, 10).nan_to_num_(0.0)
         optimizer.step()
         scheduler.step()

diff --git a/tutorials/examples/train_discreteebm.py b/tutorials/examples/train_discreteebm.py
@@ -63,15 +63,14 @@ def main(args):  # noqa: C901
     optimizer = torch.optim.Adam(module.parameters(), lr=args.lr)
 
     # 4. Train the gflownet
-
     visited_terminating_states = env.states_from_batch_shape((0,))
 
     states_visited = 0
     n_iterations = args.n_trajectories // args.batch_size
     validation_info = {"l1_dist": float("inf")}
     for iteration in trange(n_iterations):
         trajectories = gflownet.sample_trajectories(
-            env, save_logprobs=True, n_samples=args.batch_size
+            env, save_logprobs=True, n=args.batch_size
         )
         training_samples = gflownet.to_training_samples(trajectories)
 

diff --git a/tutorials/examples/train_hypergrid.py b/tutorials/examples/train_hypergrid.py
@@ -229,7 +229,7 @@ def main(args):  # noqa: C901
     for iteration in trange(n_iterations):
         trajectories = gflownet.sample_trajectories(
             env,
-            n_samples=args.batch_size,
+            n=args.batch_size,
             save_logprobs=args.replay_buffer_size == 0,
             save_estimator_outputs=False,
         )

diff --git a/tutorials/examples/train_hypergrid_simple.py b/tutorials/examples/train_hypergrid_simple.py
@@ -5,7 +5,6 @@
 from gfn.gflownet import TBGFlowNet
 from gfn.gym import HyperGrid
 from gfn.modules import DiscretePolicyEstimator
-from gfn.samplers import Sampler
 from gfn.utils import NeuralNet
 
 torch.manual_seed(0)
@@ -35,10 +34,7 @@
 pb_estimator = DiscretePolicyEstimator(
     module_PB, env.n_actions, is_backward=True, preprocessor=env.preprocessor
 )
-gflownet = TBGFlowNet(init_logZ=0.0, pf=pf_estimator, pb=pb_estimator)
-
-# Feed pf to the sampler.
-sampler = Sampler(estimator=pf_estimator)
+gflownet = TBGFlowNet(logZ=0.0, pf=pf_estimator, pb=pb_estimator)
 
 # Move the gflownet to the GPU.
 if torch.cuda.is_available():
@@ -53,9 +49,9 @@
 batch_size = int(1e5)
 
 for i in (pbar := tqdm(range(n_iterations))):
-    trajectories = sampler.sample_trajectories(
+    trajectories = gflownet.sample_trajectories(
         env,
-        n_trajectories=batch_size,
+        n=batch_size,
         save_logprobs=False,
         save_estimator_outputs=True,
         epsilon=exploration_rate,
@@ -64,4 +60,4 @@
     loss = gflownet.loss(env, trajectories)
     loss.backward()
     optimizer.step()
-    pbar.set_postfix({"loss": loss.item()})
+    pbar.set_postfix({"loss": loss.item()})
diff --git a/tutorials/examples/train_hypergrid_simple_conditional.py b/tutorials/examples/train_hypergrid_simple_conditional.py
@@ -97,9 +97,9 @@
     conditioning = torch.rand((batch_size, 1))
     conditioning = (conditioning > 0.5).to(torch.float)  # Randomly 1 and zero.
 
-    trajectories = sampler.sample_trajectories(
+    trajectories = gflownet.sample_trajectories(
         env,
-        n_trajectories=batch_size,
+        n=batch_size,
         conditioning=conditioning,
         save_logprobs=False,
         save_estimator_outputs=True,

diff --git a/tutorials/examples/train_ising.py b/tutorials/examples/train_ising.py
@@ -83,8 +83,14 @@ def ising_n_to_ij(L, n):
     # Learning
     visited_terminating_states = env.States.from_batch_shape((0,))
     states_visited = 0
+
     for i in (pbar := tqdm(range(10000))):
-        trajectories = gflownet.sample_trajectories(env, n_samples=8, off_policy=False)
+        trajectories = gflownet.sample_trajectories(
+            env,
+            n=8,
+            save_estimator_outputs=False,
+            save_logprobs=True,
+        )
         training_samples = gflownet.to_training_samples(trajectories)
         optimizer.zero_grad()
         loss = gflownet.loss(env, training_samples)

diff --git a/tutorials/examples/train_line.py b/tutorials/examples/train_line.py
@@ -227,7 +227,7 @@ def train(
         # Off Policy Sampling.
         trajectories = gflownet.sample_trajectories(
             env,
-            n_samples=batch_size,
+            n=batch_size,
             save_estimator_outputs=True,
             save_logprobs=False,
             scale_factor=scale_schedule[iteration],  # Off policy kwargs.
@@ -292,7 +292,7 @@ def train(
         policy_std_max=policy_std_max,
     )
     pb = StepEstimator(environment, pb_module, backward=True)
-    gflownet = TBGFlowNet(pf=pf, pb=pb, init_logZ=0.0)
+    gflownet = TBGFlowNet(pf=pf, pb=pb, logZ=0.0)
 
     gflownet = train(
         gflownet,