FirstModelRL / inference.py
Rohannk's picture
Create inference.py
e32f22f verified
Raw
History Blame Contribute Delete
499 Bytes
from fastapi import FastAPI
from env import DeliveryEnv
import random
app = FastAPI()
env = DeliveryEnv()
ACTIONS = ["up", "down", "left", "right", "pickup", "drop"]
@app.get("/")
def home():
return {"message": "Environment running"}
@app.get("/reset")
def reset():
return env.reset()
@app.get("/step")
def step():
action = random.choice(ACTIONS)
state, reward, done, _ = env.step(action)
return {
"state": state,
"reward": reward,
"done": done
}