Blame for brain_lab/java/BrainLabMario.java

import engine.core.*;
import engine.helper.*;
import java.io.*;
import java.nio.file.*;
import java.util.*;

/** Local adapter; the downloaded Mario AI engine handles all game physics. */
public class BrainLabMario {
    static final int ACTIONS=6, HOLD=4;
    static final String[] NAMES={"Run right","Run + jump","Move left","Left + jump","Wait","Jump in place"};
    HashMap<Long,double[]> q=new HashMap<>();
    Random rng=new Random(7);
    ArrayList<Runner> runners=new ArrayList<>();
    ArrayList<String> history=new ArrayList<>();
    static final int FINAL_LEVEL=16;
    String level,course="campaign"; int[][] originalTiles;
    int levelNumber=1; boolean campaignComplete=false,autoAdvance=true;
    int population=32,round=0,tick=0,wins=0,deaths=0,episodes=0;
    long experiences=0,updates=0;
    double explore=.35,deathPenalty=12;
    Runner solo;
    static class ObservedWorld extends MarioWorld {
        ArrayList<MarioSprite> objects=new ArrayList<>();
        ObservedWorld(){super(null);}
        @Override public void addSprite(MarioSprite s){super.addSprite(s);objects.add(s);}
        @Override public void removeSprite(MarioSprite s){super.removeSprite(s);objects.remove(s);}
    }
    static class Runner {
        ObservedWorld w; boolean done=false; int age=0,stale=0,action=0,held=0; long key;
        double best=0,reward=0,partReward=0; String reason="running";
        int levelNumber; int[][] baseTiles; int bossHP=3,bossCooldown=0; double bossX=86*16,bossY=224,bossVY=0; boolean axe=false;
        ArrayList<double[]> flames=new ArrayList<>();
        Runner(String text,int number) {
            levelNumber=number;w=new ObservedWorld();w.visuals=false;w.initializeLevel(text,180000);w.update(new boolean[5]);best=w.mario.x;
            baseTiles=new int[w.level.tileHeight][w.level.tileWidth];for(int y=0;y<baseTiles.length;y++)for(int x=0;x<baseTiles[y].length;x++)baseTiles[y][x]=w.level.getBlock(x,y);
        }
    }
    static String quote(String s){return "\""+s.replace("\\","\\\\").replace("\"","\\\"").replace("\n","\\n")+"\"";}
    static String num(double x){return Double.isFinite(x)?String.format(Locale.ROOT,"%.4f",x):"0";}
    double[] peek(long k){return q.getOrDefault(k,new double[ACTIONS]);}
    double[] values(long k){return q.computeIfAbsent(k,key->new double[ACTIONS]);}
    static int max(double[] a){int k=0;for(int i=1;i<a.length;i++)if(a[i]>a[k])k=i;return k;}
    static boolean[] action(int a){boolean[] b=new boolean[5];b[3]=true;b[1]=a==0||a==1;b[0]=a==2||a==3;b[4]=a==1||a==3||a==5;return b;}
    boolean solid(MarioWorld w,int x,int y){int t=w.level.getBlock(x,y);return t>0&&t!=15&&t!=47;}
    long key(Runner r){
        MarioWorld w=r.w;int x=(int)(w.mario.x/16),y=(int)((w.mario.y-1)/16);long k=0;
        for(int dx=0;dx<5;dx++)for(int dy=-2;dy<=1;dy++)k=(k<<1)|(solid(w,x+dx,y+dy)?1:0);
        int gap=5;for(int d=0;d<5;d++){boolean floor=false;for(int yy=y+1;yy<16;yy++)floor|=solid(w,x+d,yy);if(!floor){gap=d;break;}}
        int enemy=5;for(MarioSprite e:w.getEnemies())if(e.x>=w.mario.x-8&&Math.abs(e.y-w.mario.y)<32)enemy=Math.min(enemy,Math.max(0,(int)((e.x-w.mario.x)/16)));
        k=k*6+gap;k=k*6+enemy;k=k*2+(w.mario.onGround?1:0);k=k*2+(w.mario.mayJump?1:0);
        k=k*3+(w.mario.ya< -1?0:w.mario.ya>1?2:1);k=k*3+(w.mario.xa< -1?0:w.mario.xa>1?2:1);
        k=k*4+((int)w.mario.x%16)/4;
        if(r.levelNumber==FINAL_LEVEL){k=k*8+Math.min(7,Math.max(0,(int)((r.bossX-w.mario.x)/16)+3));k=k*4+r.bossHP;k=k*2+(r.bossY<210?1:0);k=k*2+(r.flames.stream().anyMatch(f->Math.abs(f[0]-w.mario.x)<64)?1:0);k|=1L<<60;}
        return k;
    }
    boolean awaitingNext(){return course.equals("campaign")&&!campaignComplete&&levelNumber<FINAL_LEVEL&&runners.stream().allMatch(r->r.done)&&runners.stream().anyMatch(r->r.reason.equals("win"));}
    void freshRound() throws IOException{
        if(!runners.isEmpty()){
            double progress=0,best=0,ret=0;int won=0;
            for(Runner r:runners){double p=progress(r);progress+=p;best=Math.max(best,p);ret+=r.reward;if(r.reason.equals("win"))won++;}
            history.add("{\"round\":"+round+",\"mean_progress\":"+num(progress/population)+",\"best_progress\":"+num(best)+",\"wins\":"+won+",\"mean_reward\":"+num(ret/population)+"}");
            if(history.size()>160)history.remove(0);
        }
        if(course.equals("campaign")&&runners.stream().anyMatch(r->r.reason.equals("win"))){
            if(levelNumber==FINAL_LEVEL){campaignComplete=true;return;}
            levelNumber++;level=levelText(levelNumber);history.clear();
        }
        round++;tick=0;runners.clear();for(int i=0;i<population;i++)runners.add(new Runner(level,levelNumber));
        originalTiles=runners.get(0).baseTiles;
    }
    double progress(Runner r){double start=r.w.level.marioTileX*16;return Math.max(0,Math.min(1,(r.best-start)/(r.w.level.exitTileX*16-start)));}
    void advance(Runner r,boolean learning,Integer manual){
        if(r.done)return;
        if(r.held==0){r.key=key(r);r.action=manual!=null?manual:learning&&rng.nextDouble()<explore?rng.nextInt(ACTIONS):max(learning?values(r.key):peek(r.key));r.partReward=0;}
        double oldBest=r.best,oldY=r.w.mario.y;int oldHP=r.bossHP,oldCoins=r.w.coins;r.w.update(action(r.action));r.age++;r.held++;
        if(r.levelNumber==FINAL_LEVEL)bossTick(r,oldY);
        r.best=Math.max(r.best,r.w.mario.x);r.stale=r.best>oldBest+.1?0:r.stale+1;
        double reward=(r.best-oldBest)/16-.01+2*(r.w.coins-oldCoins)+10*(oldHP-r.bossHP);
        if(r.w.gameStatus==GameStatus.WIN){r.done=true;r.reason="win";reward+=50;}
        else if(r.w.gameStatus!=GameStatus.RUNNING||r.w.mario.y>270){r.done=true;r.reason="fell or hit an enemy";reward-=deathPenalty;}
        else if((manual==null&&r.stale>180)||r.age>=Math.max(2200,r.w.level.tileWidth*24)){r.done=true;r.reason=r.stale>180?"stuck":"time limit";reward-=4;}
        r.reward+=reward;r.partReward+=reward;
        if(learning){experiences++;if(r.held>=HOLD||r.done){double[] row=values(r.key);double target=r.partReward+(r.done?0:Math.pow(.98,r.held)*Arrays.stream(values(key(r))).max().orElse(0));row[r.action]+=.18*(target-row[r.action]);updates++;}}
        if(r.held>=HOLD)r.held=0;
        if(learning&&r.done){episodes++;if(r.reason.equals("win"))wins++;else deaths++;}
    }
    String runner(Runner r,int id){return "{\"id\":"+id+",\"x\":"+num(r.w.mario.x)+",\"y\":"+num(r.w.mario.y)+",\"vx\":"+num(r.w.mario.xa)+",\"grounded\":"+r.w.mario.onGround+",\"large\":"+r.w.mario.isLarge+",\"fire\":"+r.w.mario.isFire+",\"done\":"+r.done+",\"reason\":"+quote(r.reason)+",\"progress\":"+num(progress(r))+",\"action\":"+r.action+"}";}
    String changes(Runner r){ArrayList<String> a=new ArrayList<>();if(r.baseTiles!=null)for(int y=0;y<r.w.level.tileHeight;y++)for(int x=0;x<r.w.level.tileWidth;x++){int t=r.w.level.getBlock(x,y);if(t!=r.baseTiles[y][x])a.add("["+x+","+y+","+t+"]");}return "["+String.join(",",a)+"]";}
    String objects(Runner r){ArrayList<String> a=new ArrayList<>();for(MarioSprite s:r.w.objects)if(s.alive&&s.type.getValue()>=12)a.add("["+num(s.x)+","+num(s.y)+","+s.type.getValue()+"]");return "["+String.join(",",a)+"]";}
    String enemies(Runner r){ArrayList<String>a=new ArrayList<>();for(MarioSprite s:r.w.getEnemies())a.add("["+num(s.x)+","+num(s.y)+","+s.type.getValue()+"]");return "["+String.join(",",a)+"]";}
    String levelInfo(Runner r){return "\"level_number\":"+r.levelNumber+",\"level_name\":"+quote(r.levelNumber==0?(course.equals("gaps")?"Gap school":"Jump school"):r.levelNumber==FINAL_LEVEL?"Level 16 · Bowser’s castle":"Level "+r.levelNumber)+",\"total_levels\":"+FINAL_LEVEL;}
    String map(Runner r){StringBuilder s=new StringBuilder("{\"width\":"+r.w.level.width+",\"exit_x\":"+r.w.level.exitTileX*16+",\"tiles\":[");for(int y=0;y<r.baseTiles.length;y++){if(y>0)s.append(',');s.append(Arrays.toString(r.baseTiles[y]));}return s.append("]}").toString();}
    String boss(Runner r){
        if(r.levelNumber!=FINAL_LEVEL)return "null";
        ArrayList<String> f=new ArrayList<>();for(double[] a:r.flames)f.add("["+num(a[0])+","+num(a[1])+"]");
        return "{\"x\":"+num(r.bossX)+",\"y\":"+num(r.bossY)+",\"hp\":"+r.bossHP+",\"max_hp\":3,\"axe_x\":1600,\"axe_activated\":"+r.axe+",\"fireballs\":["+String.join(",",f)+"]}";
    }
    String soloFrame(Runner r,boolean includeMap){return "{"+levelInfo(r)+",\"runner\":"+runner(r,0)+",\"enemies\":"+enemies(r)+",\"tile_changes\":"+changes(r)+",\"items\":"+objects(r)+",\"boss\":"+boss(r)+",\"next_level\":"+(course.equals("campaign")&&r.done&&r.reason.equals("win")&&r.levelNumber<FINAL_LEVEL?r.levelNumber+1:0)+",\"campaign_complete\":"+(r.levelNumber==FINAL_LEVEL&&r.reason.equals("win"))+(includeMap?",\"map\":"+map(r):"")+"}";}
    String frame(){return frame(false);}
    String frame(boolean includeMap){
        int lead=0;for(int i=1;i<runners.size();i++)if((!runners.get(i).done&&runners.get(lead).done)||runners.get(i).done==runners.get(lead).done&&runners.get(i).best>runners.get(lead).best)lead=i;
        ArrayList<String>a=new ArrayList<>();for(int i=0;i<runners.size();i++)a.add(runner(runners.get(i),i));
        return "{"+levelInfo(runners.get(lead))+",\"boss\":"+boss(runners.get(lead))+(includeMap?",\"map\":"+map(runners.get(lead)):"")+",\"round\":"+round+",\"tick\":"+tick+",\"leader\":"+lead+",\"runners\":["+String.join(",",a)+"],\"enemies\":"+enemies(runners.get(lead))+",\"tile_changes\":"+changes(runners.get(lead))+",\"items\":"+objects(runners.get(lead))+"}";
    }
    String state(boolean map){
        long active=runners.stream().filter(r->!r.done).count();Runner lead=runners.stream().filter(r->!r.done).max(Comparator.comparingDouble(r->r.best)).orElse(runners.get(0));
        StringBuilder s=new StringBuilder("{\"auto_advance\":"+autoAdvance+",\"awaiting_next_level\":"+awaitingNext()+","+levelInfo(lead)+",\"campaign_complete\":"+campaignComplete+",\"course\":"+quote(course)+",\"population\":"+population+",\"round\":"+round+",\"tick\":"+tick+",\"active\":"+active+",\"experiences\":"+experiences+",\"updates\":"+updates+",\"states\":"+q.size()+",\"episodes\":"+episodes+",\"wins\":"+wins+",\"deaths\":"+deaths+",\"exploration\":"+explore+",\"death_penalty\":"+deathPenalty+",\"q_values\":"+Arrays.toString(peek(key(lead)))+",\"history\":["+String.join(",",history)+"],\"frame\":"+frame());
        if(map){s.append(",\"width\":").append(lead.w.level.width).append(",\"exit_x\":").append(lead.w.level.exitTileX*16).append(",\"tiles\":[");for(int y=0;y<lead.w.level.tileHeight;y++){if(y>0)s.append(',');s.append('[');for(int x=0;x<lead.w.level.tileWidth;x++){if(x>0)s.append(',');s.append(originalTiles==null?lead.w.level.getBlock(x,y):originalTiles[y][x]);}s.append(']');}s.append(']');}
        return s.append('}').toString();
    }
    static String levelText(int number)throws IOException{
        if(number==FINAL_LEVEL)return castle();
        if(number<1||number>15)throw new IllegalArgumentException("Choose level 1 through 16.");
        return Files.readString(Path.of("vendor/mario-ai/levels/original/lvl-"+number+".txt"));
    }
    void reset(int n,String c,int seed)throws Exception{
        int number=c.equals("campaign")||c.equals("original")?1:c.equals("bowser")?FINAL_LEVEL:c.startsWith("level-")?Integer.parseInt(c.substring(6)):0;
        if(number==0&&!c.equals("practice")&&!c.equals("gaps"))throw new IllegalArgumentException("Unknown course");
        String text=number>0?levelText(number):practice(c.equals("gaps"));
        population=Math.max(1,Math.min(128,n));course=c;levelNumber=number;level=text;campaignComplete=false;rng=new Random(seed);q.clear();history.clear();round=0;experiences=updates=0;wins=deaths=episodes=0;runners.clear();
        freshRound();solo=null;
    }
    static String castle(){
        char[][] g=new char[16][112];for(char[] row:g)Arrays.fill(row,'-');
        for(int x=0;x<112;x++){g[14][x]='X';g[15][x]='X';g[0][x]='S';}
        g[13][2]='M';g[13][108]='F';
        for(int x:new int[]{17,51}){g[13][x]='#';g[13][x+1]='#';}
        for(int x:new int[]{32,33,68,69}){g[14][x]='-';g[15][x]='-';}
        for(int x=74;x<=97;x++){g[14][x]='#';g[15][x]='-';}
        g[10][12]='?';g[10][53]='?';for(int x:new int[]{20,40,62})g[10][x]='o';
        ArrayList<String>a=new ArrayList<>();for(char[] row:g)a.add(new String(row));return String.join("\n",a);
    }
    void bossTick(Runner r,double oldY){
        if(r.w.gameStatus==GameStatus.WIN&&r.bossHP>0){r.w.gameStatus=GameStatus.RUNNING;r.w.mario.x=1590;r.w.mario.xa=0;}
        if(r.w.gameStatus!=GameStatus.RUNNING)return;
        if(r.bossCooldown>0)r.bossCooldown--;
        if(r.bossHP>0){
            r.bossX=86*16+Math.sin(r.age/46.0)*56;
            if(r.age%150==0&&r.w.mario.x>64*16)r.bossVY=-7;
            r.bossY+=r.bossVY;r.bossVY+=.4;if(r.bossY>224){r.bossY=224;r.bossVY=0;}
            if(r.age%72==0&&r.w.mario.x>60*16)r.flames.add(new double[]{r.bossX-20,r.bossY-18});
            if(Math.abs(r.w.mario.x-r.bossX)<23&&r.w.mario.y>r.bossY-40&&r.w.mario.y-r.w.mario.height<r.bossY){
                if(r.bossCooldown==0&&oldY<=r.bossY-30&&r.w.mario.ya>0){r.bossHP--;r.bossCooldown=24;r.w.mario.ya=-11;r.w.mario.jumpTime=0;r.w.mario.onGround=false;}
                else if(r.bossCooldown==0)r.w.mario.getHurt();
            }
            if(r.w.mario.x>=100*16&&r.w.mario.y>=176){r.axe=true;r.bossHP=0;for(int x=74;x<=97;x++)r.w.level.setBlock(x,14,0);}
        }else r.bossY+=4;
        for(Iterator<double[]> it=r.flames.iterator();it.hasNext();){double[] f=it.next();f[0]-=2.8;if(f[0]<0){it.remove();continue;}if(Math.abs(f[0]-r.w.mario.x)<9&&f[1]>r.w.mario.y-r.w.mario.height&&f[1]<r.w.mario.y){r.w.mario.getHurt();it.remove();}}
    }
    static String practice(boolean gaps){
        char[][] g=new char[16][76];for(char[] row:g)Arrays.fill(row,'-');
        for(int x=0;x<76;x++){g[14][x]='X';g[15][x]='X';}
        g[13][2]='M';g[13][72]='F';
        for(int x:new int[]{16,30,48,61}){g[13][x]='#';g[13][x+1]='#';if(x==30||x==61){g[12][x]='#';g[12][x+1]='#';}}
        for(int x:new int[]{12,26,43,57,68})g[10][x]='o';
        if(gaps)for(int x:new int[]{23,24,39,40,41,54,55}){g[14][x]='-';g[15][x]='-';}
        ArrayList<String>a=new ArrayList<>();for(char[] row:g)a.add(new String(row));return String.join("\n",a);
    }
    void save(String path)throws Exception{try(DataOutputStream out=new DataOutputStream(new FileOutputStream(path))){out.writeUTF("BrainLabMarioQ3");out.writeUTF(course);out.writeInt(levelNumber);out.writeBoolean(campaignComplete);out.writeBoolean(autoAdvance);out.writeLong(experiences);out.writeLong(updates);out.writeInt(q.size());for(long key:new TreeSet<>(q.keySet())){out.writeLong(key);for(double v:q.get(key))out.writeDouble(v);}}}
    void load(String path)throws Exception{try(DataInputStream in=new DataInputStream(new FileInputStream(path))){String version=in.readUTF();if(!version.equals("BrainLabMarioQ1")&&!version.equals("BrainLabMarioQ2")&&!version.equals("BrainLabMarioQ3"))throw new IOException("Invalid checkpoint");String c=in.readUTF();int savedLevel=!version.equals("BrainLabMarioQ1")?in.readInt():0;boolean complete=!version.equals("BrainLabMarioQ1")&&in.readBoolean();boolean savedAuto=!version.equals("BrainLabMarioQ3")||in.readBoolean();reset(population,c,7);autoAdvance=savedAuto;if(c.equals("campaign")){if(savedLevel<1||savedLevel>FINAL_LEVEL)throw new IOException("Invalid campaign level");levelNumber=savedLevel;level=levelText(levelNumber);runners.clear();round=0;freshRound();}campaignComplete=complete;experiences=in.readLong();updates=in.readLong();int count=in.readInt();if(count<0||count>2000000)throw new IOException("Checkpoint too large");for(int i=0;i<count;i++){long key=in.readLong();double[] a=new double[ACTIONS];for(int j=0;j<ACTIONS;j++){a[j]=in.readDouble();if(!Double.isFinite(a[j]))throw new IOException("Invalid Q value");}q.put(key,a);}}}
    String command(String line)throws Exception{
        String[] a=line.trim().split(" +");String op=a[0];
        if(op.equals("reset")){reset(Integer.parseInt(a[1]),a[2],Integer.parseInt(a[3]));return state(true);}
        if(op.equals("progression")){autoAdvance=Boolean.parseBoolean(a[1]);return state(true);}
        if(op.equals("next_level")){if(!awaitingNext())throw new IllegalArgumentException("Clear the current training level first.");freshRound();return state(true);}
        if(op.equals("human_next")){if(solo==null||!course.equals("campaign")||!solo.done||!solo.reason.equals("win")||solo.levelNumber>=FINAL_LEVEL)throw new IllegalArgumentException("Clear the current level first.");solo=new Runner(levelText(solo.levelNumber+1),solo.levelNumber+1);return soloFrame(solo,true);}
        if(op.equals("state"))return state(true);
        if(op.equals("config")){explore=Double.parseDouble(a[1]);deathPenalty=Double.parseDouble(a[2]);return state(false);}
        if(op.equals("train")){
            int count=Math.max(1,Math.min(1200,Integer.parseInt(a[1]))),every=Math.max(1,count/24),shownLevel=-1;ArrayList<String> frames=new ArrayList<>();
            for(int t=0;t<count;t++){if(campaignComplete)break;if(runners.stream().allMatch(r->r.done)){if(!autoAdvance&&awaitingNext())break;freshRound();}if(campaignComplete)break;for(Runner r:runners)advance(r,true,null);tick++;if(t%every==0||t==count-1){frames.add(frame(shownLevel!=levelNumber));shownLevel=levelNumber;}}
            if(frames.isEmpty()||campaignComplete||(!autoAdvance&&awaitingNext()))frames.add(frame(true));String s=state(true);return s.substring(0,s.length()-1)+",\"frames\":["+String.join(",",frames)+"]}";
        }
        if(op.equals("evaluate")){
            int count=a.length>1?Integer.parseInt(a[1]):1;ArrayList<String> all=new ArrayList<>();
            for(int i=0;i<count;i++){Runner r=new Runner(level,levelNumber);ArrayList<String> frames=new ArrayList<>();while(!r.done){advance(r,false,null);if(r.age%4==0||r.done)frames.add(soloFrame(r,frames.isEmpty()));}all.add("{\"progress\":"+num(progress(r))+",\"win\":"+r.reason.equals("win")+",\"reason\":"+quote(r.reason)+",\"frames\":["+String.join(",",frames)+"]}");}
            return "{\"training_experiences\":"+experiences+",\"updates\":"+updates+",\"evaluations\":["+String.join(",",all)+"]}";
        }
        if(op.equals("human_reset")){int n=solo!=null&&solo.done&&!solo.reason.equals("win")?solo.levelNumber:1;solo=course.equals("campaign")?new Runner(levelText(n),n):new Runner(level,levelNumber);return soloFrame(solo,true);}
        if(op.equals("human")){
            if(solo==null)solo=new Runner(level,levelNumber);
            if(autoAdvance&&course.equals("campaign")&&solo.done&&solo.reason.equals("win")&&solo.levelNumber<FINAL_LEVEL)solo=new Runner(levelText(solo.levelNumber+1),solo.levelNumber+1);
            int act=Integer.parseInt(a[1]);if(act<0||act>=ACTIONS)throw new IllegalArgumentException("Invalid action");
            for(int i=0;i<4;i++){solo.held=0;advance(solo,false,act);}return soloFrame(solo,true);
        }
        if(op.equals("save")){save(a[1]);return "{\"saved\":true}";}
        if(op.equals("load")){load(a[1]);return state(true);}
        throw new IllegalArgumentException("Unknown command");
    }
    public static void main(String[] args)throws Exception{
        BrainLabMario lab=new BrainLabMario();lab.reset(32,"campaign",7);
        BufferedReader reader=new BufferedReader(new InputStreamReader(System.in));String line;
        while((line=reader.readLine())!=null){try{System.out.println(lab.command(line));}catch(Exception e){System.out.println("{\"error\":"+quote(e.toString())+"}");}System.out.flush();}
    }
}