LeanDojo as REPL for online RL

[matt-seb-ho]

Hi, I think LeanDojo is a very impressive project, and I'm interested in using LeanDojo as an RL environment (as advertised, I suppose). I noticed that initializing the environment is somewhat costly. The code snippet below took about 30s+ to initialize each theorem. I read in this issue that 1-2 minutes is quite normal. I was wondering what the recommendation is for using LeanDojo in an online RL setup given that this start-up time for each training sample is somewhat prohibitive.

I have several questions:

• Is the approach shown in the code snippet below correct? Is there a better way than initializing a new Dojo for each sample?
• Are there projects that use LeanDojo as an RL environment-- i.e. code examples I can learn from/contributors I can speak to?
• Given that my project doesn't require advanced features (e.g. premise annotations), would I be better off trying to optimize Dojo code (strip features separate from my use case, implement more caching, etc.) or using some other tool/REPL?
• Are there any general recommendations for using LeanDojo as the environment for online RL?

I am looking into these questions myself by checking papers that cite LeanDojo and trying profiling on enter/exit functions to determine where the bottlenecks are, but I would greatly appreciate any guidance!

Thanks for open sourcing this project and thanks in advance for any help!

import json
from time import perf_counter
import numpy as np
from tqdm import tqdm

from dotenv import load_dotenv
load_dotenv() # load GH token env var
from lean_dojo import Dojo, Theorem, LeanGitRepo

with open("data/leandojo_benchmark_4/novel_premises/train.json") as f:
    thm_dicts = json.load(f)

samples = 30
np.random.seed(42)
idxs = np.random.choice(range(len(thm_dicts)), samples, replace=False)
entry_times = []
exit_times = []

for idx in tqdm(idxs, total=samples):
    thm_dict = thm_dicts[idx] 
    thm = Theorem(
        repo=LeanGitRepo(url=thm_dict["url"], commit=thm_dict["commit"]),
        file_path=thm_dict["file_path"],
        full_name=thm_dict["full_name"]
    )
    start = perf_counter()
    with Dojo(thm) as (dojo, initial_state):
        entry_times.append(perf_counter() - start)
        start = perf_counter()
    exit_times.append(perf_counter() - start)

with open("outputs/entry_exit_times.json", 'w') as f:
    json.dump({"idxs": list(idxs), "entry": entry_times, "exit": exit_times}, f, indent=2)
    print("wrote to outputs/entry_exit_times.json")

[darabos]

I've looked into this a bit, and in my case the start time was due to the native "container" copying gigabytes of data several times. I have a hack that cuts this time out. (doragera#3) I think I deleted too much, and that version only works if you've processed the repo with the normal LeanDojo first. Sorry, I should clean it up someday.

[yangky11]

I see. But it looks like these copies are not truly necessary? Maybe we can avoid them using symbolic links, etc.?

[darabos]

Yes! My change completely removes the copies and it still works most of the time. It only breaks the initial tracing — unfortunately I didn't check that when I was working on this. Symbolic links are a good idea! It may be a smaller change and keep the code more tidy than what I did.

[yangky11]

I can take a look when I have time. It would be really helpful if you could share your changes as a reference!

[darabos]

The link is in the first comment! (doragera#3)

[matt-seb-ho]

I re-implemented @darabos's changes as a subclass InitOptimizedDojo with the new behaviour in hooks in this fork. Thank you both!

[yangky11]

Your code looks reasonable to me. If a theorem is in a big file and has a lot of stuff before it, the startup time is naturally going to be slow. Other than this reason, it might be worth some profiling to identify the performance bottleneck.

[yangky11]

I finally get some time to work on this. This is the draft PR. I'll make a few further improvements and merge it into main.

[yangky11]

Released! https://github.com/lean-dojo/LeanDojo/releases/tag/v2.0.0