Add more local mutation

offchain/src/bin/symreg_experiment_local.rs

@@ -0,0 +1,202 @@
+// src/bin/symreg_experiment_local.rs
+
+use anyhow::Result;
+use rand::{thread_rng, Rng};
+
+// Our GP + compiler modules (adjust paths if needed)
+use offchain::compiler::ast::{UntypedAst, Push3Ast};
+use offchain::helpers::artifact::get_creation_code;
+use offchain::runner::revm_runner::{EvmRunner, Push3InterpreterInputs};
+use offchain::compiler::push3_describtor::make_sublist_descriptor;
+
+// Our random code and local mutation
+use offchain::gp::generate_spec::ranmdom_code_fixed;
+use offchain::gp::local_mutation::local_mutation_fixed;
+
+// Ethers + ABI
+use ethers::types::U256;
+use ethers::abi::{encode, Token};
+
+/// 1) Generate (x, y) samples for f(x) = 3x^2 + x + 3
+fn generate_samples() -> Vec<(i32, i32)> {
+    let mut samples = Vec::new();
+    for x in -5..=5 {
+        let y = 3 * x * x + x + 3;
+        samples.push((x, y));
+    }
+    samples
+}
+
+/// 2) Evaluate a single AST on a single input `x`, returning the final top of the int stack.
+fn evaluate_ast_on_x(
+    runner: &mut EvmRunner,
+    ast: &UntypedAst,
+    x: i32,
+) -> i32 {
+    // Convert AST => push3 bytecode
+    let code_bytes = ast.to_bytecode();
+    let code_len = code_bytes.len() as u32;
+
+    // Build a sublist descriptor
+    let descriptor = make_sublist_descriptor(0, code_len);
+
+    // Provide x in `init_int_stack`
+    let inputs = Push3InterpreterInputs {
+        code: code_bytes,
+        init_code_stack: Vec::new(),
+        init_exec_stack: vec![descriptor],
+        init_int_stack: vec![x as i128],
+    };
+
+    // Run
+    let outputs = match runner.run_interpreter(&inputs) {
+        Ok(o) => o,
+        Err(_) => {
+            // If revert => 0
+            return 0;
+        }
+    };
+
+    if outputs.final_int_stack.is_empty() {
+        i32::MAX  // or 0, or i32::MAX to strongly penalize weird runs
+    } else {
+        *outputs.final_int_stack.last().unwrap() as i32
+    }
+}
+
+/// 3) Evaluate an AST on all samples => compute MSE
+fn evaluate_fitness(
+    runner: &mut EvmRunner,
+    ast: &UntypedAst,
+    samples: &[(i32, i32)]
+) -> f64 {
+    let mut error_sum = 0.0;
+    for &(x, target_y32) in samples {
+        // cast to i64
+        let target_y = target_y32 as i64;
+        
+        // evaluate in i64
+        let predicted_i64 = evaluate_ast_on_x(runner, ast, x) as i64;
+        
+        // wrapping_sub to avoid panic
+        let diff = predicted_i64.wrapping_sub(target_y);
+        error_sum += (diff as f64) * (diff as f64);
+    }
+    error_sum / samples.len() as f64
+}
+
+fn main() -> Result<()> {
+    // 1) Generate samples
+    let samples = generate_samples();
+
+    // 2) Setup ephemeral EVM runner
+    let creation_hex_filename = "../onchain/out/Push3Interpreter.sol/Push3Interpreter.json";
+    let creation_bytes = get_creation_code(creation_hex_filename)?;
+    let mut runner = EvmRunner::new(creation_bytes)?;
+
+    // 3) GP parameters
+    let pop_size = 3600;  // must be multiple of 3 for our scheme
+    let generations = 100;
+    let early_stop_threshold = 1.0;
+    let max_points = 12;  // random code size limit
+
+    // Ensure pop_size is multiple of 3
+    assert_eq!(pop_size % 3, 0, "pop_size must be multiple of 3 for our scheme");
+    let third = pop_size / 3;
+
+    let mut rng = thread_rng();
+
+    // 4) Initialize random population
+    let mut population: Vec<UntypedAst> = (0..pop_size)
+        .map(|_| ranmdom_code_fixed(&mut rng, max_points))
+        .collect();
+
+    // 5) Main loop
+    for gen in 0..generations {
+        // a) Evaluate
+        let mut scored: Vec<(UntypedAst, f64)> = population.into_iter()
+            .map(|ast| {
+                let err = evaluate_fitness(&mut runner, &ast, &samples);
+                (ast, err)
+            })
+            .collect();
+        
+        // b) Sort ascending by error
+        scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
+
+        let best_err = scored[0].1;
+        println!("\n=== Generation {gen} ===");
+        println!("Best error = {best_err}");
+
+        // (Optionally) Print top 5 in pretty format
+        let top_n = 5.min(scored.len());
+        println!("Top {top_n} subjects:");
+        for i in 0..top_n {
+            println!(
+                "  #{i}, err={}, AST:\n{:#?}",
+                scored[i].1, 
+                scored[i].0
+            );
+        }
+
+        // check early stop
+        if best_err < early_stop_threshold {
+            println!("Best error < {early_stop_threshold}, stopping early!");
+            population = scored.into_iter().map(|(a, _)| a).collect();
+            break;
+        }
+
+        // c) Build new population:
+        //   - 1/3 = best survivors
+        //   - 1/3 = fresh random
+        //   - 1/3 = local mutation
+        let mut new_pop = Vec::with_capacity(pop_size);
+
+        // keep best third
+        for i in 0..third {
+            new_pop.push(scored[i].0.clone());
+        }
+
+        // random third
+        for _ in 0..third {
+            let fresh = ranmdom_code_fixed(&mut rng, max_points);
+            new_pop.push(fresh);
+        }
+
+        // local mutation third
+        let top_half_count = scored.len() / 2; 
+        // we'll pick parents from top half for local mutation
+        while new_pop.len() < pop_size {
+            let idx = rng.gen_range(0..top_half_count);
+            // we assume you have local_mutation_fixed
+            // if you prefer passing a custom InstructionSet, do local_mutation(..., &instr_set)
+            let mutated = offchain::gp::local_mutation::local_mutation_fixed(
+                &scored[idx].0,
+                &mut rng
+            );
+            new_pop.push(mutated);
+        }
+
+        // store for next generation
+        population = new_pop;
+    }
+
+    // 6) Final => evaluate & sort => top 10
+    println!("\n=== Final Population (Top 10) ===");
+    let mut final_scored: Vec<(UntypedAst, f64)> = population
+        .into_iter()
+        .map(|ast| {
+            let err = evaluate_fitness(&mut runner, &ast, &samples);
+            (ast, err)
+        })
+        .collect();
+    final_scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
+
+    let top_n = 10.min(final_scored.len());
+    for i in 0..top_n {
+        let (ref ast, err) = final_scored[i];
+        println!("Subject #{i}, err={err}, AST:\n{:#?}", ast);
+    }
+
+    Ok(())
+}