Emerge-Lab / gpudrive

GPU-acceleration of Nocturne via Madrona
https://arxiv.org/abs/2408.01584
MIT License
227 stars 20 forks source link

[WIP] Unify actions in a single struct #237

Closed aaravpandya closed 2 months ago

aaravpandya commented 2 months ago

Since we now aim to support multiple dynamics models, each of the model come with their own unique set of action spaces. However, if we tried to make a new data structure for each action space we'd introduce a lot of bloat in the code. This PR makes the Action struct a union of individual action spaces. At any given time only one action space is active. This allows us to reduce the bloat, pass in only one Action component to systems, and also easily be able to add in new action spaces as needed. The new types looks like this -

struct ClassicAction {
    float acceleration;
    float steering;
    float headAngle;
};

struct DeltaAction {
    float dx;
    float dy;
    float dyaw;
};

struct StateAction {
    Position position;
    float yaw;
    Velocity velocity;  
};

union Action
{
    ClassicAction classic;
    DeltaAction delta;
    StateAction state;
};

As part of these changes, the Action tensor now is of the size of the largest type defined in the union. So for eg, before the Action tensor was of shape 3, but now it is of shape (3+1+6) as the StateAction is the biggest type here. As such, to use the action tensor, it would be important to know which bicycle model is being used, and the corresponding type size.

for eg. -

# for classic model, we can do the following 
action_tensor[:, :, :3] = classic_action_values

# for state model, we can do 
action_tensor[:, :, :] = state_action_values 

# for an aribitrary model, 
action_tensor[:, :, :k] = ... # where in k is the size of that type. 
daphnecor commented 2 months ago

Question: How do these structs correspond to the dynamics models? (i.e. classic, bicycle and deltalocal)

struct ClassicAction {
    float acceleration;
    float steering;
    float headAngle;
};

struct DeltaAction {
    float dx;
    float dy;
    float dyaw;
};

struct StateAction {
    Position position;
    float yaw;
    Velocity velocity;  
};

I understand that DeltaAction is for the DeltaLocal dynamics model, but thought that classic and bicycle share the same action space.

daphnecor commented 2 months ago

@aaravpandya

I integrated the dynamics model with a unified action tensor and verified that it works by launching a run.

Could you please review my changes (particularly the _copy_actions_to_simulator() method) and either edit / merge this PR?

daphnecor commented 2 months ago

Also, probably changes also need to be made in env_jax.py ?

The dynamics model changes are made directly in the base_env, the parent of env_jax. Action space is not updated but I think we can leave this for now.