How to integrate lean-gym with python?

[venom12138]

Great job! I want to use lean-gym as an engine to train my AI just like RL gym, but I don't know how to integrate lean-gym with python. Could you please provide an example code? Thanks!

[Purewhite2019]

The following is the code I`m using. Hope it can help you.

import subprocess
import json
from typing import Dict, List, Any
import os


class LeanGymEnvironment:
    def __init__(self, verbose: bool=False, lean_gym_base: str=None) -> None:
        self.p = subprocess.Popen(['lean', '--run', 'src/repl.lean'], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=lean_gym_base)
        out = self.p.stdout.read1().decode()
        if verbose:
            print(out)
        
    def write(self, line: str) -> None:
        if line[-1] != '\n':
            line += '\n'
        self.p.stdin.write(bytes(line, 'utf-8'))
        self.p.stdin.flush()
    
    def read(self) -> Dict[str, Any]:
        out = self.p.stdout.read1().decode()
        return json.loads(out)
    
    def init_search(self, decl_name: str, namespaces: List[str]=[]):
        if len(namespaces) == 0:
            self.write(f'["init_search", ["{decl_name}", ""]]')
        else:
            cmd = f'["init_search", ["{decl_name}", "'
            for n in namespaces[:-1]:
                cmd += n + '", "'
            cmd += namespaces[-1] + '"]]'
            self.write(cmd)
            
        return self.read()
    
    def run_tac(self, search_id, tactic_state_id, tactic):
        self.write(f'["run_tac",["{search_id}","{tactic_state_id}","{tactic}"]]')
        return self.read()
    
    def clear_search(self, search_id):
        self.write(f'["self", "{search_id}"]')
        return self.read()